openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-28 18:58:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

auto-attach convo listener

Browse Source
This commit is contained in:
Owen Lin
2025-11-02 16:11:12 -08:00
parent 698c6bd1de
commit fa03208ade
3 changed files with 465 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server/src/codex_message_processor.rs
+120 -29
View File
@@ -155,6 +155,10 @@ pub(crate) struct CodexMessageProcessor {
codex_linux_sandbox_exe: Option<PathBuf>,
config: Arc<Config>,
conversation_listeners: HashMap<Uuid, oneshot::Sender<()>>,
// Map conversations to their active listener subscription ids (legacy + v2 auto-attached).
conversation_listener_index: HashMap<ConversationId, Vec<Uuid>>,
// Reverse lookup from subscription id to conversation id for clean removal.
subscription_to_conversation: HashMap<Uuid, ConversationId>,
active_login: Arc<Mutex<Option<ActiveLogin>>>,
// Queue of pending interrupt requests per conversation. We reply when TurnAborted arrives.
pending_interrupts: Arc<Mutex<HashMap<ConversationId, Vec<RequestId>>>>,
@@ -178,6 +182,8 @@ impl CodexMessageProcessor {
codex_linux_sandbox_exe,
config,
conversation_listeners: HashMap::new(),
conversation_listener_index: HashMap::new(),
subscription_to_conversation: HashMap::new(),
active_login: Arc::new(Mutex::new(None)),
pending_interrupts: Arc::new(Mutex::new(HashMap::new())),
pending_fuzzy_searches: Arc::new(Mutex::new(HashMap::new())),
@@ -870,7 +876,7 @@ impl CodexMessageProcessor {
}
}
async fn thread_start(&self, request_id: RequestId, params: ThreadStartParams) {
async fn thread_start(&mut self, request_id: RequestId, params: ThreadStartParams) {
// Map v2 ThreadStartParams onto the existing NewConversationParams so we
// can reuse the same config derivation and conversation manager flow.
let overrides = NewConversationParams {
@@ -927,6 +933,19 @@ impl CodexMessageProcessor {
self.outgoing
.send_server_notification(ServerNotification::ThreadStarted(notif))
.await;
// Auto-attach a conversation listener when starting a thread.
// Use the same behavior as the v1 API with experimental_raw_events=false.
if let Err(err) = self
.attach_conversation_listener(new_conv.conversation_id, false)
.await
{
tracing::warn!(
"failed to attach listener for conversation {}: {}",
new_conv.conversation_id,
err.message
);
}
}
Err(err) => {
let error = JSONRPCErrorError {
@@ -1048,7 +1067,7 @@ impl CodexMessageProcessor {
}
}
async fn thread_archive(&self, request_id: RequestId, params: ThreadArchiveParams) {
async fn thread_archive(&mut self, request_id: RequestId, params: ThreadArchiveParams) {
// Resolve conversation id from v2 thread id string.
let conversation_id = match ConversationId::from_string(&params.thread_id) {
Ok(id) => id,
@@ -1106,7 +1125,7 @@ impl CodexMessageProcessor {
}
async fn archive_conversation_common(
&self,
&mut self,
conversation_id: ConversationId,
rollout_path: &Path,
) -> Result<(), JSONRPCErrorError> {
@@ -1160,6 +1179,10 @@ impl CodexMessageProcessor {
});
}
// Proactively cancel any listeners for this conversation so our
// shutdown wait does not race with listener tasks consuming events.
self.cancel_conversation_listeners(conversation_id);
// If the conversation is active, request shutdown and wait briefly.
if let Some(conversation) = self
.conversation_manager
@@ -1305,7 +1328,7 @@ impl CodexMessageProcessor {
self.outgoing.send_response(request_id, response).await;
}
async fn thread_resume(&self, request_id: RequestId, params: ThreadResumeParams) {
async fn thread_resume(&mut self, request_id: RequestId, params: ThreadResumeParams) {
// Convert thread id string into ConversationId.
let conversation_id = match ConversationId::from_string(&params.thread_id) {
Ok(id) => id,
@@ -1397,6 +1420,18 @@ impl CodexMessageProcessor {
},
};
self.outgoing.send_response(request_id, response).await;
// Auto-attach a conversation listener when resuming a thread.
if let Err(err) = self
.attach_conversation_listener(conversation_id, false)
.await
{
tracing::warn!(
"failed to attach listener for conversation {}: {}",
conversation_id,
err.message
);
}
}
Err(err) => {
let error = JSONRPCErrorError {
@@ -1764,7 +1799,11 @@ impl CodexMessageProcessor {
self.outgoing.send_error(request_id, error).await;
}
async fn archive_conversation(&self, request_id: RequestId, params: ArchiveConversationParams) {
async fn archive_conversation(
&mut self,
request_id: RequestId,
params: ArchiveConversationParams,
) {
let ArchiveConversationParams {
conversation_id,
rollout_path,
@@ -1916,34 +1955,64 @@ impl CodexMessageProcessor {
conversation_id,
experimental_raw_events,
} = params;
let Ok(conversation) = self
match self
.attach_conversation_listener(conversation_id, experimental_raw_events)
.await
{
Ok(subscription_id) => {
let response = AddConversationSubscriptionResponse { subscription_id };
self.outgoing.send_response(request_id, response).await;
}
Err(err) => {
self.outgoing.send_error(request_id, err).await;
}
}
}
async fn attach_conversation_listener(
&mut self,
conversation_id: ConversationId,
experimental_raw_events: bool,
) -> Result<Uuid, JSONRPCErrorError> {
let conversation = match self
.conversation_manager
.get_conversation(conversation_id)
.await
else {
let error = JSONRPCErrorError {
code: INVALID_REQUEST_ERROR_CODE,
message: format!("conversation not found: {conversation_id}"),
data: None,
};
self.outgoing.send_error(request_id, error).await;
return;
{
Ok(conv) => conv,
Err(_) => {
return Err(JSONRPCErrorError {
code: INVALID_REQUEST_ERROR_CODE,
message: format!("conversation not found: {conversation_id}"),
data: None,
});
}
};
let subscription_id = Uuid::new_v4();
let (cancel_tx, mut cancel_rx) = oneshot::channel();
self.conversation_listeners
.insert(subscription_id, cancel_tx);
// Track mappings so we can cleanly cancel listeners by conversation id.
self.subscription_to_conversation
.insert(subscription_id, conversation_id);
self.conversation_listener_index
.entry(conversation_id)
.or_default()
.push(subscription_id);
let outgoing_for_task = self.outgoing.clone();
let pending_interrupts = self.pending_interrupts.clone();
let conversation_for_task = conversation;
let conversation_id_for_task = conversation_id;
tokio::spawn(async move {
loop {
tokio::select! {
_ = &mut cancel_rx => {
// User has unsubscribed, so exit this task.
break;
}
event = conversation.next_event() => {
event = conversation_for_task.next_event() => {
let event = match event {
Ok(event) => event,
Err(err) => {
@@ -1957,10 +2026,6 @@ impl CodexMessageProcessor {
continue;
}
// For now, we send a notification for every event,
// JSON-serializing the `Event` as-is, but these should
// be migrated to be variants of `ServerNotification`
// instead.
let method = format!("codex/event/{}", event.msg);
let mut params = match serde_json::to_value(event.clone()) {
Ok(serde_json::Value::Object(map)) => map,
@@ -1973,21 +2038,39 @@ impl CodexMessageProcessor {
continue;
}
};
params.insert("conversationId".to_string(), conversation_id.to_string().into());
params.insert("conversationId".to_string(), conversation_id_for_task.to_string().into());
outgoing_for_task.send_notification(OutgoingNotification {
method,
params: Some(params.into()),
})
.await;
outgoing_for_task
.send_notification(OutgoingNotification {
method,
params: Some(params.into()),
})
.await;
apply_bespoke_event_handling(event.clone(), conversation_id, conversation.clone(), outgoing_for_task.clone(), pending_interrupts.clone()).await;
apply_bespoke_event_handling(
event.clone(),
conversation_id_for_task,
conversation_for_task.clone(),
outgoing_for_task.clone(),
pending_interrupts.clone(),
).await;
}
}
}
});
let response = AddConversationSubscriptionResponse { subscription_id };
self.outgoing.send_response(request_id, response).await;
Ok(subscription_id)
}
fn cancel_conversation_listeners(&mut self, conversation_id: ConversationId) {
if let Some(ids) = self.conversation_listener_index.remove(&conversation_id) {
for id in ids {
if let Some(sender) = self.conversation_listeners.remove(&id) {
let _ = sender.send(());
}
self.subscription_to_conversation.remove(&id);
}
}
}
async fn remove_conversation_listener(
@@ -2000,6 +2083,14 @@ impl CodexMessageProcessor {
Some(sender) => {
// Signal the spawned task to exit and acknowledge.
let _ = sender.send(());
// Clean reverse mappings.
if let Some(conv_id) = self.subscription_to_conversation.remove(&subscription_id)
&& let Some(list) = self.conversation_listener_index.get_mut(&conv_id) {
list.retain(|id| id != &subscription_id);
if list.is_empty() {
self.conversation_listener_index.remove(&conv_id);
}
}
let response = RemoveConversationSubscriptionResponse {};
self.outgoing.send_response(request_id, response).await;
}
codex-rs/app-server/tests/common/mcp_process.rs
+10 -1
View File
@@ -60,7 +60,16 @@ pub struct McpProcess {
impl McpProcess {
pub async fn new(codex_home: &Path) -> anyhow::Result<Self> {
Self::new_with_env(codex_home, &[]).await
// Ensure tests run with the default Codex originator (codex_cli_rs)
// regardless of any environment overrides that may be present.
Self::new_with_env(
codex_home,
&[(
codex_core::default_client::CODEX_INTERNAL_ORIGINATOR_OVERRIDE_ENV_VAR,
None,
)],
)
.await
}
/// Creates a new MCP process, allowing tests to override or remove
codex-rs/app-server/tests/suite/v2/turn_start.rs
+335 -4
View File
@@ -1,18 +1,28 @@
use anyhow::Result;
use app_test_support::McpProcess;
use app_test_support::create_final_assistant_message_sse_response;
use app_test_support::create_mock_chat_completions_server_unchecked as create_mock_chat_completions_server;
use app_test_support::create_mock_chat_completions_server;
use app_test_support::create_mock_chat_completions_server_unchecked;
use app_test_support::create_shell_sse_response;
use app_test_support::to_response;
use codex_app_server_protocol::JSONRPCNotification;
use codex_app_server_protocol::JSONRPCResponse;
use codex_app_server_protocol::RequestId;
use codex_app_server_protocol::ServerRequest;
use codex_app_server_protocol::ThreadStartParams;
use codex_app_server_protocol::ThreadStartResponse;
use codex_app_server_protocol::TurnStartParams;
use codex_app_server_protocol::TurnStartResponse;
use codex_app_server_protocol::TurnStartedNotification;
use codex_app_server_protocol::UserInput as V2UserInput;
use codex_core::protocol_config_types::ReasoningEffort;
use codex_core::protocol_config_types::ReasoningSummary;
use codex_core::spawn::CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR;
use codex_protocol::parse_command::ParsedCommand;
use codex_protocol::protocol::Event;
use codex_protocol::protocol::EventMsg;
use pretty_assertions::assert_eq;
use std::env;
use std::path::Path;
use tempfile::TempDir;
use tokio::time::timeout;
@@ -22,12 +32,13 @@ const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn turn_start_emits_notifications_and_accepts_model_override() -> Result<()> {
// Provide a mock server and config so model wiring is valid.
// Two Codex turns hit the mock model (session start + turn/start).
// Three Codex turns hit the mock model (session start + two turn/start calls).
let responses = vec![
create_final_assistant_message_sse_response("Done")?,
create_final_assistant_message_sse_response("Done")?,
create_final_assistant_message_sse_response("Done")?,
];
let server = create_mock_chat_completions_server(responses).await;
let server = create_mock_chat_completions_server_unchecked(responses).await;
let codex_home = TempDir::new()?;
create_config_toml(codex_home.path(), &server.uri())?;
@@ -117,6 +128,14 @@ async fn turn_start_emits_notifications_and_accepts_model_override() -> Result<(
)
.await??;
// And we should ultimately get a task_complete without having to add a
// legacy conversation listener explicitly (auto-attached by thread/start).
let _task_complete: JSONRPCNotification = timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_notification_message("codex/event/task_complete"),
)
.await??;
drop(server);
Ok(())
}
@@ -128,7 +147,9 @@ async fn turn_start_accepts_local_image_input() -> Result<()> {
create_final_assistant_message_sse_response("Done")?,
create_final_assistant_message_sse_response("Done")?,
];
let server = create_mock_chat_completions_server(responses).await;
// Use the unchecked variant because the request payload includes a LocalImage
// which the strict matcher does not currently cover.
let server = create_mock_chat_completions_server_unchecked(responses).await;
let codex_home = TempDir::new()?;
create_config_toml(codex_home.path(), &server.uri())?;
@@ -178,6 +199,293 @@ async fn turn_start_accepts_local_image_input() -> Result<()> {
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn turn_start_exec_approval_toggle_v2() -> Result<()> {
if env::var(CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR).is_ok() {
println!("Skipping v2 exec approval toggle test due to sandbox network disabled.");
return Ok(());
}
let tmp = TempDir::new()?;
let codex_home = tmp.path().to_path_buf();
// Mock server: first turn requests a shell call (elicitation), then completes.
// Second turn same, but we'll set approval_policy=never to avoid elicitation.
let responses = vec![
create_shell_sse_response(
vec![
"python3".to_string(),
"-c".to_string(),
"print(42)".to_string(),
],
None,
Some(5000),
"call1",
)?,
create_final_assistant_message_sse_response("done 1")?,
create_shell_sse_response(
vec![
"python3".to_string(),
"-c".to_string(),
"print(42)".to_string(),
],
None,
Some(5000),
"call2",
)?,
create_final_assistant_message_sse_response("done 2")?,
];
let server = create_mock_chat_completions_server(responses).await;
// Default approval is untrusted to force elicitation on first turn.
create_config_toml_untrusted(codex_home.as_path(), &server.uri())?;
let mut mcp = McpProcess::new(codex_home.as_path()).await?;
timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??;
// thread/start
let start_id = mcp
.send_thread_start_request(ThreadStartParams {
model: Some("mock-model".to_string()),
..Default::default()
})
.await?;
let start_resp: JSONRPCResponse = timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_response_message(RequestId::Integer(start_id)),
)
.await??;
let ThreadStartResponse { thread } = to_response::<ThreadStartResponse>(start_resp)?;
// turn/start — expect ExecCommandApproval request from server
let first_turn_id = mcp
.send_turn_start_request(TurnStartParams {
thread_id: thread.id.clone(),
input: vec![V2UserInput::Text {
text: "run python".to_string(),
}],
cwd: None,
approval_policy: None,
sandbox_policy: None,
model: None,
effort: None,
summary: None,
})
.await?;
// Acknowledge RPC
timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_response_message(RequestId::Integer(first_turn_id)),
)
.await??;
// Receive elicitation
let server_req = timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_request_message(),
)
.await??;
let ServerRequest::ExecCommandApproval { request_id, params } = server_req else {
panic!("expected ExecCommandApproval request");
};
assert_eq!(params.call_id, "call1");
assert_eq!(
params.parsed_cmd,
vec![ParsedCommand::Unknown {
cmd: "python3 -c 'print(42)'".to_string()
}]
);
// Approve and wait for task completion
mcp.send_response(
request_id,
serde_json::json!({ "decision": codex_core::protocol::ReviewDecision::Approved }),
)
.await?;
timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_notification_message("codex/event/task_complete"),
)
.await??;
// Second turn with approval_policy=never should not elicit approval
let second_turn_id = mcp
.send_turn_start_request(TurnStartParams {
thread_id: thread.id.clone(),
input: vec![V2UserInput::Text {
text: "run python again".to_string(),
}],
cwd: None,
approval_policy: Some(codex_app_server_protocol::AskForApproval::Never),
sandbox_policy: Some(codex_app_server_protocol::SandboxPolicy::DangerFullAccess),
model: Some("mock-model".to_string()),
effort: Some(ReasoningEffort::Medium),
summary: Some(ReasoningSummary::Auto),
})
.await?;
timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_response_message(RequestId::Integer(second_turn_id)),
)
.await??;
// Ensure we do NOT receive an ExecCommandApproval request before task completes
timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_notification_message("codex/event/task_complete"),
)
.await??;
drop(server);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn turn_start_updates_sandbox_and_cwd_between_turns_v2() -> Result<()> {
if env::var(CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR).is_ok() {
println!("Skipping v2 sandbox/cwd test due to sandbox network disabled.");
return Ok(());
}
let tmp = TempDir::new()?;
let codex_home = tmp.path().join("codex_home");
std::fs::create_dir(&codex_home)?;
let workspace_root = tmp.path().join("workspace");
std::fs::create_dir(&workspace_root)?;
let first_cwd = workspace_root.join("turn1");
let second_cwd = workspace_root.join("turn2");
std::fs::create_dir(&first_cwd)?;
std::fs::create_dir(&second_cwd)?;
let responses = vec![
create_shell_sse_response(
vec![
"bash".to_string(),
"-lc".to_string(),
"echo first turn".to_string(),
],
None,
Some(5000),
"call-first",
)?,
create_final_assistant_message_sse_response("done first")?,
create_shell_sse_response(
vec![
"bash".to_string(),
"-lc".to_string(),
"echo second turn".to_string(),
],
None,
Some(5000),
"call-second",
)?,
create_final_assistant_message_sse_response("done second")?,
];
let server = create_mock_chat_completions_server(responses).await;
create_config_toml_untrusted(&codex_home, &server.uri())?;
let mut mcp = McpProcess::new(&codex_home).await?;
timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??;
// thread/start
let start_id = mcp
.send_thread_start_request(ThreadStartParams {
model: Some("mock-model".to_string()),
..Default::default()
})
.await?;
let start_resp: JSONRPCResponse = timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_response_message(RequestId::Integer(start_id)),
)
.await??;
let ThreadStartResponse { thread } = to_response::<ThreadStartResponse>(start_resp)?;
// first turn with workspace-write sandbox and first_cwd
let first_turn = mcp
.send_turn_start_request(TurnStartParams {
thread_id: thread.id.clone(),
input: vec![V2UserInput::Text {
text: "first turn".to_string(),
}],
cwd: Some(first_cwd.clone()),
approval_policy: Some(codex_app_server_protocol::AskForApproval::Never),
sandbox_policy: Some(codex_app_server_protocol::SandboxPolicy::WorkspaceWrite {
writable_roots: vec![first_cwd.clone()],
network_access: false,
exclude_tmpdir_env_var: false,
exclude_slash_tmp: false,
}),
model: Some("mock-model".to_string()),
effort: Some(ReasoningEffort::Medium),
summary: Some(ReasoningSummary::Auto),
})
.await?;
timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_response_message(RequestId::Integer(first_turn)),
)
.await??;
timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_notification_message("codex/event/task_complete"),
)
.await??;
// second turn with danger-full-access and second_cwd, ensure exec begins in second_cwd
let second_turn = mcp
.send_turn_start_request(TurnStartParams {
thread_id: thread.id.clone(),
input: vec![V2UserInput::Text {
text: "second turn".to_string(),
}],
cwd: Some(second_cwd.clone()),
approval_policy: Some(codex_app_server_protocol::AskForApproval::Never),
sandbox_policy: Some(codex_app_server_protocol::SandboxPolicy::DangerFullAccess),
model: Some("mock-model".to_string()),
effort: Some(ReasoningEffort::Medium),
summary: Some(ReasoningSummary::Auto),
})
.await?;
timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_response_message(RequestId::Integer(second_turn)),
)
.await??;
let exec_begin_notification = timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_notification_message("codex/event/exec_command_begin"),
)
.await??;
let params = exec_begin_notification
.params
.clone()
.expect("exec_command_begin params");
let event: Event = serde_json::from_value(params).expect("deserialize exec begin event");
let exec_begin = match event.msg {
EventMsg::ExecCommandBegin(exec_begin) => exec_begin,
other => panic!("expected ExecCommandBegin event, got {other:?}"),
};
assert_eq!(exec_begin.cwd, second_cwd);
assert_eq!(
exec_begin.command,
vec![
"bash".to_string(),
"-lc".to_string(),
"echo second turn".to_string()
]
);
timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_notification_message("codex/event/task_complete"),
)
.await??;
drop(server);
Ok(())
}
// Helper to create a config.toml pointing at the mock model server.
fn create_config_toml(codex_home: &Path, server_uri: &str) -> std::io::Result<()> {
let config_toml = codex_home.join("config.toml");
@@ -201,3 +509,26 @@ stream_max_retries = 0
),
)
}
fn create_config_toml_untrusted(codex_home: &Path, server_uri: &str) -> std::io::Result<()> {
let config_toml = codex_home.join("config.toml");
std::fs::write(
config_toml,
format!(
r#"
model = "mock-model"
approval_policy = "untrusted"
sandbox_mode = "danger-full-access"
model_provider = "mock_provider"
[model_providers.mock_provider]
name = "Mock provider for test"
base_url = "{server_uri}/v1"
wire_api = "chat"
request_max_retries = 0
stream_max_retries = 0
"#
),
)
}