oak-gpui/crates/agent_servers/src/codex.rs

use agent_client_protocol as acp;
use anyhow::anyhow;
use collections::HashMap;
use context_server::listener::McpServerTool;
use context_server::types::requests;
use context_server::{ContextServer, ContextServerCommand, ContextServerId};
use futures::channel::{mpsc, oneshot};
use project::Project;
use settings::SettingsStore;
use smol::stream::StreamExt as _;
use std::cell::RefCell;
use std::rc::Rc;
use std::{path::Path, sync::Arc};
use util::ResultExt;

use anyhow::{Context, Result};
use gpui::{App, AppContext as _, AsyncApp, Entity, Task, WeakEntity};

use crate::mcp_server::ZedMcpServer;
use crate::{AgentServer, AgentServerCommand, AllAgentServersSettings, mcp_server};
use acp_thread::{AcpThread, AgentConnection};

#[derive(Clone)]
pub struct Codex;

impl AgentServer for Codex {
    fn name(&self) -> &'static str {
        "Codex"
    }

    fn empty_state_headline(&self) -> &'static str {
        "Welcome to Codex"
    }

    fn empty_state_message(&self) -> &'static str {
        "What can I help with?"
    }

    fn logo(&self) -> ui::IconName {
        ui::IconName::AiOpenAi
    }

    fn connect(
        &self,
        _root_dir: &Path,
        project: &Entity<Project>,
        cx: &mut App,
    ) -> Task<Result<Rc<dyn AgentConnection>>> {
        let project = project.clone();
        let working_directory = project.read(cx).active_project_directory(cx);
        cx.spawn(async move |cx| {
            let settings = cx.read_global(|settings: &SettingsStore, _| {
                settings.get::<AllAgentServersSettings>(None).codex.clone()
            })?;

            let Some(command) =
                AgentServerCommand::resolve("codex", &["mcp"], settings, &project, cx).await
            else {
                anyhow::bail!("Failed to find codex binary");
            };

            let client: Arc<ContextServer> = ContextServer::stdio(
                ContextServerId("codex-mcp-server".into()),
                ContextServerCommand {
                    path: command.path,
                    args: command.args,
                    env: command.env,
                },
                working_directory,
            )
            .into();
            ContextServer::start(client.clone(), cx).await?;

            let (notification_tx, mut notification_rx) = mpsc::unbounded();
            client
                .client()
                .context("Failed to subscribe")?
                .on_notification(acp::SESSION_UPDATE_METHOD_NAME, {
                    move |notification, _cx| {
                        let notification_tx = notification_tx.clone();
                        log::trace!(
                            "ACP Notification: {}",
                            serde_json::to_string_pretty(&notification).unwrap()
                        );

                        if let Some(notification) =
                            serde_json::from_value::<acp::SessionNotification>(notification)
                                .log_err()
                        {
                            notification_tx.unbounded_send(notification).ok();
                        }
                    }
                });

            let sessions = Rc::new(RefCell::new(HashMap::default()));

            let notification_handler_task = cx.spawn({
                let sessions = sessions.clone();
                async move |cx| {
                    while let Some(notification) = notification_rx.next().await {
                        CodexConnection::handle_session_notification(
                            notification,
                            sessions.clone(),
                            cx,
                        )
                    }
                }
            });

            let connection = CodexConnection {
                client,
                sessions,
                _notification_handler_task: notification_handler_task,
            };
            Ok(Rc::new(connection) as _)
        })
    }
}

struct CodexConnection {
    client: Arc<context_server::ContextServer>,
    sessions: Rc<RefCell<HashMap<acp::SessionId, CodexSession>>>,
    _notification_handler_task: Task<()>,
}

struct CodexSession {
    thread: WeakEntity<AcpThread>,
    cancel_tx: Option<oneshot::Sender<()>>,
    _mcp_server: ZedMcpServer,
}

