codex/codex-rs/mcp-server/src/conversation_loop.rs

use std::collections::HashSet;
use std::sync::Arc;
use tokio::sync::Mutex;

use crate::exec_approval::handle_exec_approval_request;
use crate::mcp_protocol::CodexEventNotificationParams;
use crate::mcp_protocol::ConversationId;
use crate::mcp_protocol::InitialStateNotificationParams;
use crate::mcp_protocol::InitialStatePayload;
use crate::mcp_protocol::NotificationMeta;
use crate::outgoing_message::OutgoingMessageSender;
use crate::patch_approval::handle_patch_approval_request;
use codex_core::Codex;
use codex_core::protocol::AgentMessageEvent;
use codex_core::protocol::ApplyPatchApprovalRequestEvent;
use codex_core::protocol::EventMsg;
use codex_core::protocol::ExecApprovalRequestEvent;
use mcp_types::RequestId;
use tokio::sync::watch::Receiver as WatchReceiver;
use tracing::error;
use uuid::Uuid;

pub async fn run_conversation_loop(
    codex: Arc<Codex>,
    outgoing: Arc<OutgoingMessageSender>,
    request_id: RequestId,
    mut stream_rx: WatchReceiver<bool>,
    session_id: Uuid,
    running_sessions: Arc<Mutex<HashSet<Uuid>>>,
) {
    let request_id_str = match &request_id {
        RequestId::String(s) => s.clone(),
        RequestId::Integer(n) => n.to_string(),
    };

    // Buffer all events for InitialState
    let mut buffered_events: Vec<CodexEventNotificationParams> = Vec::new();
    let mut streaming_enabled = *stream_rx.borrow();

    loop {
        tokio::select! {
            res = codex.next_event() => {
                match res {
                    Ok(event) => {
                        // Always buffer the event
                        buffered_events.push(CodexEventNotificationParams { meta: None, msg: event.msg.clone() });

                        if streaming_enabled {
                            let method = event.msg.to_string();
                            let params = CodexEventNotificationParams { meta: None, msg: event.msg.clone() };
                            if let Ok(params_val) = serde_json::to_value(&params) {
                                outgoing
                                    .send_custom_notification(&method, params_val)
                                    .await;
                            } else {
                                error!("Failed to serialize event params");
                            }
                        }

                        match event.msg {
                    EventMsg::ExecApprovalRequest(ExecApprovalRequestEvent {
                        command,
                        cwd,
                        call_id,
                        reason: _,
                    }) => {
                                if streaming_enabled {
                                    handle_exec_approval_request(
                                        command,
                                        cwd,
                                        outgoing.clone(),
                                        codex.clone(),
                                        request_id.clone(),
                                        request_id_str.clone(),
                                        event.id.clone(),
                                        call_id,
                                    )
                                    .await;
                                }
                        continue;
                    }
                    EventMsg::Error(_) => {
                        error!("Codex runtime error");
                    }
                    EventMsg::ApplyPatchApprovalRequest(ApplyPatchApprovalRequestEvent {
                        call_id,
                        reason,
                        grant_root,
                        changes,
                    }) => {
                            if streaming_enabled {
                                handle_patch_approval_request(
                                    call_id,
                                    reason,
                                    grant_root,
                                    changes,
                                    outgoing.clone(),
                                    codex.clone(),
                                    request_id.clone(),
                                    request_id_str.clone(),
                                    event.id.clone(),
                                )
                                .await;
                            }
                        continue;
                    }
                    EventMsg::TaskComplete(_) => {
                        // remove running session id
                        let mut running_sessions = running_sessions.lock().await;
                        running_sessions.remove(&session_id);
                    }
                    EventMsg::SessionConfigured(_) => {
                        tracing::error!("unexpected SessionConfigured event");
                    }
                    EventMsg::AgentMessageDelta(_) => {
                        // TODO: think how we want to support this in the MCP
                    }
                    EventMsg::AgentReasoningDelta(_) => {
                        // TODO: think how we want to support this in the MCP
                    }
                    EventMsg::AgentMessage(AgentMessageEvent { .. }) => {
                        // TODO: think how we want to support this in the MCP
                    }
                    EventMsg::TaskStarted
                    | EventMsg::TokenCount(_)
                    | EventMsg::AgentReasoning(_)
                    | EventMsg::McpToolCallBegin(_)
                    | EventMsg::McpToolCallEnd(_)
                    | EventMsg::ExecCommandBegin(_)
                    | EventMsg::ExecCommandEnd(_)
                    | EventMsg::BackgroundEvent(_)
                    | EventMsg::ExecCommandOutputDelta(_)
                    | EventMsg::PatchApplyBegin(_)
                    | EventMsg::PatchApplyEnd(_)
                    | EventMsg::GetHistoryEntryResponse(_)
                    | EventMsg::PlanUpdate(_)
                    | EventMsg::ShutdownComplete => {
                        // For now, we do not do anything extra for these
                        // events. Note that
                        // send(codex_event_to_notification(&event)) above has
                        // already dispatched these events as notifications,
                        // though we may want to do give different treatment to
                        // individual events in the future.
                    }
                        }
                    }
                    Err(e) => {
                        error!("Codex runtime error: {e}");
                    }
                }
            },
            changed = stream_rx.changed() => {
                if changed.is_ok() {
                    let now = *stream_rx.borrow();
                    if now && !streaming_enabled {
                        streaming_enabled = true;
                        // Emit InitialState with all buffered events
                        let params = InitialStateNotificationParams {
                            meta: Some(NotificationMeta { conversation_id: Some(ConversationId(session_id)), request_id: None }),
                            initial_state: InitialStatePayload { events: buffered_events.clone() },
                        };
                        if let Ok(params_val) = serde_json::to_value(&params) {
                            outgoing
                                .send_custom_notification("notifications/initial_state", params_val)
                                .await;
                        } else {
                            error!("Failed to serialize InitialState params");
                        }
                    } else if !now && streaming_enabled {
                        // streaming disabled
                        streaming_enabled = false;
                    }
                }
            }
        }
    }
}