feat: add ArchiveConversation to ClientRequest

This commit is contained in:
Michael Bolin
2025-09-08 23:07:25 -07:00
parent 4228a129cc
commit b69c2875d0
7 changed files with 185 additions and 3 deletions

View File

@@ -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<CodexConversation>`, 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<Arc<CodexConversation>> {
self.conversations.write().await.remove(conversation_id)
}
/// Fork an existing conversation by dropping the last `drop_last_messages`

View File

@@ -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;

View File

@@ -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;

View File

@@ -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;
@@ -73,12 +75,16 @@ use codex_protocol::protocol::USER_MESSAGE_BEGIN;
use mcp_types::JSONRPCErrorError;
use mcp_types::RequestId;
use std::collections::HashMap;
use std::ffi::OsStr;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use tokio::select;
use tokio::sync::Mutex;
use tokio::sync::oneshot;
use tracing::error;
use tracing::info;
use tracing::warn;
use uuid::Uuid;
// Duration before a ChatGPT login attempt is abandoned.
@@ -142,6 +148,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 +679,141 @@ impl CodexMessageProcessor {
}
}
async fn archive_conversation(&self, request_id: RequestId, params: 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);
let Some(file_name) = canonical_rollout_path.file_name().map(OsStr::to_owned) else {
let error = JSONRPCErrorError {
code: INVALID_REQUEST_ERROR_CODE,
message: format!(
"rollout path `{}` missing file name",
rollout_path.display()
),
data: None,
};
self.outgoing.send_error(request_id, error).await;
return;
};
if !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 let Some(conversation) = removed_conversation {
info!("conversation {conversation_id} was active; shutting down");
let conversation_clone = conversation.clone();
let notify = Arc::new(tokio::sync::Notify::new());
let notify_clone = notify.clone();
// Establish the listener for ShutdownComplete before submitting
// Shutdown so it is not missed.
let is_shutdown = tokio::spawn(async move {
loop {
select! {
_ = notify_clone.notified() => {
break;
}
event = conversation_clone.next_event() => {
if let Ok(event) = event && matches!(event.msg, EventMsg::ShutdownComplete) {
break;
}
}
}
}
});
// Request shutdown.
match conversation.submit(Op::Shutdown).await {
Ok(_) => {
// Successfully submitted Shutdown; wait before proceeding.
select! {
_ = is_shutdown => {
// Normal shutdown: proceed with archive.
}
_ = tokio::time::sleep(Duration::from_secs(10)) => {
warn!("conversation {conversation_id} shutdown timed out; proceeding with archive");
notify.notify_one();
}
}
}
Err(err) => {
error!("failed to submit Shutdown to conversation {conversation_id}: {err}");
notify.notify_one();
// Perhaps we lost a shutdown race, so let's continue to
// clean up the .jsonl file.
}
}
}
// Move the .jsonl file to the archived sessions subdir.
let result: std::io::Result<()> = async {
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,

View File

@@ -13,6 +13,7 @@ use anyhow::Context;
use assert_cmd::prelude::*;
use codex_mcp_server::CodexToolCallParam;
use codex_protocol::mcp_protocol::AddConversationListenerParams;
use codex_protocol::mcp_protocol::ArchiveConversationParams;
use codex_protocol::mcp_protocol::CancelLoginChatGptParams;
use codex_protocol::mcp_protocol::GetAuthStatusParams;
use codex_protocol::mcp_protocol::InterruptConversationParams;
@@ -186,6 +187,15 @@ impl McpProcess {
self.send_request("newConversation", params).await
}
/// Send an `archiveConversation` JSON-RPC request.
pub async fn send_archive_conversation_request(
&mut self,
params: ArchiveConversationParams,
) -> anyhow::Result<i64> {
let params = Some(serde_json::to_value(params)?);
self.send_request("archiveConversation", params).await
}
/// Send an `addConversationListener` JSON-RPC request.
pub async fn send_add_conversation_listener_request(
&mut self,

View File

@@ -1,4 +1,5 @@
// Aggregates all former standalone integration tests as modules.
mod archive_conversation;
mod auth;
mod codex_message_processor_flow;
mod codex_tool;

View File

@@ -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 {}