impl AgentConnection for CodexConnection {
    fn name(&self) -> &'static str {
        "Codex"
    }

    fn new_thread(
        self: Rc<Self>,
        project: Entity<Project>,
        cwd: &Path,
        cx: &mut AsyncApp,
    ) -> Task<Result<Entity<AcpThread>>> {
        let client = self.client.client();
        let sessions = self.sessions.clone();
        let cwd = cwd.to_path_buf();
        cx.spawn(async move |cx| {
            let client = client.context("MCP server is not initialized yet")?;
            let (mut thread_tx, thread_rx) = watch::channel(WeakEntity::new_invalid());

            let mcp_server = ZedMcpServer::new(thread_rx, cx).await?;

            let response = client
                .request::<requests::CallTool>(context_server::types::CallToolParams {
                    name: acp::NEW_SESSION_TOOL_NAME.into(),
                    arguments: Some(serde_json::to_value(acp::NewSessionArguments {
                        mcp_servers: [(
                            mcp_server::SERVER_NAME.to_string(),
                            mcp_server.server_config()?,
                        )]
                        .into(),
                        client_tools: acp::ClientTools {
                            request_permission: Some(acp::McpToolId {
                                mcp_server: mcp_server::SERVER_NAME.into(),
                                tool_name: mcp_server::RequestPermissionTool::NAME.into(),
                            }),
                            read_text_file: Some(acp::McpToolId {
                                mcp_server: mcp_server::SERVER_NAME.into(),
                                tool_name: mcp_server::ReadTextFileTool::NAME.into(),
                            }),
                            write_text_file: Some(acp::McpToolId {
                                mcp_server: mcp_server::SERVER_NAME.into(),
                                tool_name: mcp_server::WriteTextFileTool::NAME.into(),
                            }),
                        },
                        cwd,
                    })?),
                    meta: None,
                })
                .await?;

            if response.is_error.unwrap_or_default() {
                return Err(anyhow!(response.text_contents()));
            }

            let result = serde_json::from_value::<acp::NewSessionOutput>(
                response.structured_content.context("Empty response")?,
            )?;

            let thread =
                cx.new(|cx| AcpThread::new(self.clone(), project, result.session_id.clone(), cx))?;

            thread_tx.send(thread.downgrade())?;

            let session = CodexSession {
                thread: thread.downgrade(),
                cancel_tx: None,
                _mcp_server: mcp_server,
            };
            sessions.borrow_mut().insert(result.session_id, session);

            Ok(thread)
        })
    }

    fn authenticate(&self, _cx: &mut App) -> Task<Result<()>> {
        Task::ready(Err(anyhow!("Authentication not supported")))
    }

    fn prompt(
        &self,
        params: agent_client_protocol::PromptArguments,
        cx: &mut App,
    ) -> Task<Result<()>> {
        let client = self.client.client();
        let sessions = self.sessions.clone();

        cx.foreground_executor().spawn(async move {
            let client = client.context("MCP server is not initialized yet")?;

            let (new_cancel_tx, cancel_rx) = oneshot::channel();
            {
                let mut sessions = sessions.borrow_mut();
                let session = sessions
                    .get_mut(&params.session_id)
                    .context("Session not found")?;
                session.cancel_tx.replace(new_cancel_tx);
            }

            let result = client
                .request_with::<requests::CallTool>(
                    context_server::types::CallToolParams {
                        name: acp::PROMPT_TOOL_NAME.into(),
                        arguments: Some(serde_json::to_value(params)?),
                        meta: None,
                    },
                    Some(cancel_rx),
                    None,
                )
                .await;

            if let Err(err) = &result
                && err.is::<context_server::client::RequestCanceled>()
            {
                return Ok(());
            }

            let response = result?;

            if response.is_error.unwrap_or_default() {
                return Err(anyhow!(response.text_contents()));
            }

            Ok(())
        })
    }

    fn cancel(&self, session_id: &agent_client_protocol::SessionId, _cx: &mut App) {
        let mut sessions = self.sessions.borrow_mut();

        if let Some(cancel_tx) = sessions
            .get_mut(session_id)
            .and_then(|session| session.cancel_tx.take())
        {
            cancel_tx.send(()).ok();
        }
    }
}

impl CodexConnection {
    pub fn handle_session_notification(
        notification: acp::SessionNotification,
        threads: Rc<RefCell<HashMap<acp::SessionId, CodexSession>>>,
        cx: &mut AsyncApp,
    ) {
        let threads = threads.borrow();
        let Some(thread) = threads
            .get(&notification.session_id)
            .and_then(|session| session.thread.upgrade())
        else {
            log::error!(
                "Thread not found for session ID: {}",
                notification.session_id
            );
            return;
        };

        thread
            .update(cx, |thread, cx| {
                thread.handle_session_update(notification.update, cx)
            })
            .log_err();
    }
}

impl Drop for CodexConnection {
    fn drop(&mut self) {
        self.client.stop().log_err();
    }
}

#[cfg(test)]
pub(crate) mod tests {
    use super::*;
    use crate::AgentServerCommand;
    use std::path::Path;

    crate::common_e2e_tests!(Codex, allow_option_id = "approve");

    pub fn local_command() -> AgentServerCommand {
        let cli_path = Path::new(env!("CARGO_MANIFEST_DIR"))
            .join("../../../codex/codex-rs/target/debug/codex");

        AgentServerCommand {
            path: cli_path,
            args: vec![],
            env: None,
        }
    }
}