From cc3ab8e075a2df88e68cee6cf37ff1d0d783e251 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Mon, 8 Sep 2025 23:07:25 -0700 Subject: [PATCH] feat: add ArchiveConversation to ClientRequest --- codex-rs/core/src/conversation_manager.rs | 11 ++- codex-rs/core/src/lib.rs | 2 + codex-rs/core/src/rollout/mod.rs | 3 +- .../mcp-server/src/codex_message_processor.rs | 95 +++++++++++++++++++ codex-rs/protocol/src/mcp_protocol.rs | 17 ++++ 5 files changed, 125 insertions(+), 3 deletions(-) diff --git a/codex-rs/core/src/conversation_manager.rs b/codex-rs/core/src/conversation_manager.rs index 6fac42d59f..ec6e5a97c8 100644 --- a/codex-rs/core/src/conversation_manager.rs +++ b/codex-rs/core/src/conversation_manager.rs @@ -145,8 +145,15 @@ impl ConversationManager { self.finalize_spawn(codex, conversation_id).await } - pub async fn remove_conversation(&self, conversation_id: ConversationId) { - self.conversations.write().await.remove(&conversation_id); + /// Removes the conversation from the manager's internal map, though the + /// conversation is stored as `Arc`, it is possible that + /// other references to it exist elsewhere. Returns `true` if the + /// conversation was found and removed; otherwise, `false`. + pub async fn remove_conversation( + &self, + conversation_id: &ConversationId, + ) -> Option> { + self.conversations.write().await.remove(conversation_id) } /// Fork an existing conversation by dropping the last `drop_last_messages` diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 3b47830c50..b8cd7bc55a 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -61,7 +61,9 @@ pub mod spawn; pub mod terminal; mod tool_apply_patch; pub mod turn_diff_tracker; +pub use rollout::ARCHIVED_SESSIONS_SUBDIR; pub use rollout::RolloutRecorder; +pub use rollout::SESSIONS_SUBDIR; pub use rollout::SessionMeta; pub use rollout::list::ConversationItem; pub use rollout::list::ConversationsPage; diff --git a/codex-rs/core/src/rollout/mod.rs b/codex-rs/core/src/rollout/mod.rs index 4883517c1f..a78d760959 100644 --- a/codex-rs/core/src/rollout/mod.rs +++ b/codex-rs/core/src/rollout/mod.rs @@ -1,6 +1,7 @@ //! Rollout module: persistence and discovery of session rollout files. -pub(crate) const SESSIONS_SUBDIR: &str = "sessions"; +pub const SESSIONS_SUBDIR: &str = "sessions"; +pub const ARCHIVED_SESSIONS_SUBDIR: &str = "archived_sessions"; pub mod list; pub(crate) mod policy; diff --git a/codex-rs/mcp-server/src/codex_message_processor.rs b/codex-rs/mcp-server/src/codex_message_processor.rs index 2169a8f27f..575f1c6a73 100644 --- a/codex-rs/mcp-server/src/codex_message_processor.rs +++ b/codex-rs/mcp-server/src/codex_message_processor.rs @@ -35,6 +35,8 @@ use codex_protocol::mcp_protocol::AddConversationListenerParams; use codex_protocol::mcp_protocol::AddConversationSubscriptionResponse; use codex_protocol::mcp_protocol::ApplyPatchApprovalParams; use codex_protocol::mcp_protocol::ApplyPatchApprovalResponse; +use codex_protocol::mcp_protocol::ArchiveConversationParams; +use codex_protocol::mcp_protocol::ArchiveConversationResponse; use codex_protocol::mcp_protocol::AuthMode; use codex_protocol::mcp_protocol::AuthStatusChangeNotification; use codex_protocol::mcp_protocol::ClientRequest; @@ -79,6 +81,7 @@ use std::time::Duration; use tokio::sync::Mutex; use tokio::sync::oneshot; use tracing::error; +use tracing::warn; use uuid::Uuid; // Duration before a ChatGPT login attempt is abandoned. @@ -142,6 +145,9 @@ impl CodexMessageProcessor { ClientRequest::ResumeConversation { request_id, params } => { self.handle_resume_conversation(request_id, params).await; } + ClientRequest::ArchiveConversation { request_id, params } => { + self.archive_conversation(request_id, params).await; + } ClientRequest::SendUserMessage { request_id, params } => { self.send_user_message(request_id, params).await; } @@ -670,6 +676,95 @@ impl CodexMessageProcessor { } } + async fn archive_conversation( + &self, + request_id: RequestId, + params: codex_protocol::mcp_protocol::ArchiveConversationParams, + ) { + let ArchiveConversationParams { + conversation_id, + rollout_path, + } = params; + + // Verify that the rollout path is in the sessions directory or else + // a malicious client could specify an arbitrary path. + let rollout_folder = self.config.codex_home.join(codex_core::SESSIONS_SUBDIR); + let canonical_rollout_path = tokio::fs::canonicalize(&rollout_path).await; + let canonical_rollout_path = if let Ok(path) = canonical_rollout_path + && path.starts_with(&rollout_folder) + { + path + } else { + let error = JSONRPCErrorError { + code: INVALID_REQUEST_ERROR_CODE, + message: format!( + "rollout path `{}` must be in sessions directory", + rollout_path.display() + ), + data: None, + }; + self.outgoing.send_error(request_id, error).await; + return; + }; + + let required_suffix = format!("{}.jsonl", conversation_id.0); + if let Some(file_name) = canonical_rollout_path.file_name() + && file_name + .to_string_lossy() + .ends_with(required_suffix.as_str()) + { + let error = JSONRPCErrorError { + code: INVALID_REQUEST_ERROR_CODE, + message: format!( + "rollout path `{}` does not match conversation id {conversation_id}", + rollout_path.display() + ), + data: None, + }; + self.outgoing.send_error(request_id, error).await; + return; + } + + let removed_conversation = self + .conversation_manager + .remove_conversation(&conversation_id) + .await; + if removed_conversation.is_some() { + warn!( + "conversation {conversation_id} was active; removed from memory before archiving" + ); + } + + let result: std::io::Result<()> = async { + let file_name = canonical_rollout_path.file_name().ok_or_else(|| { + std::io::Error::other(format!("{canonical_rollout_path:?} file_name() is empty")) + })?; + let archive_folder = self + .config + .codex_home + .join(codex_core::ARCHIVED_SESSIONS_SUBDIR); + tokio::fs::create_dir_all(&archive_folder).await?; + tokio::fs::rename(&canonical_rollout_path, &archive_folder.join(file_name)).await?; + Ok(()) + } + .await; + + match result { + Ok(()) => { + let response = ArchiveConversationResponse {}; + self.outgoing.send_response(request_id, response).await; + } + Err(err) => { + let error = JSONRPCErrorError { + code: INTERNAL_ERROR_CODE, + message: format!("failed to archive conversation: {err}"), + data: None, + }; + self.outgoing.send_error(request_id, error).await; + } + } + } + async fn send_user_message(&self, request_id: RequestId, params: SendUserMessageParams) { let SendUserMessageParams { conversation_id, diff --git a/codex-rs/protocol/src/mcp_protocol.rs b/codex-rs/protocol/src/mcp_protocol.rs index 70e334036f..00391c7d01 100644 --- a/codex-rs/protocol/src/mcp_protocol.rs +++ b/codex-rs/protocol/src/mcp_protocol.rs @@ -91,6 +91,11 @@ pub enum ClientRequest { request_id: RequestId, params: ResumeConversationParams, }, + ArchiveConversation { + #[serde(rename = "id")] + request_id: RequestId, + params: ArchiveConversationParams, + }, SendUserMessage { #[serde(rename = "id")] request_id: RequestId, @@ -263,6 +268,18 @@ pub struct AddConversationSubscriptionResponse { pub subscription_id: Uuid, } +/// The [`ConversationId`] must match the `rollout_path`. +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, TS)] +#[serde(rename_all = "camelCase")] +pub struct ArchiveConversationParams { + pub conversation_id: ConversationId, + pub rollout_path: PathBuf, +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, TS)] +#[serde(rename_all = "camelCase")] +pub struct ArchiveConversationResponse {} + #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, TS)] #[serde(rename_all = "camelCase")] pub struct RemoveConversationSubscriptionResponse {}