codex/codex-rs/core/src/agent_message_board.rs

//! Bridges the board extension to the selected controller, clock and active-turn services.
//!
//! The extension owns storage and tools. This adapter never starts or restores
//! recipients and never queues a notification for an idle agent.

use crate::CodexThread;
use crate::ThreadManager;
use crate::config::Config;
use crate::context::AgentMessageBoardNotification;
use crate::context::ContextualUserFragment;
use crate::tools::MULTI_AGENT_V2_NAMESPACE_DESCRIPTION;
use chrono::DateTime;
use chrono::Utc;
use codex_agent_message_board_extension::AgentMessageBoard;
use codex_agent_message_board_extension::InMemoryMessageBoards;
use codex_agent_message_board_extension::LocalAgentMessageBoard;
use codex_agent_message_board_extension::MessageBoardHost;
use codex_agent_message_board_extension::NotificationDelivery;
use codex_agent_message_board_extension::PostPreview;
use codex_extension_api::ExtensionRegistryBuilder;
use codex_features::Feature;
use codex_protocol::AgentPath;
use codex_protocol::SessionId;
use codex_protocol::ThreadId;
use codex_protocol::error::CodexErr;
use codex_protocol::error::CodexErrorDetails;
use codex_protocol::error::Result;
use codex_protocol::protocol::InterAgentCommunication;
use futures::future::BoxFuture;
use std::sync::Arc;
use std::sync::Weak;

/// Registers the configured board for opted-in MAv2 runtimes.
/// Training can share an in-memory board per tree, including ephemeral sessions.
pub fn install_agent_message_board(
    registry: &mut ExtensionRegistryBuilder<Config>,
    manager: Weak<ThreadManager>,
) {
    let in_memory_boards = Arc::new(InMemoryMessageBoards::default());
    codex_agent_message_board_extension::install(
        registry,
        MULTI_AGENT_V2_NAMESPACE_DESCRIPTION,
        |config: &Config| config.multi_agent_v2.tool_namespace.clone(),
        move |config: &Config, tree, caller| {
            let in_memory = config.multi_agent_v2.message_board_in_memory;
            // MAv2 supplies tree paths; ephemeral runtimes must not open durable storage.
            if !config.features.enabled(Feature::AgentMessageBoard)
                || !config.features.enabled(Feature::MultiAgentV2)
                || (config.ephemeral && !in_memory)
            {
                return Box::pin(async { Ok(None) });
            }
            let sqlite = config.sqlite_config().clone();
            let in_memory_boards = Arc::clone(&in_memory_boards);
            let host = Arc::new(LocalBoardHost {
                manager: manager.clone(),
                tree,
                caller,
            });
            Box::pin(async move {
                let board: Arc<dyn AgentMessageBoard> = if in_memory {
                    Arc::new(in_memory_boards.open(tree, host).await)
                } else {
                    Arc::new(LocalAgentMessageBoard::open(&sqlite, tree, host).await?)
                };
                Ok(Some(board))
            })
        },
    );
}

struct LocalBoardHost {
    manager: Weak<ThreadManager>,
    tree: SessionId,
    caller: ThreadId,
}

impl LocalBoardHost {
    async fn actor(&self, caller: ThreadId) -> Result<Arc<CodexThread>> {
        let manager = self
            .manager
            .upgrade()
            .ok_or_else(|| CodexErr::ThreadNotFound(caller))?;
        let actor = manager.get_thread(caller).await?;
        if actor.session.session_id() != self.tree {
            return Err(CodexErr::InvalidRequest(
                "agent belongs to another message board".into(),
            ));
        }
        Ok(actor)
    }

    async fn registered_path(&self, actor: &CodexThread) -> Result<AgentPath> {
        let caller = actor.session.thread_id;
        let path = actor
            .session_source
            .get_agent_path()
            .or_else(|| (caller == ThreadId::from(self.tree)).then(AgentPath::root))
            .ok_or_else(|| CodexErr::InvalidRequest("agent has no tree path".into()))?;
        let resolved = actor
            .session
            .services
            .agent_control
            .resolve(
                caller,
                actor.session_source.parent_thread_id(),
                &actor.session_source,
                path.as_str(),
            )
            .await?;
        if resolved != caller {
            return Err(CodexErr::InvalidRequest(
                "agent does not own its tree path".into(),
            ));
        }
        Ok(path)
    }
}

impl MessageBoardHost for LocalBoardHost {
    fn agent_path(&self, caller: ThreadId) -> BoxFuture<'_, Result<AgentPath>> {
        Box::pin(async move {
            let actor = self.actor(caller).await?;
            self.registered_path(&actor).await
        })
    }

    fn resolve_agent(&self, path: AgentPath) -> BoxFuture<'_, Result<ThreadId>> {
        Box::pin(async move {
            let actor = self.actor(self.caller).await?;
            actor
                .session
                .services
                .agent_control
                .resolve(
                    self.caller,
                    actor.session_source.parent_thread_id(),
                    &actor.session_source,
                    path.as_str(),
                )
                .await
        })
    }

    fn current_time(&self, caller: ThreadId) -> BoxFuture<'_, Result<DateTime<Utc>>> {
        Box::pin(async move {
            let actor = self.actor(caller).await?;
            self.registered_path(&actor).await?;
            actor
                .session
                .services
                .time_provider
                .current_time(caller)
                .await
                .map_err(|error| CodexErr::Io(std::io::Error::other(error)))
        })
    }

    fn notify(
        &self,
        recipient_id: ThreadId,
        post: PostPreview,
    ) -> BoxFuture<'_, Result<NotificationDelivery>> {
        Box::pin(async move {
            let Some(manager) = self.manager.upgrade() else {
                return Ok(NotificationDelivery::SkippedInactive);
            };
            let recipient = match manager.get_thread(recipient_id).await {
                Ok(thread) => thread,
                Err(error) if matches!(error.details(), CodexErrorDetails::ThreadNotFound(_)) => {
                    return Ok(NotificationDelivery::SkippedInactive);
                }
                Err(error) => return Err(error),
            };
            if recipient.session.session_id() != self.tree {
                return Err(CodexErr::InvalidRequest(
                    "notification recipient belongs to another board".into(),
                ));
            }
            let recipient_path = self.registered_path(&recipient).await?;
            let notice = AgentMessageBoardNotification(post);
            let communication = InterAgentCommunication::new(
                notice.0.metadata.author.clone(),
                recipient_path,
                Vec::new(),
                notice.render(),
                /*trigger_turn*/ false,
            );
            Ok(
                match recipient
                    .inject_if_running(vec![communication.to_model_input_item()])
                    .await
                {
                    Ok(()) => NotificationDelivery::Accepted,
                    Err(_) => NotificationDelivery::SkippedInactive,
                },
            )
        })
    }
}