From 4ed19b07664d28ef67592ab5d77aa30d13d3aba0 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Tue, 17 Mar 2026 14:37:20 +0000 Subject: [PATCH 01/11] feat: rename to get more explicit close agent (#14935) https://github.com/openai/codex/issues/14907 --- .../core/src/tools/handlers/multi_agents/close_agent.rs | 6 ++++-- codex-rs/core/src/tools/handlers/multi_agents_tests.rs | 4 ++-- codex-rs/core/src/tools/spec.rs | 9 ++++++--- 3 files changed, 12 insertions(+), 7 deletions(-) diff --git a/codex-rs/core/src/tools/handlers/multi_agents/close_agent.rs b/codex-rs/core/src/tools/handlers/multi_agents/close_agent.rs index 7351299940..ead71c7046 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/close_agent.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/close_agent.rs @@ -95,13 +95,15 @@ impl ToolHandler for Handler { .await; result?; - Ok(CloseAgentResult { status }) + Ok(CloseAgentResult { + previous_status: status, + }) } } #[derive(Debug, Deserialize, Serialize)] pub(crate) struct CloseAgentResult { - pub(crate) status: AgentStatus, + pub(crate) previous_status: AgentStatus, } impl ToolOutput for CloseAgentResult { diff --git a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs index 1a46c7de57..99afe8ac25 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -971,7 +971,7 @@ async fn wait_agent_returns_final_status_without_timeout() { } #[tokio::test] -async fn close_agent_submits_shutdown_and_returns_status() { +async fn close_agent_submits_shutdown_and_returns_previous_status() { let (mut session, turn) = make_session_and_context().await; let manager = thread_manager(); session.services.agent_control = manager.agent_control(); @@ -993,7 +993,7 @@ async fn close_agent_submits_shutdown_and_returns_status() { let (content, success) = expect_text_output(output); let result: close_agent::CloseAgentResult = serde_json::from_str(&content).expect("close_agent result should be json"); - assert_eq!(result.status, status_before); + assert_eq!(result.previous_status, status_before); assert_eq!(success, Some(true)); let ops = manager.captured_ops(); diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index 09c2c50d6a..ab5cff7949 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -195,9 +195,12 @@ fn close_agent_output_schema() -> JsonValue { json!({ "type": "object", "properties": { - "status": agent_status_output_schema() + "previous_status": { + "description": "The agent status observed before shutdown was requested.", + "allOf": [agent_status_output_schema()] + } }, - "required": ["status"], + "required": ["previous_status"], "additionalProperties": false }) } @@ -1523,7 +1526,7 @@ fn create_close_agent_tool() -> ToolSpec { ToolSpec::Function(ResponsesApiTool { name: "close_agent".to_string(), - description: "Close an agent when it is no longer needed and return its last known status. Don't keep agents open for too long if they are not needed anymore.".to_string(), + description: "Close an agent when it is no longer needed and return its previous status before shutdown was requested. Don't keep agents open for too long if they are not needed anymore.".to_string(), strict: false, defer_loading: None, parameters: JsonSchema::Object { From ef36d39199c7328899e4f1f6b20a2c9ba5065f83 Mon Sep 17 00:00:00 2001 From: daveaitel-openai Date: Tue, 17 Mar 2026 10:40:14 -0400 Subject: [PATCH 02/11] Fix agent jobs finalization race and reduce status polling churn (#14843) ## Summary - make `report_agent_job_result` atomically transition an item from running to completed while storing `result_json` - remove brittle finalization grace-sleep logic and make finished-item cleanup idempotent - replace blind fixed-interval waiting with status-subscription-based waiting for active worker threads - add state runtime tests for atomic completion and late-report rejection ## Why This addresses the race and polling concerns in #13948 by removing timing-based correctness assumptions and reducing unnecessary status polling churn. ## Validation - `cd codex-rs && just fmt` - `cd codex-rs && cargo test -p codex-state` - `cd codex-rs && cargo test -p codex-core --test all suite::agent_jobs` - `cd codex-rs && cargo test` - fails in an unrelated app-server tracing test: `message_processor::tracing_tests::thread_start_jsonrpc_span_exports_server_span_and_parents_children` timed out waiting for response ## Notes - This PR supersedes #14129 with the same agent-jobs fix on a clean branch from `main`. - The earlier PR branch was stacked on unrelated history, which made the review diff include unrelated commits. Fixes #13948 --- .../core/src/tools/handlers/agent_jobs.rs | 90 +++++++++---- codex-rs/state/src/runtime/agent_jobs.rs | 124 +++++++++++++++++- 2 files changed, 186 insertions(+), 28 deletions(-) diff --git a/codex-rs/core/src/tools/handlers/agent_jobs.rs b/codex-rs/core/src/tools/handlers/agent_jobs.rs index 7c80f9383a..42cb5242d4 100644 --- a/codex-rs/core/src/tools/handlers/agent_jobs.rs +++ b/codex-rs/core/src/tools/handlers/agent_jobs.rs @@ -15,9 +15,12 @@ use crate::tools::registry::ToolHandler; use crate::tools::registry::ToolKind; use async_trait::async_trait; use codex_protocol::ThreadId; +use codex_protocol::protocol::AgentStatus; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use codex_protocol::user_input::UserInput; +use futures::StreamExt; +use futures::stream::FuturesUnordered; use serde::Deserialize; use serde::Serialize; use serde_json::Value; @@ -26,8 +29,10 @@ use std::collections::HashSet; use std::path::Path; use std::path::PathBuf; use std::sync::Arc; +use tokio::sync::watch::Receiver; use tokio::time::Duration; use tokio::time::Instant; +use tokio::time::timeout; use uuid::Uuid; pub struct BatchJobHandler; @@ -103,6 +108,7 @@ struct JobRunnerOptions { struct ActiveJobItem { item_id: String, started_at: Instant, + status_rx: Option>, } struct JobProgressEmitter { @@ -676,6 +682,12 @@ async fn run_agent_job_loop( ActiveJobItem { item_id: item.item_id.clone(), started_at: Instant::now(), + status_rx: session + .services + .agent_control + .subscribe_status(thread_id) + .await + .ok(), }, ); progressed = true; @@ -708,7 +720,7 @@ async fn run_agent_job_loop( break; } if !progressed { - tokio::time::sleep(STATUS_POLL_INTERVAL).await; + wait_for_status_change(&active_items).await; } continue; } @@ -863,6 +875,12 @@ async fn recover_running_items( ActiveJobItem { item_id: item.item_id.clone(), started_at: started_at_from_item(&item), + status_rx: session + .services + .agent_control + .subscribe_status(thread_id) + .await + .ok(), }, ); } @@ -876,13 +894,44 @@ async fn find_finished_threads( ) -> Vec<(ThreadId, String)> { let mut finished = Vec::new(); for (thread_id, item) in active_items { - if is_final(&session.services.agent_control.get_status(*thread_id).await) { + let status = active_item_status(session.as_ref(), *thread_id, item).await; + if is_final(&status) { finished.push((*thread_id, item.item_id.clone())); } } finished } +async fn active_item_status( + session: &Session, + thread_id: ThreadId, + item: &ActiveJobItem, +) -> AgentStatus { + if let Some(status_rx) = item.status_rx.as_ref() + && status_rx.has_changed().is_ok() + { + return status_rx.borrow().clone(); + } + session.services.agent_control.get_status(thread_id).await +} + +async fn wait_for_status_change(active_items: &HashMap) { + let mut waiters = FuturesUnordered::new(); + for item in active_items.values() { + if let Some(status_rx) = item.status_rx.as_ref() { + let mut status_rx = status_rx.clone(); + waiters.push(async move { + let _ = status_rx.changed().await; + }); + } + } + if waiters.is_empty() { + tokio::time::sleep(STATUS_POLL_INTERVAL).await; + return; + } + let _ = timeout(STATUS_POLL_INTERVAL, waiters.next()).await; +} + async fn reap_stale_active_items( session: Arc, db: Arc, @@ -920,37 +969,24 @@ async fn finalize_finished_item( item_id: &str, thread_id: ThreadId, ) -> anyhow::Result<()> { - let mut item = db + let item = db .get_agent_job_item(job_id, item_id) .await? .ok_or_else(|| { anyhow::anyhow!("job item not found for finalization: {job_id}/{item_id}") })?; - if item.result_json.is_none() { - tokio::time::sleep(Duration::from_millis(250)).await; - item = db - .get_agent_job_item(job_id, item_id) - .await? - .ok_or_else(|| { - anyhow::anyhow!("job item not found after grace period: {job_id}/{item_id}") - })?; - } - if item.result_json.is_some() { - if !db.mark_agent_job_item_completed(job_id, item_id).await? { - db.mark_agent_job_item_failed( - job_id, - item_id, - "worker reported result but item could not transition to completed", - ) - .await?; + if matches!(item.status, codex_state::AgentJobItemStatus::Running) { + if item.result_json.is_some() { + let _ = db.mark_agent_job_item_completed(job_id, item_id).await?; + } else { + let _ = db + .mark_agent_job_item_failed( + job_id, + item_id, + "worker finished without calling report_agent_job_result", + ) + .await?; } - } else { - db.mark_agent_job_item_failed( - job_id, - item_id, - "worker finished without calling report_agent_job_result", - ) - .await?; } let _ = session .services diff --git a/codex-rs/state/src/runtime/agent_jobs.rs b/codex-rs/state/src/runtime/agent_jobs.rs index c685605945..3f5526c58d 100644 --- a/codex-rs/state/src/runtime/agent_jobs.rs +++ b/codex-rs/state/src/runtime/agent_jobs.rs @@ -435,10 +435,13 @@ WHERE job_id = ? AND item_id = ? AND status = ? r#" UPDATE agent_job_items SET + status = ?, result_json = ?, reported_at = ?, + completed_at = ?, updated_at = ?, - last_error = NULL + last_error = NULL, + assigned_thread_id = NULL WHERE job_id = ? AND item_id = ? @@ -446,9 +449,11 @@ WHERE AND assigned_thread_id = ? "#, ) + .bind(AgentJobItemStatus::Completed.as_str()) .bind(serialized) .bind(now) .bind(now) + .bind(now) .bind(job_id) .bind(item_id) .bind(AgentJobItemStatus::Running.as_str()) @@ -560,3 +565,120 @@ WHERE job_id = ? }) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::runtime::test_support::unique_temp_dir; + use pretty_assertions::assert_eq; + use serde_json::json; + + async fn create_running_single_item_job( + runtime: &StateRuntime, + ) -> anyhow::Result<(String, String, String)> { + let job_id = "job-1".to_string(); + let item_id = "item-1".to_string(); + let thread_id = "thread-1".to_string(); + runtime + .create_agent_job( + &AgentJobCreateParams { + id: job_id.clone(), + name: "test-job".to_string(), + instruction: "Return a result".to_string(), + auto_export: true, + max_runtime_seconds: None, + output_schema_json: None, + input_headers: vec!["path".to_string()], + input_csv_path: "/tmp/in.csv".to_string(), + output_csv_path: "/tmp/out.csv".to_string(), + }, + &[AgentJobItemCreateParams { + item_id: item_id.clone(), + row_index: 0, + source_id: None, + row_json: json!({"path":"file-1"}), + }], + ) + .await?; + runtime.mark_agent_job_running(job_id.as_str()).await?; + let marked_running = runtime + .mark_agent_job_item_running_with_thread( + job_id.as_str(), + item_id.as_str(), + thread_id.as_str(), + ) + .await?; + assert!(marked_running); + Ok((job_id, item_id, thread_id)) + } + + #[tokio::test] + async fn report_agent_job_item_result_completes_item_atomically() -> anyhow::Result<()> { + let codex_home = unique_temp_dir(); + let runtime = StateRuntime::init(codex_home, "test-provider".to_string()).await?; + let (job_id, item_id, thread_id) = create_running_single_item_job(runtime.as_ref()).await?; + + let accepted = runtime + .report_agent_job_item_result( + job_id.as_str(), + item_id.as_str(), + thread_id.as_str(), + &json!({"ok": true}), + ) + .await?; + assert!(accepted); + + let item = runtime + .get_agent_job_item(job_id.as_str(), item_id.as_str()) + .await? + .expect("job item should exist"); + assert_eq!(item.status, AgentJobItemStatus::Completed); + assert_eq!(item.result_json, Some(json!({"ok": true}))); + assert_eq!(item.assigned_thread_id, None); + assert_eq!(item.last_error, None); + assert!(item.reported_at.is_some()); + assert!(item.completed_at.is_some()); + let progress = runtime.get_agent_job_progress(job_id.as_str()).await?; + assert_eq!( + progress, + AgentJobProgress { + total_items: 1, + pending_items: 0, + running_items: 0, + completed_items: 1, + failed_items: 0, + } + ); + Ok(()) + } + + #[tokio::test] + async fn report_agent_job_item_result_rejects_late_reports() -> anyhow::Result<()> { + let codex_home = unique_temp_dir(); + let runtime = StateRuntime::init(codex_home, "test-provider".to_string()).await?; + let (job_id, item_id, thread_id) = create_running_single_item_job(runtime.as_ref()).await?; + + let marked_failed = runtime + .mark_agent_job_item_failed(job_id.as_str(), item_id.as_str(), "missing report") + .await?; + assert!(marked_failed); + let accepted = runtime + .report_agent_job_item_result( + job_id.as_str(), + item_id.as_str(), + thread_id.as_str(), + &json!({"late": true}), + ) + .await?; + assert!(!accepted); + + let item = runtime + .get_agent_job_item(job_id.as_str(), item_id.as_str()) + .await? + .expect("job item should exist"); + assert_eq!(item.status, AgentJobItemStatus::Failed); + assert_eq!(item.result_json, None); + assert_eq!(item.last_error, Some("missing report".to_string())); + Ok(()) + } +} From e8add54e5dda2fc6f49757aa939378a21b8515e9 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Tue, 17 Mar 2026 16:58:58 +0000 Subject: [PATCH 03/11] feat: show effective model in spawn agent event (#14944) Show effective model after the full config layering for the sub agent --- .../app-server/tests/suite/v2/turn_start.rs | 181 ++++++++++++++++++ codex-rs/core/src/agent/control.rs | 14 ++ .../src/tools/handlers/multi_agents/spawn.rs | 32 +++- codex-rs/protocol/src/protocol.rs | 4 +- 4 files changed, 224 insertions(+), 7 deletions(-) diff --git a/codex-rs/app-server/tests/suite/v2/turn_start.rs b/codex-rs/app-server/tests/suite/v2/turn_start.rs index 441c5558bd..e235ee6b61 100644 --- a/codex-rs/app-server/tests/suite/v2/turn_start.rs +++ b/codex-rs/app-server/tests/suite/v2/turn_start.rs @@ -1857,6 +1857,187 @@ async fn turn_start_emits_spawn_agent_item_with_model_metadata_v2() -> Result<() Ok(()) } +#[tokio::test] +async fn turn_start_emits_spawn_agent_item_with_effective_role_model_metadata_v2() -> Result<()> { + skip_if_no_network!(Ok(())); + + const CHILD_PROMPT: &str = "child: do work"; + const PARENT_PROMPT: &str = "spawn a child and continue"; + const SPAWN_CALL_ID: &str = "spawn-call-1"; + const REQUESTED_MODEL: &str = "gpt-5.1"; + const REQUESTED_REASONING_EFFORT: ReasoningEffort = ReasoningEffort::Low; + const ROLE_MODEL: &str = "gpt-5.1-codex-max"; + const ROLE_REASONING_EFFORT: ReasoningEffort = ReasoningEffort::High; + + let server = responses::start_mock_server().await; + let spawn_args = serde_json::to_string(&json!({ + "message": CHILD_PROMPT, + "agent_type": "custom", + "model": REQUESTED_MODEL, + "reasoning_effort": REQUESTED_REASONING_EFFORT, + }))?; + let _parent_turn = responses::mount_sse_once_match( + &server, + |req: &wiremock::Request| body_contains(req, PARENT_PROMPT), + responses::sse(vec![ + responses::ev_response_created("resp-turn1-1"), + responses::ev_function_call(SPAWN_CALL_ID, "spawn_agent", &spawn_args), + responses::ev_completed("resp-turn1-1"), + ]), + ) + .await; + let _child_turn = responses::mount_sse_once_match( + &server, + |req: &wiremock::Request| { + body_contains(req, CHILD_PROMPT) && !body_contains(req, SPAWN_CALL_ID) + }, + responses::sse(vec![ + responses::ev_response_created("resp-child-1"), + responses::ev_assistant_message("msg-child-1", "child done"), + responses::ev_completed("resp-child-1"), + ]), + ) + .await; + let _parent_follow_up = responses::mount_sse_once_match( + &server, + |req: &wiremock::Request| body_contains(req, SPAWN_CALL_ID), + responses::sse(vec![ + responses::ev_response_created("resp-turn1-2"), + responses::ev_assistant_message("msg-turn1-2", "parent done"), + responses::ev_completed("resp-turn1-2"), + ]), + ) + .await; + + let codex_home = TempDir::new()?; + create_config_toml( + codex_home.path(), + &server.uri(), + "never", + &BTreeMap::from([(Feature::Collab, true)]), + )?; + std::fs::write( + codex_home.path().join("custom-role.toml"), + format!("model = \"{ROLE_MODEL}\"\nmodel_reasoning_effort = \"{ROLE_REASONING_EFFORT}\"\n",), + )?; + let config_path = codex_home.path().join("config.toml"); + let base_config = std::fs::read_to_string(&config_path)?; + std::fs::write( + &config_path, + format!( + r#"{base_config} + +[agents.custom] +description = "Custom role" +config_file = "./custom-role.toml" +"# + ), + )?; + + let mut mcp = McpProcess::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let thread_req = mcp + .send_thread_start_request(ThreadStartParams { + model: Some("gpt-5.2-codex".to_string()), + ..Default::default() + }) + .await?; + let thread_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), + ) + .await??; + let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; + + let turn_req = mcp + .send_turn_start_request(TurnStartParams { + thread_id: thread.id.clone(), + input: vec![V2UserInput::Text { + text: PARENT_PROMPT.to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }) + .await?; + let turn_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), + ) + .await??; + let turn: TurnStartResponse = to_response::(turn_resp)?; + + let spawn_completed = timeout(DEFAULT_READ_TIMEOUT, async { + loop { + let completed_notif = mcp + .read_stream_until_notification_message("item/completed") + .await?; + let completed: ItemCompletedNotification = + serde_json::from_value(completed_notif.params.expect("item/completed params"))?; + if let ThreadItem::CollabAgentToolCall { id, .. } = &completed.item + && id == SPAWN_CALL_ID + { + return Ok::(completed.item); + } + } + }) + .await??; + let ThreadItem::CollabAgentToolCall { + id, + tool, + status, + sender_thread_id, + receiver_thread_ids, + prompt, + model, + reasoning_effort, + agents_states, + } = spawn_completed + else { + unreachable!("loop ensures we break on collab agent tool call items"); + }; + let receiver_thread_id = receiver_thread_ids + .first() + .cloned() + .expect("spawn completion should include child thread id"); + assert_eq!(id, SPAWN_CALL_ID); + assert_eq!(tool, CollabAgentTool::SpawnAgent); + assert_eq!(status, CollabAgentToolCallStatus::Completed); + assert_eq!(sender_thread_id, thread.id); + assert_eq!(receiver_thread_ids, vec![receiver_thread_id.clone()]); + assert_eq!(prompt, Some(CHILD_PROMPT.to_string())); + assert_eq!(model, Some(ROLE_MODEL.to_string())); + assert_eq!(reasoning_effort, Some(ROLE_REASONING_EFFORT)); + assert_eq!( + agents_states, + HashMap::from([( + receiver_thread_id, + CollabAgentState { + status: CollabAgentStatus::PendingInit, + message: None, + }, + )]) + ); + + let turn_completed = timeout(DEFAULT_READ_TIMEOUT, async { + loop { + let turn_completed_notif = mcp + .read_stream_until_notification_message("turn/completed") + .await?; + let turn_completed: TurnCompletedNotification = serde_json::from_value( + turn_completed_notif.params.expect("turn/completed params"), + )?; + if turn_completed.thread_id == thread.id && turn_completed.turn.id == turn.turn.id { + return Ok::(turn_completed); + } + } + }) + .await??; + assert_eq!(turn_completed.thread_id, thread.id); + + Ok(()) +} + #[tokio::test] async fn turn_start_file_change_approval_accept_for_session_persists_v2() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index fc06fcddea..83e6a3a044 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -3,6 +3,7 @@ use crate::agent::guards::Guards; use crate::agent::role::DEFAULT_ROLE_NAME; use crate::agent::role::resolve_role_config; use crate::agent::status::is_final; +use crate::codex_thread::ThreadConfigSnapshot; use crate::error::CodexErr; use crate::error::Result as CodexResult; use crate::find_thread_path_by_id_str; @@ -360,6 +361,19 @@ impl AgentControl { )) } + pub(crate) async fn get_agent_config_snapshot( + &self, + agent_id: ThreadId, + ) -> Option { + let Ok(state) = self.upgrade() else { + return None; + }; + let Ok(thread) = state.get_thread(agent_id).await else { + return None; + }; + Some(thread.config_snapshot().await) + } + /// Subscribe to status updates for `agent_id`, yielding the latest value and changes. pub(crate) async fn subscribe_status( &self, diff --git a/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs b/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs index 26052c6d4f..7a27cd94c8 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs @@ -98,15 +98,37 @@ impl ToolHandler for Handler { ), Err(_) => (None, AgentStatus::NotFound), }; - let (new_agent_nickname, new_agent_role) = match new_thread_id { - Some(thread_id) => session + let agent_snapshot = match new_thread_id { + Some(thread_id) => { + session + .services + .agent_control + .get_agent_config_snapshot(thread_id) + .await + } + None => None, + }; + let (new_agent_nickname, new_agent_role) = match (&agent_snapshot, new_thread_id) { + (Some(snapshot), _) => ( + snapshot.session_source.get_nickname(), + snapshot.session_source.get_agent_role(), + ), + (None, Some(thread_id)) => session .services .agent_control .get_agent_nickname_and_role(thread_id) .await .unwrap_or((None, None)), - None => (None, None), + (None, None) => (None, None), }; + let effective_model = agent_snapshot + .as_ref() + .map(|snapshot| snapshot.model.clone()) + .unwrap_or_else(|| args.model.clone().unwrap_or_default()); + let effective_reasoning_effort = agent_snapshot + .as_ref() + .and_then(|snapshot| snapshot.reasoning_effort) + .unwrap_or(args.reasoning_effort.unwrap_or_default()); let nickname = new_agent_nickname.clone(); session .send_event( @@ -118,8 +140,8 @@ impl ToolHandler for Handler { new_agent_nickname, new_agent_role, prompt, - model: args.model.clone().unwrap_or_default(), - reasoning_effort: args.reasoning_effort.unwrap_or_default(), + model: effective_model, + reasoning_effort: effective_reasoning_effort, status, } .into(), diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 152743b3e1..daf3b7d74a 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -3263,9 +3263,9 @@ pub struct CollabAgentSpawnEndEvent { /// Initial prompt sent to the agent. Can be empty to prevent CoT leaking at the /// beginning. pub prompt: String, - /// Model requested for the spawned agent. + /// Effective model used by the spawned agent after inheritance and role overrides. pub model: String, - /// Reasoning effort requested for the spawned agent. + /// Effective reasoning effort used by the spawned agent after inheritance and role overrides. pub reasoning_effort: ReasoningEffortConfig, /// Last known status of the new agent reported to the sender agent. pub status: AgentStatus, From 6ea041032b500a6f3e8511d225af366d5e53439b Mon Sep 17 00:00:00 2001 From: Owen Lin Date: Tue, 17 Mar 2026 10:07:46 -0700 Subject: [PATCH 04/11] fix(core): prevent hanging turn/start due to websocket warming issues (#14838) ## Description This PR fixes a bad first-turn failure mode in app-server when the startup websocket prewarm hangs. Before this change, `initialize -> thread/start -> turn/start` could sit behind the prewarm for up to five minutes, so the client would not see `turn/started`, and even `turn/interrupt` would block because the turn had not actually started yet. Now, we: - set a (configurable) timeout of 15s for websocket startup time, exposed as `websocket_startup_timeout_ms` in config.toml - `turn/started` is sent immediately on `turn/start` even if the websocket is still connecting - `turn/interrupt` can be used to cancel a turn that is still waiting on the websocket warmup - the turn task will wait for the full 15s websocket warming timeout before falling back ## Why The old behavior made app-server feel stuck at exactly the moment the client expects turn lifecycle events to start flowing. That was especially painful for external clients, because from their point of view the server had accepted the request but then went silent for minutes. ## Configuring the websocket startup timeout Can set it in config.toml like this: ``` [model_providers.openai] supports_websockets = true websocket_connect_timeout_ms = 15000 ``` --- codex-rs/core/config.schema.json | 8 +- codex-rs/core/src/client.rs | 100 +++++--- codex-rs/core/src/codex.rs | 90 ++----- codex-rs/core/src/codex_tests.rs | 103 ++++++++ codex-rs/core/src/config/config_tests.rs | 2 + codex-rs/core/src/config/schema_tests.rs | 9 +- codex-rs/core/src/lib.rs | 1 + codex-rs/core/src/model_provider_info.rs | 14 + .../core/src/model_provider_info_tests.rs | 16 ++ .../core/src/models_manager/manager_tests.rs | 1 + codex-rs/core/src/session_startup_prewarm.rs | 241 ++++++++++++++++++ codex-rs/core/src/state/session.rs | 23 +- codex-rs/core/src/tasks/mod.rs | 1 - codex-rs/core/src/tasks/regular.rs | 76 ++---- codex-rs/core/tests/responses_headers.rs | 3 + codex-rs/core/tests/suite/client.rs | 3 + .../core/tests/suite/client_websockets.rs | 26 +- .../suite/stream_error_allows_next_turn.rs | 1 + .../core/tests/suite/stream_no_completed.rs | 1 + codex-rs/otel/src/metrics/names.rs | 5 + 20 files changed, 548 insertions(+), 176 deletions(-) create mode 100644 codex-rs/core/src/session_startup_prewarm.rs diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 2235315d43..7d3ecdaa01 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -877,6 +877,12 @@ "description": "Whether this provider supports the Responses API WebSocket transport.", "type": "boolean" }, + "websocket_connect_timeout_ms": { + "description": "Maximum time (in milliseconds) to wait for a websocket connection attempt before treating it as failed.", + "format": "uint64", + "minimum": 0.0, + "type": "integer" + }, "wire_api": { "allOf": [ { @@ -2473,4 +2479,4 @@ }, "title": "ConfigToml", "type": "object" -} \ No newline at end of file +} diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 72927b1e88..79fec5dfa1 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -117,6 +117,9 @@ const RESPONSES_WEBSOCKETS_V2_BETA_HEADER_VALUE: &str = "responses_websockets=20 const RESPONSES_ENDPOINT: &str = "/responses"; const RESPONSES_COMPACT_ENDPOINT: &str = "/responses/compact"; const MEMORIES_SUMMARIZE_ENDPOINT: &str = "/memories/trace_summarize"; +#[cfg(test)] +pub(crate) const WEBSOCKET_CONNECT_TIMEOUT: Duration = + Duration::from_millis(crate::model_provider_info::DEFAULT_WEBSOCKET_CONNECT_TIMEOUT_MS); pub fn ws_version_from_features(config: &Config) -> bool { config .features @@ -310,6 +313,27 @@ impl ModelClient { .unwrap_or_else(std::sync::PoisonError::into_inner) = websocket_session; } + pub(crate) fn force_http_fallback( + &self, + session_telemetry: &SessionTelemetry, + model_info: &ModelInfo, + ) -> bool { + let websocket_enabled = self.responses_websocket_enabled(model_info); + let activated = + websocket_enabled && !self.state.disable_websockets.swap(true, Ordering::Relaxed); + if activated { + warn!("falling back to HTTP"); + session_telemetry.counter( + "codex.transport.fallback_to_http", + /*inc*/ 1, + &[("from_wire_api", "responses_websocket")], + ); + } + + self.store_cached_websocket_session(WebsocketSession::default()); + activated + } + /// Compacts the current conversation history using the Compact endpoint. /// /// This is a unary call (no streaming) that returns a new list of @@ -538,15 +562,22 @@ impl ModelClient { auth_context, request_route_telemetry, ); + let websocket_connect_timeout = self.state.provider.websocket_connect_timeout(); let start = Instant::now(); - let result = ApiWebSocketResponsesClient::new(api_provider, api_auth) - .connect( + let result = match tokio::time::timeout( + websocket_connect_timeout, + ApiWebSocketResponsesClient::new(api_provider, api_auth).connect( headers, crate::default_client::default_headers(), turn_state, Some(websocket_telemetry), - ) - .await; + ), + ) + .await + { + Ok(result) => result, + Err(_) => Err(ApiError::Transport(TransportError::Timeout)), + }; let error_message = result.as_ref().err().map(telemetry_api_error_message); let response_debug = result .as_ref() @@ -637,13 +668,12 @@ impl Drop for ModelClientSession { } impl ModelClientSession { - fn activate_http_fallback(&self, websocket_enabled: bool) -> bool { - websocket_enabled - && !self - .client - .state - .disable_websockets - .swap(true, Ordering::Relaxed) + fn reset_websocket_session(&mut self) { + self.websocket_session.connection = None; + self.websocket_session.last_request = None; + self.websocket_session.last_response_rx = None; + self.websocket_session + .set_connection_reused(/*connection_reused*/ false); } fn build_responses_request( @@ -896,7 +926,7 @@ impl ModelClientSession { .turn_state .clone() .unwrap_or_else(|| Arc::clone(&self.turn_state)); - let new_conn = self + let new_conn = match self .client .connect_websocket( session_telemetry, @@ -907,7 +937,16 @@ impl ModelClientSession { auth_context, request_route_telemetry, ) - .await?; + .await + { + Ok(new_conn) => new_conn, + Err(err) => { + if matches!(err, ApiError::Transport(TransportError::Timeout)) { + self.reset_websocket_session(); + } + return Err(err); + } + }; self.websocket_session.connection = Some(new_conn); self.websocket_session .set_connection_reused(/*connection_reused*/ false); @@ -1130,15 +1169,12 @@ impl ModelClientSession { let ws_request = self.prepare_websocket_request(ws_payload, &request); self.websocket_session.last_request = Some(request); - let stream_result = self - .websocket_session - .connection - .as_ref() - .ok_or_else(|| { - map_api_error(ApiError::Stream( - "websocket connection is unavailable".to_string(), - )) - })? + let stream_result = self.websocket_session.connection.as_ref().ok_or_else(|| { + map_api_error(ApiError::Stream( + "websocket connection is unavailable".to_string(), + )) + })?; + let stream_result = stream_result .stream_request(ws_request, self.websocket_session.connection_reused()) .await .map_err(map_api_error)?; @@ -1296,22 +1332,10 @@ impl ModelClientSession { session_telemetry: &SessionTelemetry, model_info: &ModelInfo, ) -> bool { - let websocket_enabled = self.client.responses_websocket_enabled(model_info); - let activated = self.activate_http_fallback(websocket_enabled); - if activated { - warn!("falling back to HTTP"); - session_telemetry.counter( - "codex.transport.fallback_to_http", - /*inc*/ 1, - &[("from_wire_api", "responses_websocket")], - ); - - self.websocket_session.connection = None; - self.websocket_session.last_request = None; - self.websocket_session.last_response_rx = None; - self.websocket_session - .set_connection_reused(/*connection_reused*/ false); - } + let activated = self + .client + .force_http_fallback(session_telemetry, model_info); + self.websocket_session = WebsocketSession::default(); activated } } diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 735fe7ec4d..8ffe1d3bd1 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -105,7 +105,6 @@ use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnContextItem; use codex_protocol::protocol::TurnContextNetworkItem; -use codex_protocol::protocol::TurnStartedEvent; use codex_protocol::protocol::W3cTraceContext; use codex_protocol::request_permissions::PermissionGrantScope; use codex_protocol::request_permissions::RequestPermissionProfile; @@ -267,6 +266,7 @@ use crate::rollout::RolloutRecorderParams; use crate::rollout::map_session_init_error; use crate::rollout::metadata; use crate::rollout::policy::EventPersistenceMode; +use crate::session_startup_prewarm::SessionStartupPrewarmHandle; use crate::shell; use crate::shell_snapshot::ShellSnapshot; use crate::skills::SkillError; @@ -286,7 +286,6 @@ use crate::state::SessionServices; use crate::state::SessionState; use crate::state_db; use crate::tasks::GhostSnapshotTask; -use crate::tasks::RegularTask; use crate::tasks::ReviewTask; use crate::tasks::SessionTask; use crate::tasks::SessionTaskContext; @@ -2411,70 +2410,17 @@ impl Session { .await } - pub(crate) async fn take_startup_regular_task(&self) -> Option { - let startup_regular_task = { - let mut state = self.state.lock().await; - state.take_startup_regular_task() - }; - let startup_regular_task = startup_regular_task?; - match startup_regular_task.await { - Ok(Ok(regular_task)) => Some(regular_task), - Ok(Err(err)) => { - warn!("startup websocket prewarm setup failed: {err:#}"); - None - } - Err(err) => { - warn!("startup websocket prewarm setup join failed: {err}"); - None - } - } - } - - async fn schedule_startup_prewarm(self: &Arc, base_instructions: String) { - let sess = Arc::clone(self); - let startup_regular_task: JoinHandle> = - tokio::spawn( - async move { sess.schedule_startup_prewarm_inner(base_instructions).await }, - ); + pub(crate) async fn set_session_startup_prewarm( + &self, + startup_prewarm: SessionStartupPrewarmHandle, + ) { let mut state = self.state.lock().await; - state.set_startup_regular_task(startup_regular_task); + state.set_session_startup_prewarm(startup_prewarm); } - async fn schedule_startup_prewarm_inner( - self: &Arc, - base_instructions: String, - ) -> CodexResult { - let startup_turn_context = self - .new_default_turn_with_sub_id(INITIAL_SUBMIT_ID.to_owned()) - .await; - let startup_cancellation_token = CancellationToken::new(); - let startup_router = built_tools( - self, - startup_turn_context.as_ref(), - &[], - &HashSet::new(), - /*skills_outcome*/ None, - &startup_cancellation_token, - ) - .await?; - let startup_prompt = build_prompt( - Vec::new(), - startup_router.as_ref(), - startup_turn_context.as_ref(), - BaseInstructions { - text: base_instructions, - }, - ); - let startup_turn_metadata_header = startup_turn_context - .turn_metadata_state - .current_header_value(); - RegularTask::with_startup_prewarm( - self.services.model_client.clone(), - startup_prompt, - startup_turn_context, - startup_turn_metadata_header, - ) - .await + pub(crate) async fn take_session_startup_prewarm(&self) -> Option { + let mut state = self.state.lock().await; + state.take_session_startup_prewarm() } pub(crate) async fn get_config(&self) -> std::sync::Arc { @@ -4553,9 +4499,12 @@ mod handlers { { sess.refresh_mcp_servers_if_requested(¤t_context) .await; - let regular_task = sess.take_startup_regular_task().await.unwrap_or_default(); - sess.spawn_task(Arc::clone(¤t_context), items, regular_task) - .await; + sess.spawn_task( + Arc::clone(¤t_context), + items, + crate::tasks::RegularTask::new(), + ) + .await; } } @@ -5485,13 +5434,6 @@ pub(crate) async fn run_turn( let model_info = turn_context.model_info.clone(); let auto_compact_limit = model_info.auto_compact_token_limit().unwrap_or(i64::MAX); - - let event = EventMsg::TurnStarted(TurnStartedEvent { - turn_id: turn_context.sub_id.clone(), - model_context_window: turn_context.model_context_window(), - collaboration_mode_kind: turn_context.collaboration_mode.mode, - }); - sess.send_event(&turn_context, event).await; // TODO(ccunningham): Pre-turn compaction runs before context updates and the // new user message are recorded. Estimate pending incoming items (context // diffs/full reinjection + user input) and trigger compaction preemptively @@ -6236,7 +6178,7 @@ fn codex_apps_connector_id(tool: &crate::mcp_connection_manager::ToolInfo) -> Op tool.connector_id.as_deref() } -fn build_prompt( +pub(crate) fn build_prompt( input: Vec, router: &ToolRouter, turn_context: &TurnContext, diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index 34ed7bcd63..fab591db7d 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -39,6 +39,7 @@ use crate::protocol::TokenCountEvent; use crate::protocol::TokenUsage; use crate::protocol::TokenUsageInfo; use crate::protocol::TurnCompleteEvent; +use crate::protocol::TurnStartedEvent; use crate::protocol::UserMessageEvent; use crate::rollout::policy::EventPersistenceMode; use crate::rollout::recorder::RolloutRecorder; @@ -142,6 +143,108 @@ fn skill_message(text: &str) -> ResponseItem { } } +#[tokio::test] +async fn regular_turn_emits_turn_started_without_waiting_for_startup_prewarm() { + let (sess, tc, rx) = make_session_and_context_with_rx().await; + let (_tx, startup_prewarm_rx) = tokio::sync::oneshot::channel::<()>(); + let handle = tokio::spawn(async move { + let _ = startup_prewarm_rx.await; + Ok(test_model_client_session()) + }); + + sess.set_session_startup_prewarm( + crate::session_startup_prewarm::SessionStartupPrewarmHandle::new( + handle, + std::time::Instant::now(), + crate::client::WEBSOCKET_CONNECT_TIMEOUT, + ), + ) + .await; + sess.spawn_task( + Arc::clone(&tc), + Vec::new(), + crate::tasks::RegularTask::new(), + ) + .await; + + let first = tokio::time::timeout(std::time::Duration::from_millis(200), rx.recv()) + .await + .expect("expected turn started event without waiting for startup prewarm") + .expect("channel open"); + assert!(matches!( + first.msg, + EventMsg::TurnStarted(TurnStartedEvent { turn_id, .. }) if turn_id == tc.sub_id + )); + + sess.abort_all_tasks(TurnAbortReason::Interrupted).await; +} + +#[tokio::test] +async fn interrupting_regular_turn_waiting_on_startup_prewarm_emits_turn_aborted() { + let (sess, tc, rx) = make_session_and_context_with_rx().await; + let (_tx, startup_prewarm_rx) = tokio::sync::oneshot::channel::<()>(); + let handle = tokio::spawn(async move { + let _ = startup_prewarm_rx.await; + Ok(test_model_client_session()) + }); + + sess.set_session_startup_prewarm( + crate::session_startup_prewarm::SessionStartupPrewarmHandle::new( + handle, + std::time::Instant::now(), + crate::client::WEBSOCKET_CONNECT_TIMEOUT, + ), + ) + .await; + sess.spawn_task( + Arc::clone(&tc), + Vec::new(), + crate::tasks::RegularTask::new(), + ) + .await; + + let first = tokio::time::timeout(std::time::Duration::from_millis(200), rx.recv()) + .await + .expect("expected turn started event without waiting for startup prewarm") + .expect("channel open"); + assert!(matches!( + first.msg, + EventMsg::TurnStarted(TurnStartedEvent { turn_id, .. }) if turn_id == tc.sub_id + )); + + sess.abort_all_tasks(TurnAbortReason::Interrupted).await; + + let second = tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv()) + .await + .expect("expected turn aborted event") + .expect("channel open"); + assert!(matches!( + second.msg, + EventMsg::TurnAborted(crate::protocol::TurnAbortedEvent { + turn_id: Some(turn_id), + reason: TurnAbortReason::Interrupted, + }) if turn_id == tc.sub_id + )); +} + +fn test_model_client_session() -> crate::client::ModelClientSession { + crate::client::ModelClient::new( + None, + ThreadId::try_from("00000000-0000-4000-8000-000000000001") + .expect("test thread id should be valid"), + crate::model_provider_info::ModelProviderInfo::create_openai_provider( + /* base_url */ None, + ), + codex_protocol::protocol::SessionSource::Exec, + None, + false, + false, + false, + None, + ) + .new_session() +} + fn developer_input_texts(items: &[ResponseItem]) -> Vec<&str> { items .iter() diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index 7e82bd7da9..c6372de258 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -4076,6 +4076,7 @@ wire_api = "responses" request_max_retries = 4 # retry failed HTTP requests stream_max_retries = 10 # retry dropped SSE streams stream_idle_timeout_ms = 300000 # 5m idle timeout +websocket_connect_timeout_ms = 15000 [profiles.o3] model = "o3" @@ -4130,6 +4131,7 @@ model_verbosity = "high" request_max_retries: Some(4), stream_max_retries: Some(10), stream_idle_timeout_ms: Some(300_000), + websocket_connect_timeout_ms: Some(15_000), requires_openai_auth: false, supports_websockets: false, }; diff --git a/codex-rs/core/src/config/schema_tests.rs b/codex-rs/core/src/config/schema_tests.rs index 6205d43f40..31fabd64bd 100644 --- a/codex-rs/core/src/config/schema_tests.rs +++ b/codex-rs/core/src/config/schema_tests.rs @@ -6,6 +6,10 @@ use pretty_assertions::assert_eq; use similar::TextDiff; use tempfile::TempDir; +fn trim_single_trailing_newline(contents: &str) -> &str { + contents.strip_suffix('\n').unwrap_or(contents) +} + #[test] fn config_schema_matches_fixture() { let fixture_path = codex_utils_cargo_bin::find_resource!("config.schema.json") @@ -40,9 +44,12 @@ Run `just write-config-schema` to overwrite with your changes.\n\n{diff}" std::fs::read_to_string(&tmp_path).expect("read back config schema from temp path"); #[cfg(windows)] let fixture = fixture.replace("\r\n", "\n"); + #[cfg(windows)] + let tmp_contents = tmp_contents.replace("\r\n", "\n"); assert_eq!( - fixture, tmp_contents, + trim_single_trailing_newline(&fixture), + trim_single_trailing_newline(&tmp_contents), "fixture should match exactly with generated schema" ); } diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 2056f0dfe0..e02a346545 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -69,6 +69,7 @@ pub mod plugins; mod sandbox_tags; pub mod sandboxing; mod session_prefix; +mod session_startup_prewarm; mod shell_detect; mod stream_events_utils; pub mod test_support; diff --git a/codex-rs/core/src/model_provider_info.rs b/codex-rs/core/src/model_provider_info.rs index be7a38d27d..737a47780d 100644 --- a/codex-rs/core/src/model_provider_info.rs +++ b/codex-rs/core/src/model_provider_info.rs @@ -22,6 +22,7 @@ use std::time::Duration; const DEFAULT_STREAM_IDLE_TIMEOUT_MS: u64 = 300_000; const DEFAULT_STREAM_MAX_RETRIES: u64 = 5; const DEFAULT_REQUEST_MAX_RETRIES: u64 = 4; +pub(crate) const DEFAULT_WEBSOCKET_CONNECT_TIMEOUT_MS: u64 = 15_000; /// Hard cap for user-configured `stream_max_retries`. const MAX_STREAM_MAX_RETRIES: u64 = 100; /// Hard cap for user-configured `request_max_retries`. @@ -112,6 +113,10 @@ pub struct ModelProviderInfo { /// the connection as lost. pub stream_idle_timeout_ms: Option, + /// Maximum time (in milliseconds) to wait for a websocket connection attempt before treating + /// it as failed. + pub websocket_connect_timeout_ms: Option, + /// Does this provider require an OpenAI API Key or ChatGPT login token? If true, /// user is presented with login screen on first run, and login preference and token/key /// are stored in auth.json. If false (which is the default), login screen is skipped, @@ -227,6 +232,13 @@ impl ModelProviderInfo { .unwrap_or(Duration::from_millis(DEFAULT_STREAM_IDLE_TIMEOUT_MS)) } + /// Effective timeout for websocket connect attempts. + pub fn websocket_connect_timeout(&self) -> Duration { + self.websocket_connect_timeout_ms + .map(Duration::from_millis) + .unwrap_or(Duration::from_millis(DEFAULT_WEBSOCKET_CONNECT_TIMEOUT_MS)) + } + pub fn create_openai_provider(base_url: Option) -> ModelProviderInfo { ModelProviderInfo { name: OPENAI_PROVIDER_NAME.into(), @@ -256,6 +268,7 @@ impl ModelProviderInfo { request_max_retries: None, stream_max_retries: None, stream_idle_timeout_ms: None, + websocket_connect_timeout_ms: None, requires_openai_auth: true, supports_websockets: true, } @@ -332,6 +345,7 @@ pub fn create_oss_provider_with_base_url(base_url: &str, wire_api: WireApi) -> M request_max_retries: None, stream_max_retries: None, stream_idle_timeout_ms: None, + websocket_connect_timeout_ms: None, requires_openai_auth: false, supports_websockets: false, } diff --git a/codex-rs/core/src/model_provider_info_tests.rs b/codex-rs/core/src/model_provider_info_tests.rs index e6d5cea36b..a5309117ae 100644 --- a/codex-rs/core/src/model_provider_info_tests.rs +++ b/codex-rs/core/src/model_provider_info_tests.rs @@ -20,6 +20,7 @@ base_url = "http://localhost:11434/v1" request_max_retries: None, stream_max_retries: None, stream_idle_timeout_ms: None, + websocket_connect_timeout_ms: None, requires_openai_auth: false, supports_websockets: false, }; @@ -51,6 +52,7 @@ query_params = { api-version = "2025-04-01-preview" } request_max_retries: None, stream_max_retries: None, stream_idle_timeout_ms: None, + websocket_connect_timeout_ms: None, requires_openai_auth: false, supports_websockets: false, }; @@ -85,6 +87,7 @@ env_http_headers = { "X-Example-Env-Header" = "EXAMPLE_ENV_VAR" } request_max_retries: None, stream_max_retries: None, stream_idle_timeout_ms: None, + websocket_connect_timeout_ms: None, requires_openai_auth: false, supports_websockets: false, }; @@ -105,3 +108,16 @@ wire_api = "chat" let err = toml::from_str::(provider_toml).unwrap_err(); assert!(err.to_string().contains(CHAT_WIRE_API_REMOVED_ERROR)); } + +#[test] +fn test_deserialize_websocket_connect_timeout() { + let provider_toml = r#" +name = "OpenAI" +base_url = "https://api.openai.com/v1" +websocket_connect_timeout_ms = 15000 +supports_websockets = true + "#; + + let provider: ModelProviderInfo = toml::from_str(provider_toml).unwrap(); + assert_eq!(provider.websocket_connect_timeout_ms, Some(15_000)); +} diff --git a/codex-rs/core/src/models_manager/manager_tests.rs b/codex-rs/core/src/models_manager/manager_tests.rs index 6981d6d799..da42f2c859 100644 --- a/codex-rs/core/src/models_manager/manager_tests.rs +++ b/codex-rs/core/src/models_manager/manager_tests.rs @@ -71,6 +71,7 @@ fn provider_for(base_url: String) -> ModelProviderInfo { request_max_retries: Some(0), stream_max_retries: Some(0), stream_idle_timeout_ms: Some(5_000), + websocket_connect_timeout_ms: None, requires_openai_auth: false, supports_websockets: false, } diff --git a/codex-rs/core/src/session_startup_prewarm.rs b/codex-rs/core/src/session_startup_prewarm.rs new file mode 100644 index 0000000000..326d864f1e --- /dev/null +++ b/codex-rs/core/src/session_startup_prewarm.rs @@ -0,0 +1,241 @@ +use std::collections::HashSet; +use std::sync::Arc; +use std::time::Duration; +use std::time::Instant; + +use tokio::task::JoinHandle; +use tokio_util::sync::CancellationToken; +use tracing::info; +use tracing::warn; + +use crate::client::ModelClientSession; +use crate::codex::INITIAL_SUBMIT_ID; +use crate::codex::Session; +use crate::codex::build_prompt; +use crate::codex::built_tools; +use crate::error::Result as CodexResult; +use codex_otel::SessionTelemetry; +use codex_otel::metrics::names::STARTUP_PREWARM_AGE_AT_FIRST_TURN_METRIC; +use codex_otel::metrics::names::STARTUP_PREWARM_DURATION_METRIC; +use codex_protocol::models::BaseInstructions; + +pub(crate) struct SessionStartupPrewarmHandle { + task: JoinHandle>, + started_at: Instant, + timeout: Duration, +} + +pub(crate) enum SessionStartupPrewarmResolution { + Cancelled, + Ready(Box), + Unavailable { + status: &'static str, + prewarm_duration: Option, + }, +} + +impl SessionStartupPrewarmHandle { + pub(crate) fn new( + task: JoinHandle>, + started_at: Instant, + timeout: Duration, + ) -> Self { + Self { + task, + started_at, + timeout, + } + } + + async fn resolve( + self, + session_telemetry: &SessionTelemetry, + cancellation_token: &CancellationToken, + ) -> SessionStartupPrewarmResolution { + let Self { + mut task, + started_at, + timeout, + } = self; + let age_at_first_turn = started_at.elapsed(); + let remaining = timeout.saturating_sub(age_at_first_turn); + + let resolution = if task.is_finished() { + Self::resolution_from_join_result(task.await, started_at) + } else { + match tokio::select! { + _ = cancellation_token.cancelled() => None, + result = tokio::time::timeout(remaining, &mut task) => Some(result), + } { + Some(Ok(result)) => Self::resolution_from_join_result(result, started_at), + Some(Err(_elapsed)) => { + task.abort(); + info!("startup websocket prewarm timed out before the first turn could use it"); + SessionStartupPrewarmResolution::Unavailable { + status: "timed_out", + prewarm_duration: Some(started_at.elapsed()), + } + } + None => { + task.abort(); + session_telemetry.record_duration( + STARTUP_PREWARM_AGE_AT_FIRST_TURN_METRIC, + age_at_first_turn, + &[("status", "cancelled")], + ); + session_telemetry.record_duration( + STARTUP_PREWARM_DURATION_METRIC, + started_at.elapsed(), + &[("status", "cancelled")], + ); + return SessionStartupPrewarmResolution::Cancelled; + } + } + }; + + match resolution { + SessionStartupPrewarmResolution::Cancelled => { + SessionStartupPrewarmResolution::Cancelled + } + SessionStartupPrewarmResolution::Ready(prewarmed_session) => { + session_telemetry.record_duration( + STARTUP_PREWARM_AGE_AT_FIRST_TURN_METRIC, + age_at_first_turn, + &[("status", "consumed")], + ); + SessionStartupPrewarmResolution::Ready(prewarmed_session) + } + SessionStartupPrewarmResolution::Unavailable { + status, + prewarm_duration, + } => { + session_telemetry.record_duration( + STARTUP_PREWARM_AGE_AT_FIRST_TURN_METRIC, + age_at_first_turn, + &[("status", status)], + ); + if let Some(prewarm_duration) = prewarm_duration { + session_telemetry.record_duration( + STARTUP_PREWARM_DURATION_METRIC, + prewarm_duration, + &[("status", status)], + ); + } + SessionStartupPrewarmResolution::Unavailable { + status, + prewarm_duration, + } + } + } + } + + fn resolution_from_join_result( + result: std::result::Result, tokio::task::JoinError>, + started_at: Instant, + ) -> SessionStartupPrewarmResolution { + match result { + Ok(Ok(prewarmed_session)) => { + SessionStartupPrewarmResolution::Ready(Box::new(prewarmed_session)) + } + Ok(Err(err)) => { + warn!("startup websocket prewarm setup failed: {err:#}"); + SessionStartupPrewarmResolution::Unavailable { + status: "failed", + prewarm_duration: None, + } + } + Err(err) => { + warn!("startup websocket prewarm setup join failed: {err}"); + SessionStartupPrewarmResolution::Unavailable { + status: "join_failed", + prewarm_duration: Some(started_at.elapsed()), + } + } + } + } +} + +impl Session { + pub(crate) async fn schedule_startup_prewarm(self: &Arc, base_instructions: String) { + let session_telemetry = self.services.session_telemetry.clone(); + let websocket_connect_timeout = self.provider().await.websocket_connect_timeout(); + let started_at = Instant::now(); + let startup_prewarm_session = Arc::clone(self); + let startup_prewarm = tokio::spawn(async move { + let result = + schedule_startup_prewarm_inner(startup_prewarm_session, base_instructions).await; + let status = if result.is_ok() { "ready" } else { "failed" }; + session_telemetry.record_duration( + STARTUP_PREWARM_DURATION_METRIC, + started_at.elapsed(), + &[("status", status)], + ); + result + }); + self.set_session_startup_prewarm(SessionStartupPrewarmHandle::new( + startup_prewarm, + started_at, + websocket_connect_timeout, + )) + .await; + } + + pub(crate) async fn consume_startup_prewarm_for_regular_turn( + &self, + cancellation_token: &CancellationToken, + ) -> SessionStartupPrewarmResolution { + let Some(startup_prewarm) = self.take_session_startup_prewarm().await else { + return SessionStartupPrewarmResolution::Unavailable { + status: "not_scheduled", + prewarm_duration: None, + }; + }; + startup_prewarm + .resolve(&self.services.session_telemetry, cancellation_token) + .await + } +} + +async fn schedule_startup_prewarm_inner( + session: Arc, + base_instructions: String, +) -> CodexResult { + let startup_turn_context = session + .new_default_turn_with_sub_id(INITIAL_SUBMIT_ID.to_owned()) + .await; + let startup_cancellation_token = CancellationToken::new(); + let startup_router = built_tools( + session.as_ref(), + startup_turn_context.as_ref(), + &[], + &HashSet::new(), + /*skills_outcome*/ None, + &startup_cancellation_token, + ) + .await?; + let startup_prompt = build_prompt( + Vec::new(), + startup_router.as_ref(), + startup_turn_context.as_ref(), + BaseInstructions { + text: base_instructions, + }, + ); + let startup_turn_metadata_header = startup_turn_context + .turn_metadata_state + .current_header_value(); + let mut client_session = session.services.model_client.new_session(); + client_session + .prewarm_websocket( + &startup_prompt, + &startup_turn_context.model_info, + &startup_turn_context.session_telemetry, + startup_turn_context.reasoning_effort, + startup_turn_context.reasoning_summary, + startup_turn_context.config.service_tier, + startup_turn_metadata_header.as_deref(), + ) + .await?; + + Ok(client_session) +} diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index 40faa4b856..563e8b3403 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -4,17 +4,15 @@ use codex_protocol::models::PermissionProfile; use codex_protocol::models::ResponseItem; use std::collections::HashMap; use std::collections::HashSet; -use tokio::task::JoinHandle; use crate::codex::PreviousTurnSettings; use crate::codex::SessionConfiguration; use crate::context_manager::ContextManager; -use crate::error::Result as CodexResult; use crate::protocol::RateLimitSnapshot; use crate::protocol::TokenUsage; use crate::protocol::TokenUsageInfo; use crate::sandboxing::merge_permission_profiles; -use crate::tasks::RegularTask; +use crate::session_startup_prewarm::SessionStartupPrewarmHandle; use crate::truncate::TruncationPolicy; use codex_protocol::protocol::TurnContextItem; @@ -30,8 +28,8 @@ pub(crate) struct SessionState { /// model/realtime handling on subsequent regular turns (including full-context /// reinjection after resume or `/compact`). previous_turn_settings: Option, - /// Startup regular task pre-created during session initialization. - pub(crate) startup_regular_task: Option>>, + /// Startup prewarmed session prepared during session initialization. + pub(crate) startup_prewarm: Option, pub(crate) active_connector_selection: HashSet, pub(crate) pending_session_start_source: Option, granted_permissions: Option, @@ -49,7 +47,7 @@ impl SessionState { dependency_env: HashMap::new(), mcp_dependency_prompted: HashSet::new(), previous_turn_settings: None, - startup_regular_task: None, + startup_prewarm: None, active_connector_selection: HashSet::new(), pending_session_start_source: None, granted_permissions: None, @@ -165,14 +163,15 @@ impl SessionState { self.dependency_env.clone() } - pub(crate) fn set_startup_regular_task(&mut self, task: JoinHandle>) { - self.startup_regular_task = Some(task); + pub(crate) fn set_session_startup_prewarm( + &mut self, + startup_prewarm: SessionStartupPrewarmHandle, + ) { + self.startup_prewarm = Some(startup_prewarm); } - pub(crate) fn take_startup_regular_task( - &mut self, - ) -> Option>> { - self.startup_regular_task.take() + pub(crate) fn take_session_startup_prewarm(&mut self) -> Option { + self.startup_prewarm.take() } // Adds connector IDs to the active set and returns the merged selection. diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 2089561199..c237af4d11 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -384,7 +384,6 @@ impl Session { turn.add_task(task); *active = Some(turn); } - async fn take_active_turn(&self) -> Option { let mut active = self.active_turn.lock().await; active.take() diff --git a/codex-rs/core/src/tasks/regular.rs b/codex-rs/core/src/tasks/regular.rs index f1851b9347..6deb9abdb6 100644 --- a/codex-rs/core/src/tasks/regular.rs +++ b/codex-rs/core/src/tasks/regular.rs @@ -1,64 +1,27 @@ use std::sync::Arc; -use std::sync::Mutex; -use crate::client::ModelClient; -use crate::client::ModelClientSession; -use crate::client_common::Prompt; +use async_trait::async_trait; +use tokio_util::sync::CancellationToken; + use crate::codex::TurnContext; use crate::codex::run_turn; -use crate::error::Result as CodexResult; +use crate::protocol::EventMsg; +use crate::protocol::TurnStartedEvent; +use crate::session_startup_prewarm::SessionStartupPrewarmResolution; use crate::state::TaskKind; -use async_trait::async_trait; use codex_protocol::user_input::UserInput; -use tokio_util::sync::CancellationToken; use tracing::Instrument; use tracing::trace_span; use super::SessionTask; use super::SessionTaskContext; -pub(crate) struct RegularTask { - prewarmed_session: Mutex>, -} - -impl Default for RegularTask { - fn default() -> Self { - Self { - prewarmed_session: Mutex::new(None), - } - } -} +#[derive(Default)] +pub(crate) struct RegularTask; impl RegularTask { - pub(crate) async fn with_startup_prewarm( - model_client: ModelClient, - prompt: Prompt, - turn_context: Arc, - turn_metadata_header: Option, - ) -> CodexResult { - let mut client_session = model_client.new_session(); - client_session - .prewarm_websocket( - &prompt, - &turn_context.model_info, - &turn_context.session_telemetry, - turn_context.reasoning_effort, - turn_context.reasoning_summary, - turn_context.config.service_tier, - turn_metadata_header.as_deref(), - ) - .await?; - - Ok(Self { - prewarmed_session: Mutex::new(Some(client_session)), - }) - } - - async fn take_prewarmed_session(&self) -> Option { - self.prewarmed_session - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .take() + pub(crate) fn new() -> Self { + Self } } @@ -81,8 +44,25 @@ impl SessionTask for RegularTask { ) -> Option { let sess = session.clone_session(); let run_turn_span = trace_span!("run_turn"); + // Regular turns emit `TurnStarted` inline so first-turn lifecycle does + // not wait on startup prewarm resolution. + let event = EventMsg::TurnStarted(TurnStartedEvent { + turn_id: ctx.sub_id.clone(), + model_context_window: ctx.model_context_window(), + collaboration_mode_kind: ctx.collaboration_mode.mode, + }); + sess.send_event(ctx.as_ref(), event).await; sess.set_server_reasoning_included(/*included*/ false).await; - let prewarmed_client_session = self.take_prewarmed_session().await; + let prewarmed_client_session = match sess + .consume_startup_prewarm_for_regular_turn(&cancellation_token) + .await + { + SessionStartupPrewarmResolution::Cancelled => return None, + SessionStartupPrewarmResolution::Unavailable { .. } => None, + SessionStartupPrewarmResolution::Ready(prewarmed_client_session) => { + Some(*prewarmed_client_session) + } + }; run_turn( sess, ctx, diff --git a/codex-rs/core/tests/responses_headers.rs b/codex-rs/core/tests/responses_headers.rs index d5376fcd7d..d1c73f39d9 100644 --- a/codex-rs/core/tests/responses_headers.rs +++ b/codex-rs/core/tests/responses_headers.rs @@ -53,6 +53,7 @@ async fn responses_stream_includes_subagent_header_on_review() { request_max_retries: Some(0), stream_max_retries: Some(0), stream_idle_timeout_ms: Some(5_000), + websocket_connect_timeout_ms: None, requires_openai_auth: false, supports_websockets: false, }; @@ -165,6 +166,7 @@ async fn responses_stream_includes_subagent_header_on_other() { request_max_retries: Some(0), stream_max_retries: Some(0), stream_idle_timeout_ms: Some(5_000), + websocket_connect_timeout_ms: None, requires_openai_auth: false, supports_websockets: false, }; @@ -272,6 +274,7 @@ async fn responses_respects_model_info_overrides_from_config() { request_max_retries: Some(0), stream_max_retries: Some(0), stream_idle_timeout_ms: Some(5_000), + websocket_connect_timeout_ms: None, requires_openai_auth: false, supports_websockets: false, }; diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index 2fc5f8b9d1..7d785694dc 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -1792,6 +1792,7 @@ async fn azure_responses_request_includes_store_and_reasoning_ids() { request_max_retries: Some(0), stream_max_retries: Some(0), stream_idle_timeout_ms: Some(5_000), + websocket_connect_timeout_ms: None, requires_openai_auth: false, supports_websockets: false, }; @@ -2393,6 +2394,7 @@ async fn azure_overrides_assign_properties_used_for_responses_url() { request_max_retries: None, stream_max_retries: None, stream_idle_timeout_ms: None, + websocket_connect_timeout_ms: None, requires_openai_auth: false, supports_websockets: false, }; @@ -2477,6 +2479,7 @@ async fn env_var_overrides_loaded_auth() { request_max_retries: None, stream_max_retries: None, stream_idle_timeout_ms: None, + websocket_connect_timeout_ms: None, requires_openai_auth: false, supports_websockets: false, }; diff --git a/codex-rs/core/tests/suite/client_websockets.rs b/codex-rs/core/tests/suite/client_websockets.rs index 2c7b4d48e1..c2dd16f8b5 100755 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/codex-rs/core/tests/suite/client_websockets.rs @@ -1500,6 +1500,13 @@ fn prompt_with_input_and_instructions(input: Vec, instructions: &s } fn websocket_provider(server: &WebSocketTestServer) -> ModelProviderInfo { + websocket_provider_with_connect_timeout(server, None) +} + +fn websocket_provider_with_connect_timeout( + server: &WebSocketTestServer, + websocket_connect_timeout_ms: Option, +) -> ModelProviderInfo { ModelProviderInfo { name: "mock-ws".into(), base_url: Some(format!("{}/v1", server.uri())), @@ -1513,6 +1520,7 @@ fn websocket_provider(server: &WebSocketTestServer) -> ModelProviderInfo { request_max_retries: Some(0), stream_max_retries: Some(0), stream_idle_timeout_ms: Some(5_000), + websocket_connect_timeout_ms, requires_openai_auth: false, supports_websockets: true, } @@ -1543,7 +1551,23 @@ async fn websocket_harness_with_options( websocket_v2_enabled: bool, prefer_websockets: bool, ) -> WebsocketTestHarness { - let provider = websocket_provider(server); + websocket_harness_with_provider_options( + websocket_provider(server), + runtime_metrics_enabled, + websocket_enabled, + websocket_v2_enabled, + prefer_websockets, + ) + .await +} + +async fn websocket_harness_with_provider_options( + provider: ModelProviderInfo, + runtime_metrics_enabled: bool, + websocket_enabled: bool, + websocket_v2_enabled: bool, + prefer_websockets: bool, +) -> WebsocketTestHarness { let codex_home = TempDir::new().unwrap(); let mut config = load_default_config_for_test(&codex_home).await; config.model = Some(MODEL.to_string()); diff --git a/codex-rs/core/tests/suite/stream_error_allows_next_turn.rs b/codex-rs/core/tests/suite/stream_error_allows_next_turn.rs index a8d1b37950..23ffc4afb0 100644 --- a/codex-rs/core/tests/suite/stream_error_allows_next_turn.rs +++ b/codex-rs/core/tests/suite/stream_error_allows_next_turn.rs @@ -76,6 +76,7 @@ async fn continue_after_stream_error() { request_max_retries: Some(1), stream_max_retries: Some(1), stream_idle_timeout_ms: Some(2_000), + websocket_connect_timeout_ms: None, requires_openai_auth: false, supports_websockets: false, }; diff --git a/codex-rs/core/tests/suite/stream_no_completed.rs b/codex-rs/core/tests/suite/stream_no_completed.rs index e6fc7ee8cb..5d1b214811 100644 --- a/codex-rs/core/tests/suite/stream_no_completed.rs +++ b/codex-rs/core/tests/suite/stream_no_completed.rs @@ -61,6 +61,7 @@ async fn retries_on_early_close() { request_max_retries: Some(0), stream_max_retries: Some(1), stream_idle_timeout_ms: Some(2000), + websocket_connect_timeout_ms: None, requires_openai_auth: false, supports_websockets: false, }; diff --git a/codex-rs/otel/src/metrics/names.rs b/codex-rs/otel/src/metrics/names.rs index 5063001f2c..569cdc8256 100644 --- a/codex-rs/otel/src/metrics/names.rs +++ b/codex-rs/otel/src/metrics/names.rs @@ -25,4 +25,9 @@ pub const TURN_TTFM_DURATION_METRIC: &str = "codex.turn.ttfm.duration_ms"; pub const TURN_NETWORK_PROXY_METRIC: &str = "codex.turn.network_proxy"; pub const TURN_TOOL_CALL_METRIC: &str = "codex.turn.tool.call"; pub const TURN_TOKEN_USAGE_METRIC: &str = "codex.turn.token_usage"; +/// Total runtime of a startup prewarm attempt until it completes, tagged by final status. +pub const STARTUP_PREWARM_DURATION_METRIC: &str = "codex.startup_prewarm.duration_ms"; +/// Age of the startup prewarm attempt when the first real turn resolves it, tagged by outcome. +pub const STARTUP_PREWARM_AGE_AT_FIRST_TURN_METRIC: &str = + "codex.startup_prewarm.age_at_first_turn_ms"; pub const THREAD_STARTED_METRIC: &str = "codex.thread.started"; From 8e258eb3f57a42477b5811a54321263185136a6a Mon Sep 17 00:00:00 2001 From: Shijie Rao Date: Tue, 17 Mar 2026 10:14:34 -0700 Subject: [PATCH 05/11] Feat: CXA-1831 Persist latest model and reasoning effort in sqlite (#14859) ### Summary The goal is for us to get the latest turn model and reasoning effort on thread/resume is no override is provided on the thread/resume func call. This is the part 1 which we write the model and reasoning effort for a thread to the sqlite db and there will be a followup PR to consume the two new fields on thread/resume. [part 2 PR is currently WIP](https://github.com/openai/codex/pull/14888) and this one can be merged independently. --- codex-rs/core/src/realtime_context_tests.rs | 2 + codex-rs/protocol/src/openai_models.rs | 24 ++++ .../0020_threads_model_reasoning_effort.sql | 2 + codex-rs/state/src/extract.rs | 72 +++++++++++- codex-rs/state/src/model/thread_metadata.rs | 106 ++++++++++++++++++ codex-rs/state/src/runtime/memories.rs | 2 + codex-rs/state/src/runtime/test_support.rs | 4 + codex-rs/state/src/runtime/threads.rs | 28 ++++- 8 files changed, 237 insertions(+), 3 deletions(-) create mode 100644 codex-rs/state/migrations/0020_threads_model_reasoning_effort.sql diff --git a/codex-rs/core/src/realtime_context_tests.rs b/codex-rs/core/src/realtime_context_tests.rs index 1e23b73b32..a04b771396 100644 --- a/codex-rs/core/src/realtime_context_tests.rs +++ b/codex-rs/core/src/realtime_context_tests.rs @@ -26,6 +26,8 @@ fn thread_metadata(cwd: &str, title: &str, first_user_message: &str) -> ThreadMe agent_nickname: None, agent_role: None, model_provider: "test-provider".to_string(), + model: Some("gpt-5".to_string()), + reasoning_effort: None, cwd: PathBuf::from(cwd), cli_version: "test".to_string(), title: title.to_string(), diff --git a/codex-rs/protocol/src/openai_models.rs b/codex-rs/protocol/src/openai_models.rs index 04ca8dc9de..a113460c1a 100644 --- a/codex-rs/protocol/src/openai_models.rs +++ b/codex-rs/protocol/src/openai_models.rs @@ -4,6 +4,7 @@ //! are used to preserve compatibility when older payloads omit newly introduced attributes. use std::collections::HashMap; +use std::str::FromStr; use schemars::JsonSchema; use serde::Deserialize; @@ -48,6 +49,15 @@ pub enum ReasoningEffort { XHigh, } +impl FromStr for ReasoningEffort { + type Err = String; + + fn from_str(s: &str) -> Result { + serde_json::from_value(serde_json::Value::String(s.to_string())) + .map_err(|_| format!("invalid reasoning_effort: {s}")) + } +} + /// Canonical user-input modality tags advertised by a model. #[derive( Debug, @@ -552,6 +562,20 @@ mod tests { } } + #[test] + fn reasoning_effort_from_str_accepts_known_values() { + assert_eq!("high".parse(), Ok(ReasoningEffort::High)); + assert_eq!("minimal".parse(), Ok(ReasoningEffort::Minimal)); + } + + #[test] + fn reasoning_effort_from_str_rejects_unknown_values() { + assert_eq!( + "unsupported".parse::(), + Err("invalid reasoning_effort: unsupported".to_string()) + ); + } + #[test] fn get_model_instructions_uses_template_when_placeholder_present() { let model = test_model(Some(ModelMessages { diff --git a/codex-rs/state/migrations/0020_threads_model_reasoning_effort.sql b/codex-rs/state/migrations/0020_threads_model_reasoning_effort.sql new file mode 100644 index 0000000000..b15f4be37f --- /dev/null +++ b/codex-rs/state/migrations/0020_threads_model_reasoning_effort.sql @@ -0,0 +1,2 @@ +ALTER TABLE threads ADD COLUMN model TEXT; +ALTER TABLE threads ADD COLUMN reasoning_effort TEXT; diff --git a/codex-rs/state/src/extract.rs b/codex-rs/state/src/extract.rs index ba425adbec..037b1f5d22 100644 --- a/codex-rs/state/src/extract.rs +++ b/codex-rs/state/src/extract.rs @@ -70,6 +70,8 @@ fn apply_turn_context(metadata: &mut ThreadMetadata, turn_ctx: &TurnContextItem) if metadata.cwd.as_os_str().is_empty() { metadata.cwd = turn_ctx.cwd.clone(); } + metadata.model = Some(turn_ctx.model.clone()); + metadata.reasoning_effort = turn_ctx.effort; metadata.sandbox_policy = enum_to_string(&turn_ctx.sandbox_policy); metadata.approval_mode = enum_to_string(&turn_ctx.approval_policy); } @@ -141,6 +143,7 @@ mod tests { use codex_protocol::config_types::ReasoningSummary; use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseItem; + use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::RolloutItem; @@ -312,7 +315,7 @@ mod tests { personality: None, collaboration_mode: None, realtime_active: None, - effort: None, + effort: Some(ReasoningEffort::High), summary: ReasoningSummary::Auto, user_instructions: None, developer_instructions: None, @@ -325,6 +328,71 @@ mod tests { assert_eq!(metadata.cwd, PathBuf::from("/fallback/workspace")); } + #[test] + fn turn_context_sets_model_and_reasoning_effort() { + let mut metadata = metadata_for_test(); + + apply_rollout_item( + &mut metadata, + &RolloutItem::TurnContext(TurnContextItem { + turn_id: Some("turn-1".to_string()), + trace_id: None, + cwd: PathBuf::from("/fallback/workspace"), + current_date: None, + timezone: None, + approval_policy: AskForApproval::OnRequest, + sandbox_policy: SandboxPolicy::new_read_only_policy(), + network: None, + model: "gpt-5".to_string(), + personality: None, + collaboration_mode: None, + realtime_active: None, + effort: Some(ReasoningEffort::High), + summary: ReasoningSummary::Auto, + user_instructions: None, + developer_instructions: None, + final_output_json_schema: None, + truncation_policy: None, + }), + "test-provider", + ); + + assert_eq!(metadata.model.as_deref(), Some("gpt-5")); + assert_eq!(metadata.reasoning_effort, Some(ReasoningEffort::High)); + } + + #[test] + fn session_meta_does_not_set_model_or_reasoning_effort() { + let mut metadata = metadata_for_test(); + let thread_id = metadata.id; + + apply_rollout_item( + &mut metadata, + &RolloutItem::SessionMeta(SessionMetaLine { + meta: SessionMeta { + id: thread_id, + forked_from_id: None, + timestamp: "2026-02-26T00:00:00.000Z".to_string(), + cwd: PathBuf::from("/workspace"), + originator: "codex_cli_rs".to_string(), + cli_version: "0.0.0".to_string(), + source: SessionSource::Cli, + agent_nickname: None, + agent_role: None, + model_provider: Some("openai".to_string()), + base_instructions: None, + dynamic_tools: None, + memory_mode: None, + }, + git: None, + }), + "test-provider", + ); + + assert_eq!(metadata.model, None); + assert_eq!(metadata.reasoning_effort, None); + } + fn metadata_for_test() -> ThreadMetadata { let id = ThreadId::from_string(&Uuid::from_u128(42).to_string()).expect("thread id"); let created_at = DateTime::::from_timestamp(1_735_689_600, 0).expect("timestamp"); @@ -337,6 +405,8 @@ mod tests { agent_nickname: None, agent_role: None, model_provider: "openai".to_string(), + model: None, + reasoning_effort: None, cwd: PathBuf::from("/tmp"), cli_version: "0.0.0".to_string(), title: String::new(), diff --git a/codex-rs/state/src/model/thread_metadata.rs b/codex-rs/state/src/model/thread_metadata.rs index c4362a8df2..db4a2d95e7 100644 --- a/codex-rs/state/src/model/thread_metadata.rs +++ b/codex-rs/state/src/model/thread_metadata.rs @@ -3,6 +3,7 @@ use chrono::DateTime; use chrono::Timelike; use chrono::Utc; use codex_protocol::ThreadId; +use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::SessionSource; @@ -70,6 +71,10 @@ pub struct ThreadMetadata { pub agent_role: Option, /// The model provider identifier. pub model_provider: String, + /// The latest observed model for the thread. + pub model: Option, + /// The latest observed reasoning effort for the thread. + pub reasoning_effort: Option, /// The working directory for the thread. pub cwd: PathBuf, /// Version of the CLI that created the thread. @@ -181,6 +186,8 @@ impl ThreadMetadataBuilder { .model_provider .clone() .unwrap_or_else(|| default_provider.to_string()), + model: None, + reasoning_effort: None, cwd: self.cwd.clone(), cli_version: self.cli_version.clone().unwrap_or_default(), title: String::new(), @@ -237,6 +244,12 @@ impl ThreadMetadata { if self.model_provider != other.model_provider { diffs.push("model_provider"); } + if self.model != other.model { + diffs.push("model"); + } + if self.reasoning_effort != other.reasoning_effort { + diffs.push("reasoning_effort"); + } if self.cwd != other.cwd { diffs.push("cwd"); } @@ -288,6 +301,8 @@ pub(crate) struct ThreadRow { agent_nickname: Option, agent_role: Option, model_provider: String, + model: Option, + reasoning_effort: Option, cwd: String, cli_version: String, title: String, @@ -312,6 +327,8 @@ impl ThreadRow { agent_nickname: row.try_get("agent_nickname")?, agent_role: row.try_get("agent_role")?, model_provider: row.try_get("model_provider")?, + model: row.try_get("model")?, + reasoning_effort: row.try_get("reasoning_effort")?, cwd: row.try_get("cwd")?, cli_version: row.try_get("cli_version")?, title: row.try_get("title")?, @@ -340,6 +357,8 @@ impl TryFrom for ThreadMetadata { agent_nickname, agent_role, model_provider, + model, + reasoning_effort, cwd, cli_version, title, @@ -361,6 +380,9 @@ impl TryFrom for ThreadMetadata { agent_nickname, agent_role, model_provider, + model, + reasoning_effort: reasoning_effort + .and_then(|value| value.parse::().ok()), cwd: PathBuf::from(cwd), cli_version, title, @@ -404,3 +426,87 @@ pub struct BackfillStats { /// The number of rows that failed to upsert. pub failed: usize, } + +#[cfg(test)] +mod tests { + use super::ThreadMetadata; + use super::ThreadRow; + use chrono::DateTime; + use chrono::Utc; + use codex_protocol::ThreadId; + use codex_protocol::openai_models::ReasoningEffort; + use pretty_assertions::assert_eq; + use std::path::PathBuf; + + fn thread_row(reasoning_effort: Option<&str>) -> ThreadRow { + ThreadRow { + id: "00000000-0000-0000-0000-000000000123".to_string(), + rollout_path: "/tmp/rollout-123.jsonl".to_string(), + created_at: 1_700_000_000, + updated_at: 1_700_000_100, + source: "cli".to_string(), + agent_nickname: None, + agent_role: None, + model_provider: "openai".to_string(), + model: Some("gpt-5".to_string()), + reasoning_effort: reasoning_effort.map(str::to_string), + cwd: "/tmp/workspace".to_string(), + cli_version: "0.0.0".to_string(), + title: String::new(), + sandbox_policy: "read-only".to_string(), + approval_mode: "on-request".to_string(), + tokens_used: 1, + first_user_message: String::new(), + archived_at: None, + git_sha: None, + git_branch: None, + git_origin_url: None, + } + } + + fn expected_thread_metadata(reasoning_effort: Option) -> ThreadMetadata { + ThreadMetadata { + id: ThreadId::from_string("00000000-0000-0000-0000-000000000123") + .expect("valid thread id"), + rollout_path: PathBuf::from("/tmp/rollout-123.jsonl"), + created_at: DateTime::::from_timestamp(1_700_000_000, 0).expect("timestamp"), + updated_at: DateTime::::from_timestamp(1_700_000_100, 0).expect("timestamp"), + source: "cli".to_string(), + agent_nickname: None, + agent_role: None, + model_provider: "openai".to_string(), + model: Some("gpt-5".to_string()), + reasoning_effort, + cwd: PathBuf::from("/tmp/workspace"), + cli_version: "0.0.0".to_string(), + title: String::new(), + sandbox_policy: "read-only".to_string(), + approval_mode: "on-request".to_string(), + tokens_used: 1, + first_user_message: None, + archived_at: None, + git_sha: None, + git_branch: None, + git_origin_url: None, + } + } + + #[test] + fn thread_row_parses_reasoning_effort() { + let metadata = ThreadMetadata::try_from(thread_row(Some("high"))) + .expect("thread metadata should parse"); + + assert_eq!( + metadata, + expected_thread_metadata(Some(ReasoningEffort::High)) + ); + } + + #[test] + fn thread_row_ignores_unknown_reasoning_effort_values() { + let metadata = ThreadMetadata::try_from(thread_row(Some("future"))) + .expect("thread metadata should parse"); + + assert_eq!(metadata, expected_thread_metadata(None)); + } +} diff --git a/codex-rs/state/src/runtime/memories.rs b/codex-rs/state/src/runtime/memories.rs index 9046079140..5ca33885f3 100644 --- a/codex-rs/state/src/runtime/memories.rs +++ b/codex-rs/state/src/runtime/memories.rs @@ -169,6 +169,8 @@ SELECT agent_nickname, agent_role, model_provider, + model, + reasoning_effort, cwd, cli_version, title, diff --git a/codex-rs/state/src/runtime/test_support.rs b/codex-rs/state/src/runtime/test_support.rs index d749fe2bfb..229ece64b4 100644 --- a/codex-rs/state/src/runtime/test_support.rs +++ b/codex-rs/state/src/runtime/test_support.rs @@ -5,6 +5,8 @@ use chrono::Utc; #[cfg(test)] use codex_protocol::ThreadId; #[cfg(test)] +use codex_protocol::openai_models::ReasoningEffort; +#[cfg(test)] use codex_protocol::protocol::AskForApproval; #[cfg(test)] use codex_protocol::protocol::SandboxPolicy; @@ -49,6 +51,8 @@ pub(super) fn test_thread_metadata( agent_nickname: None, agent_role: None, model_provider: "test-provider".to_string(), + model: Some("gpt-5".to_string()), + reasoning_effort: Some(ReasoningEffort::Medium), cwd, cli_version: "0.0.0".to_string(), title: String::new(), diff --git a/codex-rs/state/src/runtime/threads.rs b/codex-rs/state/src/runtime/threads.rs index a4de8636a4..6373568e26 100644 --- a/codex-rs/state/src/runtime/threads.rs +++ b/codex-rs/state/src/runtime/threads.rs @@ -13,6 +13,8 @@ SELECT agent_nickname, agent_role, model_provider, + model, + reasoning_effort, cwd, cli_version, title, @@ -125,6 +127,8 @@ SELECT agent_nickname, agent_role, model_provider, + model, + reasoning_effort, cwd, cli_version, title, @@ -223,6 +227,8 @@ INSERT INTO threads ( agent_nickname, agent_role, model_provider, + model, + reasoning_effort, cwd, cli_version, title, @@ -236,7 +242,7 @@ INSERT INTO threads ( git_branch, git_origin_url, memory_mode -) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(id) DO NOTHING "#, ) @@ -248,6 +254,13 @@ ON CONFLICT(id) DO NOTHING .bind(metadata.agent_nickname.as_deref()) .bind(metadata.agent_role.as_deref()) .bind(metadata.model_provider.as_str()) + .bind(metadata.model.as_deref()) + .bind( + metadata + .reasoning_effort + .as_ref() + .map(crate::extract::enum_to_string), + ) .bind(metadata.cwd.display().to_string()) .bind(metadata.cli_version.as_str()) .bind(metadata.title.as_str()) @@ -337,6 +350,8 @@ INSERT INTO threads ( agent_nickname, agent_role, model_provider, + model, + reasoning_effort, cwd, cli_version, title, @@ -350,7 +365,7 @@ INSERT INTO threads ( git_branch, git_origin_url, memory_mode -) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(id) DO UPDATE SET rollout_path = excluded.rollout_path, created_at = excluded.created_at, @@ -359,6 +374,8 @@ ON CONFLICT(id) DO UPDATE SET agent_nickname = excluded.agent_nickname, agent_role = excluded.agent_role, model_provider = excluded.model_provider, + model = excluded.model, + reasoning_effort = excluded.reasoning_effort, cwd = excluded.cwd, cli_version = excluded.cli_version, title = excluded.title, @@ -381,6 +398,13 @@ ON CONFLICT(id) DO UPDATE SET .bind(metadata.agent_nickname.as_deref()) .bind(metadata.agent_role.as_deref()) .bind(metadata.model_provider.as_str()) + .bind(metadata.model.as_deref()) + .bind( + metadata + .reasoning_effort + .as_ref() + .map(crate::extract::enum_to_string), + ) .bind(metadata.cwd.display().to_string()) .bind(metadata.cli_version.as_str()) .bind(metadata.title.as_str()) From 78e8ee4591d4ff42d180000fbad29d5fb3bcd2a5 Mon Sep 17 00:00:00 2001 From: Felipe Coury Date: Tue, 17 Mar 2026 14:16:08 -0300 Subject: [PATCH 06/11] fix(tui): restore remote resume and fork history (#14930) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Problem When the TUI connects to a **remote** app-server (via WebSocket), resume and fork operations lost all conversation history. `AppServerStartedThread` carried only the `SessionConfigured` event, not the full `Thread` snapshot. After resume or fork, the chat transcript was empty — prior turns were silently discarded. A secondary issue: `primary_session_configured` was not cleared on reset, causing stale session state after reconnection. ## Approach: TUI-side only, zero app-server changes The app-server **already returns** the full `Thread` object (with populated `turns: Vec`) in its `ThreadStartResponse`, `ThreadResumeResponse`, and `ThreadForkResponse`. The data was always there — the TUI was simply throwing it away. The old `AppServerStartedThread` struct only kept the `SessionConfiguredEvent`, discarding the rich turn history that the server had already provided. This PR fixes the problem entirely within `tui_app_server` (3 files changed, 0 changes to `app-server`, `app-server-protocol`, or any other crate). Rather than modifying the server to send history in a different format or adding a new endpoint, the fix preserves the existing `Thread` snapshot and replays it through the TUI's standard event pipeline — making restored sessions indistinguishable from live ones. ## Solution Add a **thread snapshot replay** path. When the server hands back a `Thread` object (on start, resume, or fork), `restore_started_app_server_thread` converts its historical turns into the same core `Event` sequence the TUI already processes for live interactions, then replays them into the event store so the chat widget renders them. Key changes: - **`AppServerStartedThread` now carries the full `Thread`** — `started_thread_from_{start,resume,fork}_response` clone the thread into the struct alongside the existing `SessionConfiguredEvent`. - **`thread_snapshot_events()`** walks the thread's turns and items, producing `TurnStarted` → `ItemCompleted`* → `TurnComplete`/`TurnAborted` event sequences that the TUI already knows how to render. - **`restore_started_app_server_thread()`** pushes the session event + history events into the thread channel's store, activates the channel, and replays the snapshot — used for initial startup, resume, and fork. - **`primary_session_configured` cleared on reset** to prevent stale session state after reconnection. ## Tradeoffs - **`Thread` is cloned into `AppServerStartedThread`**: The full thread snapshot (including all historical turns) is cloned at startup. For long-lived threads this could be large, but it's a one-time cost and avoids lifetime gymnastics with the response. ## Tests - `restore_started_app_server_thread_replays_remote_history` — end-to-end: constructs a `Thread` with one completed turn, restores it, and asserts user/agent messages appear in the transcript. - `bridges_thread_snapshot_turns_for_resume_restore` — unit: verifies `thread_snapshot_events` produces the correct event sequence for completed and interrupted turns. ## Test plan - [ ] Verify `cargo check -p codex-tui-app-server` passes - [ ] Verify `cargo test -p codex-tui-app-server` passes - [ ] Manual: connect to a remote app-server, resume an existing thread, confirm history renders in the chat widget - [ ] Manual: fork a thread via remote, confirm prior turns appear --- codex-rs/tui_app_server/src/app.rs | 486 +++++++++++++++++- .../src/app/app_server_adapter.rs | 461 +++++++++++++++-- .../tui_app_server/src/app_server_session.rs | 214 ++------ codex-rs/tui_app_server/src/chatwidget.rs | 27 +- .../tui_app_server/src/chatwidget/tests.rs | 1 + 5 files changed, 958 insertions(+), 231 deletions(-) diff --git a/codex-rs/tui_app_server/src/app.rs b/codex-rs/tui_app_server/src/app.rs index 9fd3f1bd3d..c08e74c337 100644 --- a/codex-rs/tui_app_server/src/app.rs +++ b/codex-rs/tui_app_server/src/app.rs @@ -118,6 +118,7 @@ mod pending_interactive_replay; use self::agent_navigation::AgentNavigationDirection; use self::agent_navigation::AgentNavigationState; +use self::app_server_adapter::thread_snapshot_events; use self::app_server_requests::PendingAppServerRequests; use self::pending_interactive_replay::PendingInteractiveReplayState; @@ -2050,6 +2051,7 @@ impl App { self.active_thread_id = None; self.active_thread_rx = None; self.primary_thread_id = None; + self.primary_session_configured = None; self.pending_primary_events.clear(); self.pending_app_server_requests.clear(); self.chat_widget.set_pending_thread_approvals(Vec::new()); @@ -2117,11 +2119,66 @@ impl App { let init = self.chatwidget_init_for_forked_or_resumed_thread(tui, self.config.clone()); self.chat_widget = ChatWidget::new_with_app_event(init); self.reset_thread_event_state(); - self.enqueue_primary_event(Event { + self.restore_started_app_server_thread(started).await + } + + /// Hydrate thread state from an `AppServerStartedThread` returned by the + /// app-server start/resume/fork handshake. + /// + /// This is the single path that every session-start variant funnels + /// through. It performs four things in order: + /// + /// 1. Converts the `Thread` snapshot into protocol-level `Event`s. + /// 2. Builds a **lossless** replay snapshot from a temporary store so that + /// the initial render sees all history even when the thread has more + /// turns than the bounded channel capacity. + /// 3. Pushes the same events into the real channel store for backtrack and + /// navigation. + /// 4. Activates the thread channel and replays the snapshot into the chat + /// widget. + async fn restore_started_app_server_thread( + &mut self, + started: AppServerStartedThread, + ) -> Result<()> { + let session_configured = started.session_configured; + let thread_id = session_configured.session_id; + let session_event = Event { id: String::new(), - msg: EventMsg::SessionConfigured(started.session_configured), - }) - .await + msg: EventMsg::SessionConfigured(session_configured.clone()), + }; + let history_events = + thread_snapshot_events(&started.thread, started.show_raw_agent_reasoning); + let replay_snapshot = { + let mut replay_store = ThreadEventStore::new(history_events.len().saturating_add(1)); + replay_store.push_event(session_event.clone()); + for event in &history_events { + replay_store.push_event(event.clone()); + } + replay_store.snapshot() + }; + + self.primary_thread_id = Some(thread_id); + self.primary_session_configured = Some(session_configured); + self.upsert_agent_picker_thread( + thread_id, /*agent_nickname*/ None, /*agent_role*/ None, + /*is_closed*/ false, + ); + + let store = { + let channel = self.ensure_thread_channel(thread_id); + Arc::clone(&channel.store) + }; + { + let mut store = store.lock().await; + store.push_event(session_event); + for event in history_events { + store.push_event(event); + } + } + + self.activate_thread_channel(thread_id).await; + self.replay_thread_snapshot(replay_snapshot, /*resume_restored_queue*/ false); + Ok(()) } fn fresh_session_config(&self) -> Config { @@ -2187,6 +2244,8 @@ impl App { snapshot: ThreadEventSnapshot, resume_restored_queue: bool, ) { + self.chat_widget + .set_initial_user_message_submit_suppressed(/*suppressed*/ true); if let Some(event) = snapshot.session_configured { self.handle_codex_event_replay(event); } @@ -2199,6 +2258,9 @@ impl App { } self.chat_widget .set_queue_autosend_suppressed(/*suppressed*/ false); + self.chat_widget + .set_initial_user_message_submit_suppressed(/*suppressed*/ false); + self.chat_widget.submit_initial_user_message_if_pending(); if resume_restored_queue { self.chat_widget.maybe_send_next_queued_input(); } @@ -2313,7 +2375,7 @@ impl App { let enhanced_keys_supported = tui.enhanced_keys_supported(); let wait_for_initial_session_configured = Self::should_wait_for_initial_session(&session_selection); - let (mut chat_widget, initial_session_configured) = match session_selection { + let (mut chat_widget, initial_started_thread) = match session_selection { SessionSelection::StartFresh | SessionSelection::Exit => { let started = app_server.start_thread(&config).await?; let startup_tooltip_override = @@ -2342,10 +2404,7 @@ impl App { status_line_invalid_items_warned: status_line_invalid_items_warned.clone(), session_telemetry: session_telemetry.clone(), }; - ( - ChatWidget::new_with_app_event(init), - Some(started.session_configured), - ) + (ChatWidget::new_with_app_event(init), started) } SessionSelection::Resume(target_session) => { let resumed = app_server @@ -2378,10 +2437,7 @@ impl App { status_line_invalid_items_warned: status_line_invalid_items_warned.clone(), session_telemetry: session_telemetry.clone(), }; - ( - ChatWidget::new_with_app_event(init), - Some(resumed.session_configured), - ) + (ChatWidget::new_with_app_event(init), resumed) } SessionSelection::Fork(target_session) => { session_telemetry.counter( @@ -2419,10 +2475,7 @@ impl App { status_line_invalid_items_warned: status_line_invalid_items_warned.clone(), session_telemetry: session_telemetry.clone(), }; - ( - ChatWidget::new_with_app_event(init), - Some(forked.session_configured), - ) + (ChatWidget::new_with_app_event(init), forked) } }; @@ -2474,13 +2527,8 @@ impl App { pending_primary_events: VecDeque::new(), pending_app_server_requests: PendingAppServerRequests::default(), }; - if let Some(session_configured) = initial_session_configured { - app.enqueue_primary_event(Event { - id: String::new(), - msg: EventMsg::SessionConfigured(session_configured), - }) + app.restore_started_app_server_thread(initial_started_thread) .await?; - } // On startup, if Agent mode (workspace-write) or ReadOnly is active, warn about world-writable dirs on Windows. #[cfg(target_os = "windows")] @@ -4427,6 +4475,7 @@ mod tests { use crate::app_backtrack::BacktrackSelection; use crate::app_backtrack::BacktrackState; use crate::app_backtrack::user_count; + use crate::app_server_session::AppServerStartedThread; use crate::chatwidget::tests::make_chatwidget_manual_with_sender; use crate::chatwidget::tests::set_chatgpt_auth; use crate::file_search::FileSearchManager; @@ -4437,6 +4486,11 @@ mod tests { use crate::multi_agents::AgentPickerThreadEntry; use assert_matches::assert_matches; + use codex_app_server_protocol::Thread; + use codex_app_server_protocol::ThreadItem; + use codex_app_server_protocol::ThreadStatus; + use codex_app_server_protocol::Turn; + use codex_app_server_protocol::TurnStatus; use codex_core::config::ConfigBuilder; use codex_core::config::ConfigOverrides; use codex_core::config::types::ModelAvailabilityNuxConfig; @@ -6680,6 +6734,392 @@ guardian_approval = true ) } + #[tokio::test] + async fn restore_started_app_server_thread_replays_remote_history() -> Result<()> { + let (mut app, mut app_event_rx, _op_rx) = make_test_app_with_channels().await; + let thread_id = ThreadId::new(); + + app.restore_started_app_server_thread(AppServerStartedThread { + thread: Thread { + id: thread_id.to_string(), + preview: "hello".to_string(), + ephemeral: false, + model_provider: "test-provider".to_string(), + created_at: 0, + updated_at: 0, + status: ThreadStatus::Idle, + path: None, + cwd: PathBuf::from("/tmp/project"), + cli_version: "test".to_string(), + source: SessionSource::Cli.into(), + agent_nickname: None, + agent_role: None, + git_info: None, + name: Some("restored".to_string()), + turns: vec![Turn { + id: "turn-1".to_string(), + items: vec![ + ThreadItem::UserMessage { + id: "user-1".to_string(), + content: vec![codex_app_server_protocol::UserInput::Text { + text: "hello from remote".to_string(), + text_elements: Vec::new(), + }], + }, + ThreadItem::AgentMessage { + id: "assistant-1".to_string(), + text: "restored response".to_string(), + phase: None, + }, + ], + status: TurnStatus::Completed, + error: None, + }], + }, + session_configured: SessionConfiguredEvent { + session_id: thread_id, + forked_from_id: None, + thread_name: Some("restored".to_string()), + model: "gpt-test".to_string(), + model_provider_id: "test-provider".to_string(), + service_tier: None, + approval_policy: AskForApproval::Never, + approvals_reviewer: ApprovalsReviewer::User, + sandbox_policy: SandboxPolicy::new_read_only_policy(), + cwd: PathBuf::from("/tmp/project"), + reasoning_effort: None, + history_log_id: 0, + history_entry_count: 0, + initial_messages: None, + network_proxy: None, + rollout_path: Some(PathBuf::new()), + }, + show_raw_agent_reasoning: false, + }) + .await?; + + while let Ok(event) = app_event_rx.try_recv() { + if let AppEvent::InsertHistoryCell(cell) = event { + let cell: Arc = cell.into(); + app.transcript_cells.push(cell); + } + } + + assert_eq!(app.primary_thread_id, Some(thread_id)); + assert_eq!(app.active_thread_id, Some(thread_id)); + + let user_messages: Vec = app + .transcript_cells + .iter() + .filter_map(|cell| { + cell.as_any() + .downcast_ref::() + .map(|cell| cell.message.clone()) + }) + .collect(); + let agent_messages: Vec = app + .transcript_cells + .iter() + .filter_map(|cell| { + cell.as_any() + .downcast_ref::() + .map(|cell| { + cell.display_lines(80) + .into_iter() + .map(|line| line.to_string()) + .collect::>() + .join("\n") + }) + }) + .collect(); + + assert_eq!(user_messages, vec!["hello from remote".to_string()]); + assert_eq!(agent_messages, vec!["• restored response".to_string()]); + + Ok(()) + } + + #[tokio::test] + async fn restore_started_app_server_thread_submits_initial_prompt_after_history_replay() + -> Result<()> { + let (mut app, mut app_event_rx, mut op_rx) = make_test_app_with_channels().await; + let thread_id = ThreadId::new(); + app.chat_widget.set_initial_user_message_for_test( + crate::chatwidget::create_initial_user_message( + Some("resume prompt".to_string()), + Vec::new(), + Vec::new(), + ), + ); + + app.restore_started_app_server_thread(AppServerStartedThread { + thread: Thread { + id: thread_id.to_string(), + preview: "hello".to_string(), + ephemeral: false, + model_provider: "test-provider".to_string(), + created_at: 0, + updated_at: 0, + status: ThreadStatus::Idle, + path: None, + cwd: PathBuf::from("/tmp/project"), + cli_version: "test".to_string(), + source: SessionSource::Cli.into(), + agent_nickname: None, + agent_role: None, + git_info: None, + name: Some("restored".to_string()), + turns: vec![Turn { + id: "turn-1".to_string(), + items: vec![ + ThreadItem::UserMessage { + id: "user-1".to_string(), + content: vec![codex_app_server_protocol::UserInput::Text { + text: "hello from remote".to_string(), + text_elements: Vec::new(), + }], + }, + ThreadItem::AgentMessage { + id: "assistant-1".to_string(), + text: "restored response".to_string(), + phase: None, + }, + ], + status: TurnStatus::Completed, + error: None, + }], + }, + session_configured: SessionConfiguredEvent { + session_id: thread_id, + forked_from_id: None, + thread_name: Some("restored".to_string()), + model: "gpt-test".to_string(), + model_provider_id: "test-provider".to_string(), + service_tier: None, + approval_policy: AskForApproval::Never, + approvals_reviewer: ApprovalsReviewer::User, + sandbox_policy: SandboxPolicy::new_read_only_policy(), + cwd: PathBuf::from("/tmp/project"), + reasoning_effort: None, + history_log_id: 0, + history_entry_count: 0, + initial_messages: None, + network_proxy: None, + rollout_path: Some(PathBuf::new()), + }, + show_raw_agent_reasoning: false, + }) + .await?; + + while let Ok(event) = app_event_rx.try_recv() { + if let AppEvent::InsertHistoryCell(cell) = event { + let cell: Arc = cell.into(); + app.transcript_cells.push(cell); + } + } + + let user_messages: Vec = app + .transcript_cells + .iter() + .filter_map(|cell| { + cell.as_any() + .downcast_ref::() + .map(|cell| cell.message.clone()) + }) + .collect(); + + assert_eq!( + user_messages, + vec!["hello from remote".to_string(), "resume prompt".to_string()] + ); + match next_user_turn_op(&mut op_rx) { + Op::UserTurn { items, .. } => assert_eq!( + items, + vec![UserInput::Text { + text: "resume prompt".to_string(), + text_elements: Vec::new(), + }] + ), + other => panic!("expected resume prompt submission, got {other:?}"), + } + + Ok(()) + } + + #[tokio::test] + async fn restore_started_app_server_thread_replays_history_beyond_store_capacity() -> Result<()> + { + let (mut app, mut app_event_rx, _op_rx) = make_test_app_with_channels().await; + let thread_id = ThreadId::new(); + let turn_count = THREAD_EVENT_CHANNEL_CAPACITY + 5; + + let turns = (0..turn_count) + .map(|index| Turn { + id: format!("turn-{index}"), + items: vec![ThreadItem::UserMessage { + id: format!("user-{index}"), + content: vec![codex_app_server_protocol::UserInput::Text { + text: format!("message {index}"), + text_elements: Vec::new(), + }], + }], + status: TurnStatus::Completed, + error: None, + }) + .collect(); + + app.restore_started_app_server_thread(AppServerStartedThread { + thread: Thread { + id: thread_id.to_string(), + preview: "hello".to_string(), + ephemeral: false, + model_provider: "test-provider".to_string(), + created_at: 0, + updated_at: 0, + status: ThreadStatus::Idle, + path: None, + cwd: PathBuf::from("/tmp/project"), + cli_version: "test".to_string(), + source: SessionSource::Cli.into(), + agent_nickname: None, + agent_role: None, + git_info: None, + name: Some("restored".to_string()), + turns, + }, + session_configured: SessionConfiguredEvent { + session_id: thread_id, + forked_from_id: None, + thread_name: Some("restored".to_string()), + model: "gpt-test".to_string(), + model_provider_id: "test-provider".to_string(), + service_tier: None, + approval_policy: AskForApproval::Never, + approvals_reviewer: ApprovalsReviewer::User, + sandbox_policy: SandboxPolicy::new_read_only_policy(), + cwd: PathBuf::from("/tmp/project"), + reasoning_effort: None, + history_log_id: 0, + history_entry_count: 0, + initial_messages: None, + network_proxy: None, + rollout_path: Some(PathBuf::new()), + }, + show_raw_agent_reasoning: false, + }) + .await?; + + while let Ok(event) = app_event_rx.try_recv() { + if let AppEvent::InsertHistoryCell(cell) = event { + let cell: Arc = cell.into(); + app.transcript_cells.push(cell); + } + } + + let user_messages: Vec = app + .transcript_cells + .iter() + .filter_map(|cell| { + cell.as_any() + .downcast_ref::() + .map(|cell| cell.message.clone()) + }) + .collect(); + + assert_eq!(user_messages.len(), turn_count); + assert_eq!(user_messages.first().map(String::as_str), Some("message 0")); + let last_message = format!("message {}", turn_count - 1); + assert_eq!( + user_messages.last().map(String::as_str), + Some(last_message.as_str()) + ); + + Ok(()) + } + + #[tokio::test] + async fn restore_started_app_server_thread_replays_raw_reasoning_when_enabled() -> Result<()> { + let (mut app, mut app_event_rx, _op_rx) = make_test_app_with_channels().await; + let thread_id = ThreadId::new(); + + app.restore_started_app_server_thread(AppServerStartedThread { + thread: Thread { + id: thread_id.to_string(), + preview: "hello".to_string(), + ephemeral: false, + model_provider: "test-provider".to_string(), + created_at: 0, + updated_at: 0, + status: ThreadStatus::Idle, + path: None, + cwd: PathBuf::from("/tmp/project"), + cli_version: "test".to_string(), + source: SessionSource::Cli.into(), + agent_nickname: None, + agent_role: None, + git_info: None, + name: Some("restored".to_string()), + turns: vec![Turn { + id: "turn-1".to_string(), + items: vec![ThreadItem::Reasoning { + id: "reasoning-1".to_string(), + summary: vec!["summary reasoning".to_string()], + content: vec!["raw reasoning".to_string()], + }], + status: TurnStatus::Completed, + error: None, + }], + }, + session_configured: SessionConfiguredEvent { + session_id: thread_id, + forked_from_id: None, + thread_name: Some("restored".to_string()), + model: "gpt-test".to_string(), + model_provider_id: "test-provider".to_string(), + service_tier: None, + approval_policy: AskForApproval::Never, + approvals_reviewer: ApprovalsReviewer::User, + sandbox_policy: SandboxPolicy::new_read_only_policy(), + cwd: PathBuf::from("/tmp/project"), + reasoning_effort: None, + history_log_id: 0, + history_entry_count: 0, + initial_messages: None, + network_proxy: None, + rollout_path: Some(PathBuf::new()), + }, + show_raw_agent_reasoning: true, + }) + .await?; + + while let Ok(event) = app_event_rx.try_recv() { + if let AppEvent::InsertHistoryCell(cell) = event { + let cell: Arc = cell.into(); + app.transcript_cells.push(cell); + } + } + + let channel = app + .thread_event_channels + .get(&thread_id) + .expect("restored thread channel should exist"); + let snapshot = channel.store.lock().await.snapshot(); + let replayed_raw_reasoning = snapshot.events.iter().any(|event| { + matches!( + &event.msg, + EventMsg::AgentReasoningRawContent(raw) if raw.text == "raw reasoning" + ) + }); + + assert!( + replayed_raw_reasoning, + "expected restored snapshot to keep raw reasoning event: {:?}", + snapshot.events + ); + + Ok(()) + } + #[test] fn thread_event_store_tracks_active_turn_lifecycle() { let mut store = ThreadEventStore::new(8); diff --git a/codex-rs/tui_app_server/src/app/app_server_adapter.rs b/codex-rs/tui_app_server/src/app/app_server_adapter.rs index d21542b8bf..0fff49fd20 100644 --- a/codex-rs/tui_app_server/src/app/app_server_adapter.rs +++ b/codex-rs/tui_app_server/src/app/app_server_adapter.rs @@ -19,7 +19,10 @@ use crate::app_server_session::status_account_display_from_auth_mode; use codex_app_server_client::AppServerEvent; use codex_app_server_protocol::JSONRPCErrorError; use codex_app_server_protocol::ServerNotification; +use codex_app_server_protocol::Thread; use codex_app_server_protocol::ThreadItem; +use codex_app_server_protocol::Turn; +use codex_app_server_protocol::TurnStatus; use codex_protocol::ThreadId; use codex_protocol::config_types::ModeKind; use codex_protocol::items::AgentMessageContent; @@ -48,6 +51,8 @@ use codex_protocol::protocol::ThreadNameUpdatedEvent; use codex_protocol::protocol::TokenCountEvent; use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TokenUsageInfo; +use codex_protocol::protocol::TurnAbortReason; +use codex_protocol::protocol::TurnAbortedEvent; use codex_protocol::protocol::TurnCompleteEvent; use codex_protocol::protocol::TurnStartedEvent; use serde_json::Value; @@ -196,6 +201,31 @@ impl App { } } +/// Convert a `Thread` snapshot into a flat sequence of protocol `Event`s +/// suitable for replaying into the TUI event store. +/// +/// Each turn is expanded into `TurnStarted`, zero or more `ItemCompleted`, +/// and a terminal event that matches the turn's `TurnStatus`. Returns an +/// empty vec (with a warning log) if the thread ID is not a valid UUID. +pub(super) fn thread_snapshot_events( + thread: &Thread, + show_raw_agent_reasoning: bool, +) -> Vec { + let Ok(thread_id) = ThreadId::from_string(&thread.id) else { + tracing::warn!( + thread_id = %thread.id, + "ignoring app-server thread snapshot with invalid thread id" + ); + return Vec::new(); + }; + + thread + .turns + .iter() + .flat_map(|turn| turn_snapshot_events(thread_id, turn, show_raw_agent_reasoning)) + .collect() +} + fn legacy_thread_event(params: Option) -> Option<(ThreadId, Event)> { let Value::Object(mut params) = params? else { return None; @@ -286,16 +316,16 @@ fn server_notification_thread_events( }), }], )), - ServerNotification::TurnCompleted(notification) => Some(( - ThreadId::from_string(¬ification.thread_id).ok()?, - vec![Event { - id: String::new(), - msg: EventMsg::TurnComplete(TurnCompleteEvent { - turn_id: notification.turn.id, - last_agent_message: None, - }), - }], - )), + ServerNotification::TurnCompleted(notification) => { + let thread_id = ThreadId::from_string(¬ification.thread_id).ok()?; + let mut events = Vec::new(); + append_terminal_turn_events( + &mut events, + ¬ification.turn, + /*include_failed_error*/ false, + ); + Some((thread_id, events)) + } ServerNotification::ItemStarted(notification) => Some(( ThreadId::from_string(¬ification.thread_id).ok()?, vec![Event { @@ -303,7 +333,7 @@ fn server_notification_thread_events( msg: EventMsg::ItemStarted(ItemStartedEvent { thread_id: ThreadId::from_string(¬ification.thread_id).ok()?, turn_id: notification.turn_id, - item: thread_item_to_core(notification.item)?, + item: thread_item_to_core(¬ification.item)?, }), }], )), @@ -314,7 +344,7 @@ fn server_notification_thread_events( msg: EventMsg::ItemCompleted(ItemCompletedEvent { thread_id: ThreadId::from_string(¬ification.thread_id).ok()?, turn_id: notification.turn_id, - item: thread_item_to_core(notification.item)?, + item: thread_item_to_core(¬ification.item)?, }), }], )), @@ -418,36 +448,150 @@ fn token_usage_from_app_server( } } -fn thread_item_to_core(item: ThreadItem) -> Option { +/// Expand a single `Turn` into the event sequence the TUI would have +/// observed if it had been connected for the turn's entire lifetime. +/// +/// Snapshot replay keeps committed-item semantics for user / plan / +/// agent-message items, while replaying the legacy events that still +/// drive rendering for reasoning, web-search, image-generation, and +/// context-compaction history cells. +fn turn_snapshot_events( + thread_id: ThreadId, + turn: &Turn, + show_raw_agent_reasoning: bool, +) -> Vec { + let mut events = vec![Event { + id: String::new(), + msg: EventMsg::TurnStarted(TurnStartedEvent { + turn_id: turn.id.clone(), + model_context_window: None, + collaboration_mode_kind: ModeKind::default(), + }), + }]; + + for item in &turn.items { + let Some(item) = thread_item_to_core(item) else { + continue; + }; + match item { + TurnItem::UserMessage(_) | TurnItem::Plan(_) | TurnItem::AgentMessage(_) => { + events.push(Event { + id: String::new(), + msg: EventMsg::ItemCompleted(ItemCompletedEvent { + thread_id, + turn_id: turn.id.clone(), + item, + }), + }); + } + TurnItem::Reasoning(_) + | TurnItem::WebSearch(_) + | TurnItem::ImageGeneration(_) + | TurnItem::ContextCompaction(_) => { + events.extend( + item.as_legacy_events(show_raw_agent_reasoning) + .into_iter() + .map(|msg| Event { + id: String::new(), + msg, + }), + ); + } + } + } + + append_terminal_turn_events(&mut events, turn, /*include_failed_error*/ true); + + events +} + +/// Append the terminal event(s) for a turn based on its `TurnStatus`. +/// +/// This function is shared between the live notification bridge +/// (`TurnCompleted` handling) and the snapshot replay path so that both +/// produce identical `EventMsg` sequences for the same turn status. +/// +/// - `Completed` → `TurnComplete` +/// - `Interrupted` → `TurnAborted { reason: Interrupted }` +/// - `Failed` → `Error` (if present) then `TurnComplete` +/// - `InProgress` → no events (the turn is still running) +fn append_terminal_turn_events(events: &mut Vec, turn: &Turn, include_failed_error: bool) { + match turn.status { + TurnStatus::Completed => events.push(Event { + id: String::new(), + msg: EventMsg::TurnComplete(TurnCompleteEvent { + turn_id: turn.id.clone(), + last_agent_message: None, + }), + }), + TurnStatus::Interrupted => events.push(Event { + id: String::new(), + msg: EventMsg::TurnAborted(TurnAbortedEvent { + turn_id: Some(turn.id.clone()), + reason: TurnAbortReason::Interrupted, + }), + }), + TurnStatus::Failed => { + if include_failed_error && let Some(error) = &turn.error { + events.push(Event { + id: String::new(), + msg: EventMsg::Error(ErrorEvent { + message: error.message.clone(), + codex_error_info: error + .codex_error_info + .clone() + .and_then(app_server_codex_error_info_to_core), + }), + }); + } + events.push(Event { + id: String::new(), + msg: EventMsg::TurnComplete(TurnCompleteEvent { + turn_id: turn.id.clone(), + last_agent_message: None, + }), + }); + } + TurnStatus::InProgress => { + // Preserve unfinished turns during snapshot replay without emitting completion events. + } + } +} + +fn thread_item_to_core(item: &ThreadItem) -> Option { match item { ThreadItem::UserMessage { id, content } => Some(TurnItem::UserMessage(UserMessageItem { - id, + id: id.clone(), content: content - .into_iter() + .iter() + .cloned() .map(codex_app_server_protocol::UserInput::into_core) .collect(), })), ThreadItem::AgentMessage { id, text, phase } => { Some(TurnItem::AgentMessage(AgentMessageItem { - id, - content: vec![AgentMessageContent::Text { text }], - phase, + id: id.clone(), + content: vec![AgentMessageContent::Text { text: text.clone() }], + phase: phase.clone(), })) } - ThreadItem::Plan { id, text } => Some(TurnItem::Plan(PlanItem { id, text })), + ThreadItem::Plan { id, text } => Some(TurnItem::Plan(PlanItem { + id: id.clone(), + text: text.clone(), + })), ThreadItem::Reasoning { id, summary, content, } => Some(TurnItem::Reasoning(ReasoningItem { - id, - summary_text: summary, - raw_content: content, + id: id.clone(), + summary_text: summary.clone(), + raw_content: content.clone(), })), ThreadItem::WebSearch { id, query, action } => Some(TurnItem::WebSearch(WebSearchItem { - id, - query, - action: app_server_web_search_action_to_core(action?)?, + id: id.clone(), + query: query.clone(), + action: app_server_web_search_action_to_core(action.clone()?)?, })), ThreadItem::ImageGeneration { id, @@ -455,14 +599,16 @@ fn thread_item_to_core(item: ThreadItem) -> Option { revised_prompt, result, } => Some(TurnItem::ImageGeneration(ImageGenerationItem { - id, - status, - revised_prompt, - result, + id: id.clone(), + status: status.clone(), + revised_prompt: revised_prompt.clone(), + result: result.clone(), saved_path: None, })), ThreadItem::ContextCompaction { id } => { - Some(TurnItem::ContextCompaction(ContextCompactionItem { id })) + Some(TurnItem::ContextCompaction(ContextCompactionItem { + id: id.clone(), + })) } ThreadItem::CommandExecution { .. } | ThreadItem::FileChange { .. } @@ -491,7 +637,9 @@ fn app_server_web_search_action_to_core( codex_app_server_protocol::WebSearchAction::FindInPage { url, pattern } => { Some(codex_protocol::models::WebSearchAction::FindInPage { url, pattern }) } - codex_app_server_protocol::WebSearchAction::Other => None, + codex_app_server_protocol::WebSearchAction::Other => { + Some(codex_protocol::models::WebSearchAction::Other) + } } } @@ -504,13 +652,19 @@ fn app_server_codex_error_info_to_core( #[cfg(test)] mod tests { use super::server_notification_thread_events; + use super::thread_snapshot_events; + use super::turn_snapshot_events; use codex_app_server_protocol::AgentMessageDeltaNotification; + use codex_app_server_protocol::CodexErrorInfo; use codex_app_server_protocol::ItemCompletedNotification; use codex_app_server_protocol::ReasoningSummaryTextDeltaNotification; use codex_app_server_protocol::ServerNotification; + use codex_app_server_protocol::Thread; use codex_app_server_protocol::ThreadItem; + use codex_app_server_protocol::ThreadStatus; use codex_app_server_protocol::Turn; use codex_app_server_protocol::TurnCompletedNotification; + use codex_app_server_protocol::TurnError; use codex_app_server_protocol::TurnStatus; use codex_protocol::ThreadId; use codex_protocol::items::AgentMessageContent; @@ -518,7 +672,11 @@ mod tests { use codex_protocol::items::TurnItem; use codex_protocol::models::MessagePhase; use codex_protocol::protocol::EventMsg; + use codex_protocol::protocol::SessionSource; + use codex_protocol::protocol::TurnAbortReason; + use codex_protocol::protocol::TurnAbortedEvent; use pretty_assertions::assert_eq; + use std::path::PathBuf; #[test] fn bridges_completed_agent_messages_from_server_notifications() { @@ -601,6 +759,74 @@ mod tests { assert_eq!(completed.last_agent_message, None); } + #[test] + fn bridges_interrupted_turn_completion_from_server_notifications() { + let thread_id = "019cee8c-b993-7e33-88c0-014d4e62612d".to_string(); + let turn_id = "019cee8c-b9b4-7f10-a1b0-38caa876a012".to_string(); + + let (actual_thread_id, events) = server_notification_thread_events( + ServerNotification::TurnCompleted(TurnCompletedNotification { + thread_id: thread_id.clone(), + turn: Turn { + id: turn_id.clone(), + items: Vec::new(), + status: TurnStatus::Interrupted, + error: None, + }, + }), + ) + .expect("notification should bridge"); + + assert_eq!( + actual_thread_id, + ThreadId::from_string(&thread_id).expect("valid thread id") + ); + let [event] = events.as_slice() else { + panic!("expected one bridged event"); + }; + let EventMsg::TurnAborted(aborted) = &event.msg else { + panic!("expected turn aborted event"); + }; + assert_eq!(aborted.turn_id.as_deref(), Some(turn_id.as_str())); + assert_eq!(aborted.reason, TurnAbortReason::Interrupted); + } + + #[test] + fn bridges_failed_turn_completion_from_server_notifications() { + let thread_id = "019cee8c-b993-7e33-88c0-014d4e62612d".to_string(); + let turn_id = "019cee8c-b9b4-7f10-a1b0-38caa876a012".to_string(); + + let (actual_thread_id, events) = server_notification_thread_events( + ServerNotification::TurnCompleted(TurnCompletedNotification { + thread_id: thread_id.clone(), + turn: Turn { + id: turn_id.clone(), + items: Vec::new(), + status: TurnStatus::Failed, + error: Some(TurnError { + message: "request failed".to_string(), + codex_error_info: Some(CodexErrorInfo::Other), + additional_details: None, + }), + }, + }), + ) + .expect("notification should bridge"); + + assert_eq!( + actual_thread_id, + ThreadId::from_string(&thread_id).expect("valid thread id") + ); + let [complete_event] = events.as_slice() else { + panic!("expected turn completion only"); + }; + let EventMsg::TurnComplete(completed) = &complete_event.msg else { + panic!("expected turn complete event"); + }; + assert_eq!(completed.turn_id, turn_id); + assert_eq!(completed.last_agent_message, None); + } + #[test] fn bridges_text_deltas_from_server_notifications() { let thread_id = "019cee8c-b993-7e33-88c0-014d4e62612d".to_string(); @@ -642,4 +868,177 @@ mod tests { }; assert_eq!(delta.delta, "Thinking"); } + + #[test] + fn bridges_thread_snapshot_turns_for_resume_restore() { + let thread_id = ThreadId::new(); + let events = thread_snapshot_events( + &Thread { + id: thread_id.to_string(), + preview: "hello".to_string(), + ephemeral: false, + model_provider: "openai".to_string(), + created_at: 0, + updated_at: 0, + status: ThreadStatus::Idle, + path: None, + cwd: PathBuf::from("/tmp/project"), + cli_version: "test".to_string(), + source: SessionSource::Cli.into(), + agent_nickname: None, + agent_role: None, + git_info: None, + name: Some("restore".to_string()), + turns: vec![ + Turn { + id: "turn-complete".to_string(), + items: vec![ + ThreadItem::UserMessage { + id: "user-1".to_string(), + content: vec![codex_app_server_protocol::UserInput::Text { + text: "hello".to_string(), + text_elements: Vec::new(), + }], + }, + ThreadItem::AgentMessage { + id: "assistant-1".to_string(), + text: "hi".to_string(), + phase: Some(MessagePhase::FinalAnswer), + }, + ], + status: TurnStatus::Completed, + error: None, + }, + Turn { + id: "turn-interrupted".to_string(), + items: Vec::new(), + status: TurnStatus::Interrupted, + error: None, + }, + Turn { + id: "turn-failed".to_string(), + items: Vec::new(), + status: TurnStatus::Failed, + error: Some(TurnError { + message: "request failed".to_string(), + codex_error_info: Some(CodexErrorInfo::Other), + additional_details: None, + }), + }, + ], + }, + /*show_raw_agent_reasoning*/ false, + ); + + assert_eq!(events.len(), 9); + assert!(matches!(events[0].msg, EventMsg::TurnStarted(_))); + assert!(matches!(events[1].msg, EventMsg::ItemCompleted(_))); + assert!(matches!(events[2].msg, EventMsg::ItemCompleted(_))); + assert!(matches!(events[3].msg, EventMsg::TurnComplete(_))); + assert!(matches!(events[4].msg, EventMsg::TurnStarted(_))); + let EventMsg::TurnAborted(TurnAbortedEvent { turn_id, reason }) = &events[5].msg else { + panic!("expected interrupted turn replay"); + }; + assert_eq!(turn_id.as_deref(), Some("turn-interrupted")); + assert_eq!(*reason, TurnAbortReason::Interrupted); + assert!(matches!(events[6].msg, EventMsg::TurnStarted(_))); + let EventMsg::Error(error) = &events[7].msg else { + panic!("expected failed turn error replay"); + }; + assert_eq!(error.message, "request failed"); + assert_eq!( + error.codex_error_info, + Some(codex_protocol::protocol::CodexErrorInfo::Other) + ); + assert!(matches!(events[8].msg, EventMsg::TurnComplete(_))); + } + + #[test] + fn bridges_non_message_snapshot_items_via_legacy_events() { + let events = turn_snapshot_events( + ThreadId::new(), + &Turn { + id: "turn-complete".to_string(), + items: vec![ + ThreadItem::Reasoning { + id: "reasoning-1".to_string(), + summary: vec!["Need to inspect config".to_string()], + content: vec!["hidden chain".to_string()], + }, + ThreadItem::WebSearch { + id: "search-1".to_string(), + query: "ratatui stylize".to_string(), + action: Some(codex_app_server_protocol::WebSearchAction::Other), + }, + ThreadItem::ImageGeneration { + id: "image-1".to_string(), + status: "completed".to_string(), + revised_prompt: Some("diagram".to_string()), + result: "image.png".to_string(), + }, + ThreadItem::ContextCompaction { + id: "compact-1".to_string(), + }, + ], + status: TurnStatus::Completed, + error: None, + }, + /*show_raw_agent_reasoning*/ false, + ); + + assert_eq!(events.len(), 6); + assert!(matches!(events[0].msg, EventMsg::TurnStarted(_))); + let EventMsg::AgentReasoning(reasoning) = &events[1].msg else { + panic!("expected reasoning replay"); + }; + assert_eq!(reasoning.text, "Need to inspect config"); + let EventMsg::WebSearchEnd(web_search) = &events[2].msg else { + panic!("expected web search replay"); + }; + assert_eq!(web_search.call_id, "search-1"); + assert_eq!(web_search.query, "ratatui stylize"); + assert_eq!( + web_search.action, + codex_protocol::models::WebSearchAction::Other + ); + let EventMsg::ImageGenerationEnd(image_generation) = &events[3].msg else { + panic!("expected image generation replay"); + }; + assert_eq!(image_generation.call_id, "image-1"); + assert_eq!(image_generation.status, "completed"); + assert_eq!(image_generation.revised_prompt.as_deref(), Some("diagram")); + assert_eq!(image_generation.result, "image.png"); + assert!(matches!(events[4].msg, EventMsg::ContextCompacted(_))); + assert!(matches!(events[5].msg, EventMsg::TurnComplete(_))); + } + + #[test] + fn bridges_raw_reasoning_snapshot_items_when_enabled() { + let events = turn_snapshot_events( + ThreadId::new(), + &Turn { + id: "turn-complete".to_string(), + items: vec![ThreadItem::Reasoning { + id: "reasoning-1".to_string(), + summary: vec!["Need to inspect config".to_string()], + content: vec!["hidden chain".to_string()], + }], + status: TurnStatus::Completed, + error: None, + }, + /*show_raw_agent_reasoning*/ true, + ); + + assert_eq!(events.len(), 4); + assert!(matches!(events[0].msg, EventMsg::TurnStarted(_))); + let EventMsg::AgentReasoning(reasoning) = &events[1].msg else { + panic!("expected reasoning replay"); + }; + assert_eq!(reasoning.text, "Need to inspect config"); + let EventMsg::AgentReasoningRawContent(raw_reasoning) = &events[2].msg else { + panic!("expected raw reasoning replay"); + }; + assert_eq!(raw_reasoning.text, "hidden chain"); + assert!(matches!(events[3].msg, EventMsg::TurnComplete(_))); + } } diff --git a/codex-rs/tui_app_server/src/app_server_session.rs b/codex-rs/tui_app_server/src/app_server_session.rs index 19c882caf0..514005193b 100644 --- a/codex-rs/tui_app_server/src/app_server_session.rs +++ b/codex-rs/tui_app_server/src/app_server_session.rs @@ -55,15 +55,6 @@ use codex_app_server_protocol::TurnSteerResponse; use codex_core::config::Config; use codex_otel::TelemetryAuthMode; use codex_protocol::ThreadId; -use codex_protocol::items::AgentMessageContent; -use codex_protocol::items::AgentMessageItem; -use codex_protocol::items::ContextCompactionItem; -use codex_protocol::items::ImageGenerationItem; -use codex_protocol::items::PlanItem; -use codex_protocol::items::ReasoningItem; -use codex_protocol::items::TurnItem; -use codex_protocol::items::UserMessageItem; -use codex_protocol::items::WebSearchItem; use codex_protocol::openai_models::ModelAvailabilityNux; use codex_protocol::openai_models::ModelPreset; use codex_protocol::openai_models::ModelUpgrade; @@ -73,8 +64,6 @@ use codex_protocol::protocol::ConversationAudioParams; use codex_protocol::protocol::ConversationStartParams; use codex_protocol::protocol::ConversationTextParams; use codex_protocol::protocol::CreditsSnapshot; -use codex_protocol::protocol::EventMsg; -use codex_protocol::protocol::ItemCompletedEvent; use codex_protocol::protocol::RateLimitSnapshot; use codex_protocol::protocol::RateLimitWindow; use codex_protocol::protocol::ReviewRequest; @@ -123,8 +112,18 @@ impl ThreadParamsMode { } } +/// Result of starting, resuming, or forking an app-server thread. +/// +/// Carries the full `Thread` snapshot returned by the server alongside the +/// derived `SessionConfiguredEvent`. The snapshot's `turns` are used by +/// `App::restore_started_app_server_thread` to seed the event store and +/// replay transcript history — this is the only source of prior-turn data +/// for remote sessions, where historical websocket notifications are not +/// re-sent after the handshake. pub(crate) struct AppServerStartedThread { + pub(crate) thread: Thread, pub(crate) session_configured: SessionConfiguredEvent, + pub(crate) show_raw_agent_reasoning: bool, } impl AppServerSession { @@ -267,7 +266,7 @@ impl AppServerSession { }) .await .wrap_err("thread/start failed during TUI bootstrap")?; - started_thread_from_start_response(&response) + started_thread_from_start_response(response) } pub(crate) async fn resume_thread( @@ -289,7 +288,7 @@ impl AppServerSession { }) .await .wrap_err("thread/resume failed during TUI bootstrap")?; - started_thread_from_resume_response(&response, show_raw_agent_reasoning) + started_thread_from_resume_response(response, show_raw_agent_reasoning) } pub(crate) async fn fork_thread( @@ -311,7 +310,7 @@ impl AppServerSession { }) .await .wrap_err("thread/fork failed during TUI bootstrap")?; - started_thread_from_fork_response(&response, show_raw_agent_reasoning) + started_thread_from_fork_response(response, show_raw_agent_reasoning) } fn thread_params_mode(&self) -> ThreadParamsMode { @@ -836,46 +835,42 @@ fn thread_cwd_from_config(config: &Config, thread_params_mode: ThreadParamsMode) } fn started_thread_from_start_response( - response: &ThreadStartResponse, + response: ThreadStartResponse, ) -> Result { - let session_configured = session_configured_from_thread_start_response(response) + let session_configured = session_configured_from_thread_start_response(&response) .map_err(color_eyre::eyre::Report::msg)?; - Ok(AppServerStartedThread { session_configured }) + Ok(AppServerStartedThread { + thread: response.thread, + session_configured, + show_raw_agent_reasoning: false, + }) } fn started_thread_from_resume_response( - response: &ThreadResumeResponse, + response: ThreadResumeResponse, show_raw_agent_reasoning: bool, ) -> Result { - let session_configured = session_configured_from_thread_resume_response(response) + let session_configured = session_configured_from_thread_resume_response(&response) .map_err(color_eyre::eyre::Report::msg)?; + let thread = response.thread; Ok(AppServerStartedThread { - session_configured: SessionConfiguredEvent { - initial_messages: thread_initial_messages( - &session_configured.session_id, - &response.thread.turns, - show_raw_agent_reasoning, - ), - ..session_configured - }, + thread, + session_configured, + show_raw_agent_reasoning, }) } fn started_thread_from_fork_response( - response: &ThreadForkResponse, + response: ThreadForkResponse, show_raw_agent_reasoning: bool, ) -> Result { - let session_configured = session_configured_from_thread_fork_response(response) + let session_configured = session_configured_from_thread_fork_response(&response) .map_err(color_eyre::eyre::Report::msg)?; + let thread = response.thread; Ok(AppServerStartedThread { - session_configured: SessionConfiguredEvent { - initial_messages: thread_initial_messages( - &session_configured.session_id, - &response.thread.turns, - show_raw_agent_reasoning, - ), - ..session_configured - }, + thread, + session_configured, + show_raw_agent_reasoning, }) } @@ -992,121 +987,6 @@ fn session_configured_from_thread_response( }) } -fn thread_initial_messages( - thread_id: &ThreadId, - turns: &[codex_app_server_protocol::Turn], - show_raw_agent_reasoning: bool, -) -> Option> { - let events: Vec = turns - .iter() - .flat_map(|turn| turn_initial_messages(thread_id, turn, show_raw_agent_reasoning)) - .collect(); - (!events.is_empty()).then_some(events) -} - -fn turn_initial_messages( - thread_id: &ThreadId, - turn: &codex_app_server_protocol::Turn, - show_raw_agent_reasoning: bool, -) -> Vec { - turn.items - .iter() - .cloned() - .filter_map(app_server_thread_item_to_core) - .flat_map(|item| match item { - TurnItem::UserMessage(item) => vec![item.as_legacy_event()], - TurnItem::Plan(item) => vec![EventMsg::ItemCompleted(ItemCompletedEvent { - thread_id: *thread_id, - turn_id: turn.id.clone(), - item: TurnItem::Plan(item), - })], - item => item.as_legacy_events(show_raw_agent_reasoning), - }) - .collect() -} - -fn app_server_thread_item_to_core(item: codex_app_server_protocol::ThreadItem) -> Option { - match item { - codex_app_server_protocol::ThreadItem::UserMessage { id, content } => { - Some(TurnItem::UserMessage(UserMessageItem { - id, - content: content - .into_iter() - .map(codex_app_server_protocol::UserInput::into_core) - .collect(), - })) - } - codex_app_server_protocol::ThreadItem::AgentMessage { id, text, phase } => { - Some(TurnItem::AgentMessage(AgentMessageItem { - id, - content: vec![AgentMessageContent::Text { text }], - phase, - })) - } - codex_app_server_protocol::ThreadItem::Plan { id, text } => { - Some(TurnItem::Plan(PlanItem { id, text })) - } - codex_app_server_protocol::ThreadItem::Reasoning { - id, - summary, - content, - } => Some(TurnItem::Reasoning(ReasoningItem { - id, - summary_text: summary, - raw_content: content, - })), - codex_app_server_protocol::ThreadItem::WebSearch { id, query, action } => { - Some(TurnItem::WebSearch(WebSearchItem { - id, - query, - action: app_server_web_search_action_to_core(action?)?, - })) - } - codex_app_server_protocol::ThreadItem::ImageGeneration { - id, - status, - revised_prompt, - result, - } => Some(TurnItem::ImageGeneration(ImageGenerationItem { - id, - status, - revised_prompt, - result, - saved_path: None, - })), - codex_app_server_protocol::ThreadItem::ContextCompaction { id } => { - Some(TurnItem::ContextCompaction(ContextCompactionItem { id })) - } - codex_app_server_protocol::ThreadItem::CommandExecution { .. } - | codex_app_server_protocol::ThreadItem::FileChange { .. } - | codex_app_server_protocol::ThreadItem::McpToolCall { .. } - | codex_app_server_protocol::ThreadItem::DynamicToolCall { .. } - | codex_app_server_protocol::ThreadItem::CollabAgentToolCall { .. } - | codex_app_server_protocol::ThreadItem::ImageView { .. } - | codex_app_server_protocol::ThreadItem::EnteredReviewMode { .. } - | codex_app_server_protocol::ThreadItem::ExitedReviewMode { .. } => None, - } -} - -fn app_server_web_search_action_to_core( - action: codex_app_server_protocol::WebSearchAction, -) -> Option { - match action { - codex_app_server_protocol::WebSearchAction::Search { query, queries } => { - Some(codex_protocol::models::WebSearchAction::Search { query, queries }) - } - codex_app_server_protocol::WebSearchAction::OpenPage { url } => { - Some(codex_protocol::models::WebSearchAction::OpenPage { url }) - } - codex_app_server_protocol::WebSearchAction::FindInPage { url, pattern } => { - Some(codex_protocol::models::WebSearchAction::FindInPage { url, pattern }) - } - codex_app_server_protocol::WebSearchAction::Other => { - Some(codex_protocol::models::WebSearchAction::Other) - } - } -} - fn app_server_rate_limit_snapshots_to_core( response: GetAccountRateLimitsResponse, ) -> Vec { @@ -1204,7 +1084,7 @@ mod tests { } #[test] - fn resume_response_restores_initial_messages_from_turn_items() { + fn resume_response_relies_on_snapshot_replay_not_initial_messages() { let thread_id = ThreadId::new(); let response = ThreadResumeResponse { thread: codex_app_server_protocol::Thread { @@ -1254,29 +1134,11 @@ mod tests { }; let started = - started_thread_from_resume_response(&response, /*show_raw_agent_reasoning*/ false) + started_thread_from_resume_response(response, /*show_raw_agent_reasoning*/ false) .expect("resume response should map"); - let initial_messages = started - .session_configured - .initial_messages - .expect("resume response should restore replay history"); - - assert_eq!(initial_messages.len(), 2); - match &initial_messages[0] { - EventMsg::UserMessage(event) => { - assert_eq!(event.message, "hello from history"); - assert_eq!(event.images.as_ref(), Some(&Vec::new())); - assert!(event.local_images.is_empty()); - assert!(event.text_elements.is_empty()); - } - other => panic!("expected replayed user message, got {other:?}"), - } - match &initial_messages[1] { - EventMsg::AgentMessage(event) => { - assert_eq!(event.message, "assistant reply"); - assert_eq!(event.phase, None); - } - other => panic!("expected replayed agent message, got {other:?}"), - } + assert!(started.session_configured.initial_messages.is_none()); + assert!(!started.show_raw_agent_reasoning); + assert_eq!(started.thread.turns.len(), 1); + assert_eq!(started.thread.turns[0].items.len(), 2); } } diff --git a/codex-rs/tui_app_server/src/chatwidget.rs b/codex-rs/tui_app_server/src/chatwidget.rs index 0b4fb7c184..e4101bb11e 100644 --- a/codex-rs/tui_app_server/src/chatwidget.rs +++ b/codex-rs/tui_app_server/src/chatwidget.rs @@ -719,6 +719,10 @@ pub(crate) struct ChatWidget { // When resuming an existing session (selected via resume picker), avoid an // immediate redraw on SessionConfigured to prevent a gratuitous UI flicker. suppress_session_configured_redraw: bool, + // During snapshot restore, defer startup prompt submission until replayed + // history has been rendered so resumed/forked prompts keep chronological + // order. + suppress_initial_user_message_submit: bool, // User messages queued while a turn is in progress queued_user_messages: VecDeque, // Steers already submitted to core but not yet committed into history. @@ -1427,7 +1431,11 @@ impl ChatWidget { self.prefetch_connectors(); } if let Some(user_message) = self.initial_user_message.take() { - self.submit_user_message(user_message); + if self.suppress_initial_user_message_submit { + self.initial_user_message = Some(user_message); + } else { + self.submit_user_message(user_message); + } } if let Some(forked_from_id) = forked_from_id { self.emit_forked_thread_event(forked_from_id); @@ -1437,6 +1445,21 @@ impl ChatWidget { } } + pub(crate) fn set_initial_user_message_submit_suppressed(&mut self, suppressed: bool) { + self.suppress_initial_user_message_submit = suppressed; + } + + #[cfg(test)] + pub(crate) fn set_initial_user_message_for_test(&mut self, user_message: Option) { + self.initial_user_message = user_message; + } + + pub(crate) fn submit_initial_user_message_if_pending(&mut self) { + if let Some(user_message) = self.initial_user_message.take() { + self.submit_user_message(user_message); + } + } + fn emit_forked_thread_event(&self, forked_from_id: ThreadId) { let app_event_tx = self.app_event_tx.clone(); let codex_home = self.config.codex_home.clone(); @@ -3624,6 +3647,7 @@ impl ChatWidget { show_welcome_banner: is_first_run, startup_tooltip_override, suppress_session_configured_redraw: false, + suppress_initial_user_message_submit: false, pending_notification: None, quit_shortcut_expires_at: None, quit_shortcut_key: None, @@ -3816,6 +3840,7 @@ impl ChatWidget { show_welcome_banner: false, startup_tooltip_override: None, suppress_session_configured_redraw: true, + suppress_initial_user_message_submit: false, pending_notification: None, quit_shortcut_expires_at: None, quit_shortcut_key: None, diff --git a/codex-rs/tui_app_server/src/chatwidget/tests.rs b/codex-rs/tui_app_server/src/chatwidget/tests.rs index 07770182b2..6468e3de47 100644 --- a/codex-rs/tui_app_server/src/chatwidget/tests.rs +++ b/codex-rs/tui_app_server/src/chatwidget/tests.rs @@ -1898,6 +1898,7 @@ async fn make_chatwidget_manual( submit_pending_steers_after_interrupt: false, queued_message_edit_binding: crate::key_hint::alt(KeyCode::Up), suppress_session_configured_redraw: false, + suppress_initial_user_message_submit: false, pending_notification: None, quit_shortcut_expires_at: None, quit_shortcut_key: None, From f26ad3c92c3ac1bd1c63325d74924053d3cd0c01 Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Tue, 17 Mar 2026 10:52:16 -0700 Subject: [PATCH 07/11] Fix fuzzy search notification buffering in app-server tests (#14955) ## What is flaky `codex-rs/app-server/tests/suite/fuzzy_file_search.rs` intermittently loses the expected `fuzzyFileSearch/sessionUpdated` and `fuzzyFileSearch/sessionCompleted` notifications when multiple fuzzy-search sessions are active and CI delivers notifications out of order. ## Why it was flaky The wait helpers were keyed only by JSON-RPC method name. - `wait_for_session_updated` consumed the next `fuzzyFileSearch/sessionUpdated` notification even when it belonged to a different search session. - `wait_for_session_completed` did the same for `fuzzyFileSearch/sessionCompleted`. - Once an unmatched notification was read, it was dropped permanently instead of buffered. - That meant a valid completion for the target search could arrive slightly early, be consumed by the wrong waiter, and disappear before the test started waiting for it. The result depended on notification ordering and runner scheduling instead of on the actual product behavior. ## How this PR fixes it - Add a buffered notification reader in `codex-rs/app-server/tests/common/mcp_process.rs`. - Match fuzzy-search notifications on the identifying payload fields instead of matching only on method name. - Preserve unmatched notifications in the in-process queue so later waiters can still consume them. - Include pending notification methods in timeout failures to make future diagnosis concrete. ## Why this fix fixes the flakiness The test now behaves like a real consumer of an out-of-order event stream: notifications for other sessions stay buffered until the correct waiter asks for them. Reordering no longer loses the target event, so the test result is determined by whether the server emitted the right notifications, not by which one happened to be read first. Co-authored-by: Ahmed Ibrahim <219906144+aibrahim-oai@users.noreply.github.com> Co-authored-by: Codex --- .../app-server/tests/common/mcp_process.rs | 35 ++++++ .../tests/suite/fuzzy_file_search.rs | 110 +++++++++++------- 2 files changed, 106 insertions(+), 39 deletions(-) diff --git a/codex-rs/app-server/tests/common/mcp_process.rs b/codex-rs/app-server/tests/common/mcp_process.rs index b2289493ac..430a400a2c 100644 --- a/codex-rs/app-server/tests/common/mcp_process.rs +++ b/codex-rs/app-server/tests/common/mcp_process.rs @@ -1031,6 +1031,31 @@ impl McpProcess { Ok(notification) } + pub async fn read_stream_until_matching_notification( + &mut self, + description: &str, + predicate: F, + ) -> anyhow::Result + where + F: Fn(&JSONRPCNotification) -> bool, + { + eprintln!("in read_stream_until_matching_notification({description})"); + + let message = self + .read_stream_until_message(|message| { + matches!( + message, + JSONRPCMessage::Notification(notification) if predicate(notification) + ) + }) + .await?; + + let JSONRPCMessage::Notification(notification) = message else { + unreachable!("expected JSONRPCMessage::Notification, got {message:?}"); + }; + Ok(notification) + } + pub async fn read_next_message(&mut self) -> anyhow::Result { self.read_stream_until_message(|_| true).await } @@ -1043,6 +1068,16 @@ impl McpProcess { self.pending_messages.clear(); } + pub fn pending_notification_methods(&self) -> Vec { + self.pending_messages + .iter() + .filter_map(|message| match message { + JSONRPCMessage::Notification(notification) => Some(notification.method.clone()), + _ => None, + }) + .collect() + } + /// Reads the stream until a message matches `predicate`, buffering any non-matching messages /// for later reads. async fn read_stream_until_message(&mut self, predicate: F) -> anyhow::Result diff --git a/codex-rs/app-server/tests/suite/fuzzy_file_search.rs b/codex-rs/app-server/tests/suite/fuzzy_file_search.rs index 7341a5a5f7..0070c2b30b 100644 --- a/codex-rs/app-server/tests/suite/fuzzy_file_search.rs +++ b/codex-rs/app-server/tests/suite/fuzzy_file_search.rs @@ -52,54 +52,86 @@ async fn wait_for_session_updated( query: &str, file_expectation: FileExpectation, ) -> Result { - for _ in 0..20 { - let notification = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message(SESSION_UPDATED_METHOD), - ) - .await??; - let params = notification - .params - .ok_or_else(|| anyhow!("missing notification params"))?; - let payload = serde_json::from_value::(params)?; - if payload.session_id != session_id || payload.query != query { - continue; + let description = format!("session update for sessionId={session_id}, query={query}"); + let notification = match timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_matching_notification(&description, |notification| { + if notification.method != SESSION_UPDATED_METHOD { + return false; + } + let Some(params) = notification.params.as_ref() else { + return false; + }; + let Ok(payload) = + serde_json::from_value::(params.clone()) + else { + return false; + }; + let files_match = match file_expectation { + FileExpectation::Any => true, + FileExpectation::Empty => payload.files.is_empty(), + FileExpectation::NonEmpty => !payload.files.is_empty(), + }; + payload.session_id == session_id && payload.query == query && files_match + }), + ) + .await + { + Ok(result) => result?, + Err(_) => { + anyhow::bail!( + "timed out waiting for {description}; buffered notifications={:?}", + mcp.pending_notification_methods() + ) } - let files_match = match file_expectation { - FileExpectation::Any => true, - FileExpectation::Empty => payload.files.is_empty(), - FileExpectation::NonEmpty => !payload.files.is_empty(), - }; - if files_match { - return Ok(payload); - } - } - anyhow::bail!( - "did not receive expected session update for sessionId={session_id}, query={query}" - ); + }; + let params = notification + .params + .ok_or_else(|| anyhow!("missing notification params"))?; + Ok(serde_json::from_value::< + FuzzyFileSearchSessionUpdatedNotification, + >(params)?) } async fn wait_for_session_completed( mcp: &mut McpProcess, session_id: &str, ) -> Result { - for _ in 0..20 { - let notification = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message(SESSION_COMPLETED_METHOD), - ) - .await??; - let params = notification - .params - .ok_or_else(|| anyhow!("missing notification params"))?; - let payload = - serde_json::from_value::(params)?; - if payload.session_id == session_id { - return Ok(payload); + let description = format!("session completion for sessionId={session_id}"); + let notification = match timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_matching_notification(&description, |notification| { + if notification.method != SESSION_COMPLETED_METHOD { + return false; + } + let Some(params) = notification.params.as_ref() else { + return false; + }; + let Ok(payload) = serde_json::from_value::( + params.clone(), + ) else { + return false; + }; + payload.session_id == session_id + }), + ) + .await + { + Ok(result) => result?, + Err(_) => { + anyhow::bail!( + "timed out waiting for {description}; buffered notifications={:?}", + mcp.pending_notification_methods() + ) } - } + }; - anyhow::bail!("did not receive expected session completion for sessionId={session_id}"); + let params = notification + .params + .ok_or_else(|| anyhow!("missing notification params"))?; + Ok(serde_json::from_value::< + FuzzyFileSearchSessionCompletedNotification, + >(params)?) } async fn assert_update_request_fails_for_missing_session( From d484bb57d9baea4603df0a89ad4f602cee79871d Mon Sep 17 00:00:00 2001 From: jif-oai Date: Tue, 17 Mar 2026 17:59:27 +0000 Subject: [PATCH 08/11] feat: add suffix to shell snapshot name (#14938) https://github.com/openai/codex/issues/14906 --- codex-rs/core/src/shell_snapshot.rs | 27 ++++++---- codex-rs/core/src/shell_snapshot_tests.rs | 66 +++++++++++++++++++++-- 2 files changed, 80 insertions(+), 13 deletions(-) diff --git a/codex-rs/core/src/shell_snapshot.rs b/codex-rs/core/src/shell_snapshot.rs index 3bf1515add..29b50cb9e8 100644 --- a/codex-rs/core/src/shell_snapshot.rs +++ b/codex-rs/core/src/shell_snapshot.rs @@ -120,13 +120,13 @@ impl ShellSnapshot { ShellType::PowerShell => "ps1", _ => "sh", }; - let path = codex_home - .join(SNAPSHOT_DIR) - .join(format!("{session_id}.{extension}")); let nonce = SystemTime::now() .duration_since(SystemTime::UNIX_EPOCH) .map(|duration| duration.as_nanos()) .unwrap_or(0); + let path = codex_home + .join(SNAPSHOT_DIR) + .join(format!("{session_id}.{nonce}.{extension}")); let temp_path = codex_home .join(SNAPSHOT_DIR) .join(format!("{session_id}.tmp-{nonce}")); @@ -510,12 +510,9 @@ pub async fn cleanup_stale_snapshots(codex_home: &Path, active_session_id: Threa let file_name = entry.file_name(); let file_name = file_name.to_string_lossy(); - let (session_id, _) = match file_name.rsplit_once('.') { - Some((stem, ext)) => (stem, ext), - None => { - remove_snapshot_file(&path).await; - continue; - } + let Some(session_id) = snapshot_session_id_from_file_name(&file_name) else { + remove_snapshot_file(&path).await; + continue; }; if session_id == active_session_id { continue; @@ -556,6 +553,18 @@ async fn remove_snapshot_file(path: &Path) { } } +fn snapshot_session_id_from_file_name(file_name: &str) -> Option<&str> { + let (stem, extension) = file_name.rsplit_once('.')?; + match extension { + "sh" | "ps1" => Some( + stem.split_once('.') + .map_or(stem, |(session_id, _generation)| session_id), + ), + _ if extension.starts_with("tmp-") => Some(stem), + _ => None, + } +} + #[cfg(test)] #[path = "shell_snapshot_tests.rs"] mod tests; diff --git a/codex-rs/core/src/shell_snapshot_tests.rs b/codex-rs/core/src/shell_snapshot_tests.rs index 558a40e2e2..2819f67d3b 100644 --- a/codex-rs/core/src/shell_snapshot_tests.rs +++ b/codex-rs/core/src/shell_snapshot_tests.rs @@ -98,6 +98,28 @@ fn strip_snapshot_preamble_requires_marker() { assert!(result.is_err()); } +#[test] +fn snapshot_file_name_parser_supports_legacy_and_suffixed_names() { + let session_id = "019cf82b-6a62-7700-bbbd-46909794ef89"; + + assert_eq!( + snapshot_session_id_from_file_name(&format!("{session_id}.sh")), + Some(session_id) + ); + assert_eq!( + snapshot_session_id_from_file_name(&format!("{session_id}.123.sh")), + Some(session_id) + ); + assert_eq!( + snapshot_session_id_from_file_name(&format!("{session_id}.tmp-123")), + Some(session_id) + ); + assert_eq!( + snapshot_session_id_from_file_name("not-a-snapshot.txt"), + None + ); +} + #[cfg(unix)] #[test] fn bash_snapshot_filters_invalid_exports() -> Result<()> { @@ -186,6 +208,42 @@ async fn try_new_creates_and_deletes_snapshot_file() -> Result<()> { Ok(()) } +#[cfg(unix)] +#[tokio::test] +async fn try_new_uses_distinct_generation_paths() -> Result<()> { + let dir = tempdir()?; + let session_id = ThreadId::new(); + let shell = Shell { + shell_type: ShellType::Bash, + shell_path: PathBuf::from("/bin/bash"), + shell_snapshot: crate::shell::empty_shell_snapshot_receiver(), + }; + + let initial_snapshot = ShellSnapshot::try_new(dir.path(), session_id, dir.path(), &shell) + .await + .expect("initial snapshot should be created"); + let refreshed_snapshot = ShellSnapshot::try_new(dir.path(), session_id, dir.path(), &shell) + .await + .expect("refreshed snapshot should be created"); + let initial_path = initial_snapshot.path.clone(); + let refreshed_path = refreshed_snapshot.path.clone(); + + assert_ne!(initial_path, refreshed_path); + assert_eq!(initial_path.exists(), true); + assert_eq!(refreshed_path.exists(), true); + + drop(initial_snapshot); + + assert_eq!(initial_path.exists(), false); + assert_eq!(refreshed_path.exists(), true); + + drop(refreshed_snapshot); + + assert_eq!(refreshed_path.exists(), false); + + Ok(()) +} + #[cfg(unix)] #[tokio::test] async fn snapshot_shell_does_not_inherit_stdin() -> Result<()> { @@ -339,8 +397,8 @@ async fn cleanup_stale_snapshots_removes_orphans_and_keeps_live() -> Result<()> let live_session = ThreadId::new(); let orphan_session = ThreadId::new(); - let live_snapshot = snapshot_dir.join(format!("{live_session}.sh")); - let orphan_snapshot = snapshot_dir.join(format!("{orphan_session}.sh")); + let live_snapshot = snapshot_dir.join(format!("{live_session}.123.sh")); + let orphan_snapshot = snapshot_dir.join(format!("{orphan_session}.456.sh")); let invalid_snapshot = snapshot_dir.join("not-a-snapshot.txt"); write_rollout_stub(codex_home, live_session).await?; @@ -365,7 +423,7 @@ async fn cleanup_stale_snapshots_removes_stale_rollouts() -> Result<()> { fs::create_dir_all(&snapshot_dir).await?; let stale_session = ThreadId::new(); - let stale_snapshot = snapshot_dir.join(format!("{stale_session}.sh")); + let stale_snapshot = snapshot_dir.join(format!("{stale_session}.123.sh")); let rollout_path = write_rollout_stub(codex_home, stale_session).await?; fs::write(&stale_snapshot, "stale").await?; @@ -386,7 +444,7 @@ async fn cleanup_stale_snapshots_skips_active_session() -> Result<()> { fs::create_dir_all(&snapshot_dir).await?; let active_session = ThreadId::new(); - let active_snapshot = snapshot_dir.join(format!("{active_session}.sh")); + let active_snapshot = snapshot_dir.join(format!("{active_session}.123.sh")); let rollout_path = write_rollout_stub(codex_home, active_session).await?; fs::write(&active_snapshot, "active").await?; From 0d531c05f2cc497d29da8e478f6770850cdb51bc Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Tue, 17 Mar 2026 11:09:12 -0700 Subject: [PATCH 09/11] Fix code mode yield startup race (#14959) --- codex-rs/core/src/tools/code_mode/runner.cjs | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/codex-rs/core/src/tools/code_mode/runner.cjs b/codex-rs/core/src/tools/code_mode/runner.cjs index 9ebb9c9882..48b9d5a9af 100644 --- a/codex-rs/core/src/tools/code_mode/runner.cjs +++ b/codex-rs/core/src/tools/code_mode/runner.cjs @@ -465,6 +465,7 @@ function codeModeWorkerMain() { writable: false, }); + parentPort.postMessage({ type: 'started' }); try { await runModule(context, start, callTool, helpers); parentPort.postMessage({ @@ -639,6 +640,10 @@ function startSession(protocol, sessions, start) { content_items: [], default_yield_time_ms: normalizeYieldTime(start.default_yield_time_ms), id: start.cell_id, + initial_yield_time_ms: + start.yield_time_ms == null + ? normalizeYieldTime(start.default_yield_time_ms) + : normalizeYieldTime(start.yield_time_ms), initial_yield_timer: null, initial_yield_triggered: false, max_output_tokens_per_exec_call: maxOutputTokensPerExecCall, @@ -651,11 +656,6 @@ function startSession(protocol, sessions, start) { }), }; sessions.set(session.id, session); - const initialYieldTime = - start.yield_time_ms == null - ? session.default_yield_time_ms - : normalizeYieldTime(start.yield_time_ms); - scheduleInitialYield(protocol, session, initialYieldTime); session.worker.on('message', (message) => { void handleWorkerMessage(protocol, sessions, session, message).catch((error) => { @@ -694,6 +694,11 @@ async function handleWorkerMessage(protocol, sessions, session, message) { return; } + if (message.type === 'started') { + scheduleInitialYield(protocol, session, session.initial_yield_time_ms); + return; + } + if (message.type === 'yield') { void sendYielded(protocol, session); return; From 904dbd414f223027ecdb3d54a8444d3c94395aa6 Mon Sep 17 00:00:00 2001 From: Keyan Zhang <2268452+keyz@users.noreply.github.com> Date: Tue, 17 Mar 2026 11:19:42 -0700 Subject: [PATCH 10/11] generate an internal json schema for `RolloutLine` (#14434) ### Why i'm working on something that parses and analyzes codex rollout logs, and i'd like to have a schema for generating a parser/validator. `codex app-server generate-internal-json-schema` writes an `RolloutLine.json` file while doing this, i noticed we have a writer <> reader mismatch issue on `FunctionCallOutputPayload` and reasoning item ID -- added some schemars annotations to fix those ### Test ``` $ just codex app-server generate-internal-json-schema --out ./foo ``` generates an `RolloutLine.json` file, which i validated against jsonl files on disk `just codex app-server --help` doesn't expose the `generate-internal-json-schema` option by default, but you can do `just codex app-server generate-internal-json-schema --help` if you know the command everything else still works --------- Co-authored-by: Codex --- .../schema/json/ClientRequest.json | 27 ++----------------- .../codex_app_server_protocol.schemas.json | 27 ++----------------- .../codex_app_server_protocol.v2.schemas.json | 27 ++----------------- .../RawResponseItemCompletedNotification.json | 27 ++----------------- .../schema/json/v2/ThreadResumeParams.json | 27 ++----------------- .../typescript/FunctionCallOutputPayload.ts | 12 --------- .../schema/typescript/ResponseItem.ts | 4 +-- .../schema/typescript/index.ts | 1 - codex-rs/app-server-protocol/src/export.rs | 7 +++++ codex-rs/app-server-protocol/src/lib.rs | 1 + codex-rs/cli/src/main.rs | 15 +++++++++++ codex-rs/protocol/src/models.rs | 9 +++++++ 12 files changed, 44 insertions(+), 140 deletions(-) delete mode 100644 codex-rs/app-server-protocol/schema/typescript/FunctionCallOutputPayload.ts diff --git a/codex-rs/app-server-protocol/schema/json/ClientRequest.json b/codex-rs/app-server-protocol/schema/json/ClientRequest.json index dd5c955cfa..9c2004f057 100644 --- a/codex-rs/app-server-protocol/schema/json/ClientRequest.json +++ b/codex-rs/app-server-protocol/schema/json/ClientRequest.json @@ -871,24 +871,6 @@ } ] }, - "FunctionCallOutputPayload": { - "description": "The payload we send back to OpenAI when reporting a tool call result.\n\n`body` serializes directly as the wire value for `function_call_output.output`. `success` remains internal metadata for downstream handling.", - "properties": { - "body": { - "$ref": "#/definitions/FunctionCallOutputBody" - }, - "success": { - "type": [ - "boolean", - "null" - ] - } - }, - "required": [ - "body" - ], - "type": "object" - }, "FuzzyFileSearchParams": { "properties": { "cancellationToken": { @@ -1583,10 +1565,6 @@ "null" ] }, - "id": { - "type": "string", - "writeOnly": true - }, "summary": { "items": { "$ref": "#/definitions/ReasoningItemReasoningSummary" @@ -1602,7 +1580,6 @@ } }, "required": [ - "id", "summary", "type" ], @@ -1736,7 +1713,7 @@ "type": "string" }, "output": { - "$ref": "#/definitions/FunctionCallOutputPayload" + "$ref": "#/definitions/FunctionCallOutputBody" }, "type": { "enum": [ @@ -1801,7 +1778,7 @@ "type": "string" }, "output": { - "$ref": "#/definitions/FunctionCallOutputPayload" + "$ref": "#/definitions/FunctionCallOutputBody" }, "type": { "enum": [ diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json index 4709bae11c..e370546dc2 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json @@ -7731,24 +7731,6 @@ } ] }, - "FunctionCallOutputPayload": { - "description": "The payload we send back to OpenAI when reporting a tool call result.\n\n`body` serializes directly as the wire value for `function_call_output.output`. `success` remains internal metadata for downstream handling.", - "properties": { - "body": { - "$ref": "#/definitions/v2/FunctionCallOutputBody" - }, - "success": { - "type": [ - "boolean", - "null" - ] - } - }, - "required": [ - "body" - ], - "type": "object" - }, "GetAccountParams": { "$schema": "http://json-schema.org/draft-07/schema#", "properties": { @@ -10195,10 +10177,6 @@ "null" ] }, - "id": { - "type": "string", - "writeOnly": true - }, "summary": { "items": { "$ref": "#/definitions/v2/ReasoningItemReasoningSummary" @@ -10214,7 +10192,6 @@ } }, "required": [ - "id", "summary", "type" ], @@ -10348,7 +10325,7 @@ "type": "string" }, "output": { - "$ref": "#/definitions/v2/FunctionCallOutputPayload" + "$ref": "#/definitions/v2/FunctionCallOutputBody" }, "type": { "enum": [ @@ -10413,7 +10390,7 @@ "type": "string" }, "output": { - "$ref": "#/definitions/v2/FunctionCallOutputPayload" + "$ref": "#/definitions/v2/FunctionCallOutputBody" }, "type": { "enum": [ diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json index 601306fe14..b069d3e5e7 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json @@ -4375,24 +4375,6 @@ } ] }, - "FunctionCallOutputPayload": { - "description": "The payload we send back to OpenAI when reporting a tool call result.\n\n`body` serializes directly as the wire value for `function_call_output.output`. `success` remains internal metadata for downstream handling.", - "properties": { - "body": { - "$ref": "#/definitions/FunctionCallOutputBody" - }, - "success": { - "type": [ - "boolean", - "null" - ] - } - }, - "required": [ - "body" - ], - "type": "object" - }, "FuzzyFileSearchParams": { "$schema": "http://json-schema.org/draft-07/schema#", "properties": { @@ -6983,10 +6965,6 @@ "null" ] }, - "id": { - "type": "string", - "writeOnly": true - }, "summary": { "items": { "$ref": "#/definitions/ReasoningItemReasoningSummary" @@ -7002,7 +6980,6 @@ } }, "required": [ - "id", "summary", "type" ], @@ -7136,7 +7113,7 @@ "type": "string" }, "output": { - "$ref": "#/definitions/FunctionCallOutputPayload" + "$ref": "#/definitions/FunctionCallOutputBody" }, "type": { "enum": [ @@ -7201,7 +7178,7 @@ "type": "string" }, "output": { - "$ref": "#/definitions/FunctionCallOutputPayload" + "$ref": "#/definitions/FunctionCallOutputBody" }, "type": { "enum": [ diff --git a/codex-rs/app-server-protocol/schema/json/v2/RawResponseItemCompletedNotification.json b/codex-rs/app-server-protocol/schema/json/v2/RawResponseItemCompletedNotification.json index 19f0fe34f2..621332c3c2 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/RawResponseItemCompletedNotification.json +++ b/codex-rs/app-server-protocol/schema/json/v2/RawResponseItemCompletedNotification.json @@ -133,24 +133,6 @@ } ] }, - "FunctionCallOutputPayload": { - "description": "The payload we send back to OpenAI when reporting a tool call result.\n\n`body` serializes directly as the wire value for `function_call_output.output`. `success` remains internal metadata for downstream handling.", - "properties": { - "body": { - "$ref": "#/definitions/FunctionCallOutputBody" - }, - "success": { - "type": [ - "boolean", - "null" - ] - } - }, - "required": [ - "body" - ], - "type": "object" - }, "GhostCommit": { "description": "Details of a ghost commit created from a repository state.", "properties": { @@ -413,10 +395,6 @@ "null" ] }, - "id": { - "type": "string", - "writeOnly": true - }, "summary": { "items": { "$ref": "#/definitions/ReasoningItemReasoningSummary" @@ -432,7 +410,6 @@ } }, "required": [ - "id", "summary", "type" ], @@ -566,7 +543,7 @@ "type": "string" }, "output": { - "$ref": "#/definitions/FunctionCallOutputPayload" + "$ref": "#/definitions/FunctionCallOutputBody" }, "type": { "enum": [ @@ -631,7 +608,7 @@ "type": "string" }, "output": { - "$ref": "#/definitions/FunctionCallOutputPayload" + "$ref": "#/definitions/FunctionCallOutputBody" }, "type": { "enum": [ diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadResumeParams.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadResumeParams.json index b21c5a78ee..3524f86b26 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadResumeParams.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadResumeParams.json @@ -191,24 +191,6 @@ } ] }, - "FunctionCallOutputPayload": { - "description": "The payload we send back to OpenAI when reporting a tool call result.\n\n`body` serializes directly as the wire value for `function_call_output.output`. `success` remains internal metadata for downstream handling.", - "properties": { - "body": { - "$ref": "#/definitions/FunctionCallOutputBody" - }, - "success": { - "type": [ - "boolean", - "null" - ] - } - }, - "required": [ - "body" - ], - "type": "object" - }, "GhostCommit": { "description": "Details of a ghost commit created from a repository state.", "properties": { @@ -479,10 +461,6 @@ "null" ] }, - "id": { - "type": "string", - "writeOnly": true - }, "summary": { "items": { "$ref": "#/definitions/ReasoningItemReasoningSummary" @@ -498,7 +476,6 @@ } }, "required": [ - "id", "summary", "type" ], @@ -632,7 +609,7 @@ "type": "string" }, "output": { - "$ref": "#/definitions/FunctionCallOutputPayload" + "$ref": "#/definitions/FunctionCallOutputBody" }, "type": { "enum": [ @@ -697,7 +674,7 @@ "type": "string" }, "output": { - "$ref": "#/definitions/FunctionCallOutputPayload" + "$ref": "#/definitions/FunctionCallOutputBody" }, "type": { "enum": [ diff --git a/codex-rs/app-server-protocol/schema/typescript/FunctionCallOutputPayload.ts b/codex-rs/app-server-protocol/schema/typescript/FunctionCallOutputPayload.ts deleted file mode 100644 index 6376c5b8eb..0000000000 --- a/codex-rs/app-server-protocol/schema/typescript/FunctionCallOutputPayload.ts +++ /dev/null @@ -1,12 +0,0 @@ -// GENERATED CODE! DO NOT MODIFY BY HAND! - -// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. -import type { FunctionCallOutputBody } from "./FunctionCallOutputBody"; - -/** - * The payload we send back to OpenAI when reporting a tool call result. - * - * `body` serializes directly as the wire value for `function_call_output.output`. - * `success` remains internal metadata for downstream handling. - */ -export type FunctionCallOutputPayload = { body: FunctionCallOutputBody, success: boolean | null, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/ResponseItem.ts b/codex-rs/app-server-protocol/schema/typescript/ResponseItem.ts index 2464037a50..48e6b69d73 100644 --- a/codex-rs/app-server-protocol/schema/typescript/ResponseItem.ts +++ b/codex-rs/app-server-protocol/schema/typescript/ResponseItem.ts @@ -2,7 +2,7 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. import type { ContentItem } from "./ContentItem"; -import type { FunctionCallOutputPayload } from "./FunctionCallOutputPayload"; +import type { FunctionCallOutputBody } from "./FunctionCallOutputBody"; import type { GhostCommit } from "./GhostCommit"; import type { LocalShellAction } from "./LocalShellAction"; import type { LocalShellStatus } from "./LocalShellStatus"; @@ -15,4 +15,4 @@ export type ResponseItem = { "type": "message", role: string, content: Array Result<()> { generate_json_with_experimental(out_dir, /*experimental_api*/ false) } +pub fn generate_internal_json_schema(out_dir: &Path) -> Result<()> { + ensure_dir(out_dir)?; + write_json_schema::(out_dir, "RolloutLine")?; + Ok(()) +} + pub fn generate_json_with_experimental(out_dir: &Path, experimental_api: bool) -> Result<()> { ensure_dir(out_dir)?; let envelope_emitters: Vec = vec![ diff --git a/codex-rs/app-server-protocol/src/lib.rs b/codex-rs/app-server-protocol/src/lib.rs index 067dcb0369..3c5fa6dc2e 100644 --- a/codex-rs/app-server-protocol/src/lib.rs +++ b/codex-rs/app-server-protocol/src/lib.rs @@ -6,6 +6,7 @@ mod schema_fixtures; pub use experimental_api::*; pub use export::GenerateTsOptions; +pub use export::generate_internal_json_schema; pub use export::generate_json; pub use export::generate_json_with_experimental; pub use export::generate_ts; diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index 972c2854f8..05b568b7b7 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -351,12 +351,17 @@ struct AppServerCommand { } #[derive(Debug, clap::Subcommand)] +#[allow(clippy::enum_variant_names)] enum AppServerSubcommand { /// [experimental] Generate TypeScript bindings for the app server protocol. GenerateTs(GenerateTsCommand), /// [experimental] Generate JSON Schema for the app server protocol. GenerateJsonSchema(GenerateJsonSchemaCommand), + + /// [internal] Generate internal JSON Schema artifacts for Codex tooling. + #[clap(hide = true)] + GenerateInternalJsonSchema(GenerateInternalJsonSchemaCommand), } #[derive(Debug, Args)] @@ -385,6 +390,13 @@ struct GenerateJsonSchemaCommand { experimental: bool, } +#[derive(Debug, Args)] +struct GenerateInternalJsonSchemaCommand { + /// Output directory where internal JSON Schema artifacts will be written + #[arg(short = 'o', long = "out", value_name = "DIR")] + out_dir: PathBuf, +} + #[derive(Debug, Parser)] struct StdioToUdsCommand { /// Path to the Unix domain socket to connect to. @@ -665,6 +677,9 @@ async fn cli_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { gen_cli.experimental, )?; } + Some(AppServerSubcommand::GenerateInternalJsonSchema(gen_cli)) => { + codex_app_server_protocol::generate_internal_json_schema(&gen_cli.out_dir)?; + } }, #[cfg(target_os = "macos")] Some(Subcommand::App(app_cli)) => { diff --git a/codex-rs/protocol/src/models.rs b/codex-rs/protocol/src/models.rs index 700c1f80e0..1e5a7445e3 100644 --- a/codex-rs/protocol/src/models.rs +++ b/codex-rs/protocol/src/models.rs @@ -231,6 +231,8 @@ pub enum ResponseInputItem { }, FunctionCallOutput { call_id: String, + #[ts(as = "FunctionCallOutputBody")] + #[schemars(with = "FunctionCallOutputBody")] output: FunctionCallOutputPayload, }, McpToolCallOutput { @@ -239,6 +241,8 @@ pub enum ResponseInputItem { }, CustomToolCallOutput { call_id: String, + #[ts(as = "FunctionCallOutputBody")] + #[schemars(with = "FunctionCallOutputBody")] output: FunctionCallOutputPayload, }, ToolSearchOutput { @@ -306,6 +310,7 @@ pub enum ResponseItem { Reasoning { #[serde(default, skip_serializing)] #[ts(skip)] + #[schemars(skip)] id: String, summary: Vec, #[serde(default, skip_serializing_if = "should_serialize_reasoning_content")] @@ -356,6 +361,8 @@ pub enum ResponseItem { // We keep this behavior centralized in `FunctionCallOutputPayload`. FunctionCallOutput { call_id: String, + #[ts(as = "FunctionCallOutputBody")] + #[schemars(with = "FunctionCallOutputBody")] output: FunctionCallOutputPayload, }, CustomToolCall { @@ -375,6 +382,8 @@ pub enum ResponseItem { // text or structured content items. CustomToolCallOutput { call_id: String, + #[ts(as = "FunctionCallOutputBody")] + #[schemars(with = "FunctionCallOutputBody")] output: FunctionCallOutputPayload, }, ToolSearchOutput { From cbec70060314259cf275baaee759f8ad59419d62 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Tue, 17 Mar 2026 14:27:32 -0700 Subject: [PATCH 11/11] feat: introduce sandbox-aware filesystem abstraction --- codex-rs/Cargo.lock | 14 ++ codex-rs/Cargo.toml | 2 + codex-rs/arg0/Cargo.toml | 1 + codex-rs/arg0/src/lib.rs | 11 + codex-rs/core/Cargo.toml | 1 + codex-rs/core/src/lib.rs | 1 + codex-rs/core/src/sandboxed_fs.rs | 234 ++++++++++++++++++ .../core/src/tools/handlers/view_image.rs | 57 +++-- codex-rs/core/tests/suite/view_image.rs | 111 +++++++++ codex-rs/fs-ops/BUILD.bazel | 6 + codex-rs/fs-ops/Cargo.toml | 22 ++ codex-rs/fs-ops/src/lib.rs | 229 +++++++++++++++++ codex-rs/protocol/src/models.rs | 24 +- codex-rs/utils/image/src/lib.rs | 9 +- 14 files changed, 702 insertions(+), 20 deletions(-) create mode 100644 codex-rs/core/src/sandboxed_fs.rs create mode 100644 codex-rs/fs-ops/BUILD.bazel create mode 100644 codex-rs/fs-ops/Cargo.toml create mode 100644 codex-rs/fs-ops/src/lib.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 84f911bd44..32d0414ebb 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1557,6 +1557,7 @@ version = "0.0.0" dependencies = [ "anyhow", "codex-apply-patch", + "codex-fs-ops", "codex-linux-sandbox", "codex-shell-escalation", "codex-utils-home-dir", @@ -1843,6 +1844,7 @@ dependencies = [ "codex-connectors", "codex-execpolicy", "codex-file-search", + "codex-fs-ops", "codex-git", "codex-hooks", "codex-keyring-store", @@ -2063,6 +2065,18 @@ dependencies = [ "tokio", ] +[[package]] +name = "codex-fs-ops" +version = "0.0.0" +dependencies = [ + "anyhow", + "base64 0.22.1", + "pretty_assertions", + "serde", + "serde_json", + "tempfile", +] + [[package]] name = "codex-git" version = "0.0.0" diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index b1e0fcf371..b7018679f4 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -11,6 +11,7 @@ members = [ "apply-patch", "arg0", "feedback", + "fs-ops", "codex-backend-openapi-models", "cloud-requirements", "cloud-tasks", @@ -107,6 +108,7 @@ codex-exec = { path = "exec" } codex-execpolicy = { path = "execpolicy" } codex-experimental-api-macros = { path = "codex-experimental-api-macros" } codex-feedback = { path = "feedback" } +codex-fs-ops = { path = "fs-ops" } codex-file-search = { path = "file-search" } codex-git = { path = "utils/git" } codex-hooks = { path = "hooks" } diff --git a/codex-rs/arg0/Cargo.toml b/codex-rs/arg0/Cargo.toml index abe7277d94..10c31eac5e 100644 --- a/codex-rs/arg0/Cargo.toml +++ b/codex-rs/arg0/Cargo.toml @@ -14,6 +14,7 @@ workspace = true [dependencies] anyhow = { workspace = true } codex-apply-patch = { workspace = true } +codex-fs-ops = { workspace = true } codex-linux-sandbox = { workspace = true } codex-shell-escalation = { workspace = true } codex-utils-home-dir = { workspace = true } diff --git a/codex-rs/arg0/src/lib.rs b/codex-rs/arg0/src/lib.rs index d8dbaaa30d..aee33e9649 100644 --- a/codex-rs/arg0/src/lib.rs +++ b/codex-rs/arg0/src/lib.rs @@ -4,6 +4,7 @@ use std::path::Path; use std::path::PathBuf; use codex_apply_patch::CODEX_CORE_APPLY_PATCH_ARG1; +use codex_fs_ops::CODEX_CORE_FS_OPS_ARG1; use codex_utils_home_dir::find_codex_home; #[cfg(unix)] use std::os::unix::fs::symlink; @@ -105,6 +106,16 @@ pub fn arg0_dispatch() -> Option { }; std::process::exit(exit_code); } + if argv1 == CODEX_CORE_FS_OPS_ARG1 { + let exit_code = match codex_fs_ops::run_from_args(args) { + Ok(()) => 0, + Err(err) => { + eprintln!("Error: failed to run fs helper: {err}"); + 1 + } + }; + std::process::exit(exit_code); + } // This modifies the environment, which is not thread-safe, so do this // before creating any threads/the Tokio runtime. diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index ef6b8a0132..ee297aa53b 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -37,6 +37,7 @@ codex-config = { workspace = true } codex-shell-command = { workspace = true } codex-skills = { workspace = true } codex-execpolicy = { workspace = true } +codex-fs-ops = { workspace = true } codex-file-search = { workspace = true } codex-git = { workspace = true } codex-hooks = { workspace = true } diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index e02a346545..04989ee8da 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -112,6 +112,7 @@ pub mod default_client; pub mod project_doc; mod rollout; pub(crate) mod safety; +mod sandboxed_fs; pub mod seatbelt; pub mod shell; pub mod shell_snapshot; diff --git a/codex-rs/core/src/sandboxed_fs.rs b/codex-rs/core/src/sandboxed_fs.rs new file mode 100644 index 0000000000..c438abd5af --- /dev/null +++ b/codex-rs/core/src/sandboxed_fs.rs @@ -0,0 +1,234 @@ +use crate::codex::Session; +use crate::codex::TurnContext; +use crate::exec::ExecExpiration; +use crate::sandboxing::CommandSpec; +use crate::sandboxing::SandboxPermissions; +use crate::sandboxing::execute_env; +use crate::sandboxing::merge_permission_profiles; +use crate::tools::sandboxing::SandboxAttempt; +use crate::tools::sandboxing::SandboxablePreference; +use base64::Engine; +use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; +use codex_fs_ops::CODEX_CORE_FS_OPS_ARG1; +use codex_fs_ops::FsError; +use codex_fs_ops::FsErrorKind; +use codex_fs_ops::FsPayload; +use codex_fs_ops::FsResponse; +use codex_protocol::models::PermissionProfile; +use std::collections::HashMap; +use std::path::Path; +use std::path::PathBuf; +use std::sync::Arc; +use std::time::Duration; + +const SANDBOXED_FS_TIMEOUT: Duration = Duration::from_secs(30); + +#[derive(Debug, thiserror::Error)] +pub(crate) enum SandboxedFsError { + #[error("failed to determine codex executable: {message}")] + ResolveExe { message: String }, + #[error("sandboxed fs helper timed out while accessing `{path}`")] + TimedOut { path: PathBuf }, + #[error("sandboxed fs helper exited with code {exit_code} while accessing `{path}`: {message}")] + ProcessFailed { + path: PathBuf, + exit_code: i32, + message: String, + }, + #[error("sandboxed fs helper returned invalid output for `{path}`: {message}")] + InvalidResponse { path: PathBuf, message: String }, + #[error("sandboxed fs helper could not access `{path}`: {error}")] + Operation { path: PathBuf, error: FsError }, +} + +impl SandboxedFsError { + pub(crate) fn operation_error_kind(&self) -> Option<&FsErrorKind> { + match self { + Self::Operation { error, .. } => Some(&error.kind), + _ => None, + } + } + + pub(crate) fn operation_error_message(&self) -> Option<&str> { + match self { + Self::Operation { error, .. } => Some(error.message.as_str()), + _ => None, + } + } + + #[allow(dead_code)] + pub(crate) fn to_io_error(&self) -> std::io::Error { + match self { + Self::Operation { error, .. } => error.to_io_error(), + Self::TimedOut { .. } => { + std::io::Error::new(std::io::ErrorKind::TimedOut, self.to_string()) + } + Self::ResolveExe { .. } | Self::ProcessFailed { .. } | Self::InvalidResponse { .. } => { + std::io::Error::other(self.to_string()) + } + } + } +} + +pub(crate) async fn read_bytes( + session: &Arc, + turn: &Arc, + path: &Path, +) -> Result, SandboxedFsError> { + let path_buf = path.to_path_buf(); + let payload = run_request(session, turn, "read_bytes", &path_buf).await?; + let FsPayload::Bytes { base64 } = payload else { + return Err(SandboxedFsError::InvalidResponse { + path: path_buf, + message: "expected bytes payload".to_string(), + }); + }; + BASE64_STANDARD + .decode(base64) + .map_err(|error| SandboxedFsError::InvalidResponse { + path: path.to_path_buf(), + message: format!("failed to decode base64 payload: {error}"), + }) +} + +#[allow(dead_code)] +pub(crate) async fn read_text( + session: &Arc, + turn: &Arc, + path: &Path, +) -> Result { + let path_buf = path.to_path_buf(); + let payload = run_request(session, turn, "read_text", &path_buf).await?; + let FsPayload::Text { text } = payload else { + return Err(SandboxedFsError::InvalidResponse { + path: path_buf, + message: "expected text payload".to_string(), + }); + }; + Ok(text) +} + +async fn run_request( + session: &Arc, + turn: &Arc, + operation: &str, + path: &Path, +) -> Result { + let exe = resolve_codex_exe(&turn.config.codex_home)?; + let additional_permissions = effective_granted_permissions(session).await; + let sandbox_manager = crate::sandboxing::SandboxManager::new(); + let attempt = SandboxAttempt { + sandbox: sandbox_manager.select_initial( + &turn.file_system_sandbox_policy, + turn.network_sandbox_policy, + SandboxablePreference::Auto, + turn.windows_sandbox_level, + /*has_managed_network_requirements*/ false, + ), + policy: &turn.sandbox_policy, + file_system_policy: &turn.file_system_sandbox_policy, + network_policy: turn.network_sandbox_policy, + enforce_managed_network: false, + manager: &sandbox_manager, + sandbox_cwd: &turn.cwd, + codex_linux_sandbox_exe: turn.codex_linux_sandbox_exe.as_ref(), + use_legacy_landlock: turn.features.use_legacy_landlock(), + windows_sandbox_level: turn.windows_sandbox_level, + windows_sandbox_private_desktop: turn.config.permissions.windows_sandbox_private_desktop, + }; + let exec_request = attempt + .env_for( + CommandSpec { + program: exe.to_string_lossy().to_string(), + args: vec![ + CODEX_CORE_FS_OPS_ARG1.to_string(), + operation.to_string(), + path.to_string_lossy().to_string(), + ], + cwd: turn.cwd.clone(), + env: HashMap::new(), + expiration: ExecExpiration::Timeout(SANDBOXED_FS_TIMEOUT), + sandbox_permissions: SandboxPermissions::UseDefault, + additional_permissions, + justification: None, + }, + None, + ) + .map_err(|error| SandboxedFsError::ProcessFailed { + path: path.to_path_buf(), + exit_code: -1, + message: error.to_string(), + })?; + let output = + execute_env(exec_request, None) + .await + .map_err(|error| SandboxedFsError::ProcessFailed { + path: path.to_path_buf(), + exit_code: -1, + message: error.to_string(), + })?; + + if output.timed_out { + return Err(SandboxedFsError::TimedOut { + path: path.to_path_buf(), + }); + } + if output.exit_code != 0 { + let stderr = output.stderr.text.trim(); + let stdout = output.stdout.text.trim(); + let message = if !stderr.is_empty() { + stderr.to_string() + } else if !stdout.is_empty() { + stdout.to_string() + } else { + "no error details emitted".to_string() + }; + return Err(SandboxedFsError::ProcessFailed { + path: path.to_path_buf(), + exit_code: output.exit_code, + message, + }); + } + + let response: FsResponse = + serde_json::from_str(output.stdout.text.trim()).map_err(|error| { + SandboxedFsError::InvalidResponse { + path: path.to_path_buf(), + message: format!("failed to parse helper response: {error}"), + } + })?; + + match response { + FsResponse::Success { payload } => Ok(payload), + FsResponse::Error { error } => Err(SandboxedFsError::Operation { + path: path.to_path_buf(), + error, + }), + } +} + +async fn effective_granted_permissions(session: &Session) -> Option { + let granted_session_permissions = session.granted_session_permissions().await; + let granted_turn_permissions = session.granted_turn_permissions().await; + merge_permission_profiles( + granted_session_permissions.as_ref(), + granted_turn_permissions.as_ref(), + ) +} + +fn resolve_codex_exe(codex_home: &Path) -> Result { + #[cfg(target_os = "windows")] + { + Ok(codex_windows_sandbox::resolve_current_exe_for_launch( + codex_home, + "codex.exe", + )) + } + #[cfg(not(target_os = "windows"))] + { + let _ = codex_home; + std::env::current_exe().map_err(|error| SandboxedFsError::ResolveExe { + message: error.to_string(), + }) + } +} diff --git a/codex-rs/core/src/tools/handlers/view_image.rs b/codex-rs/core/src/tools/handlers/view_image.rs index 5757aeb4e1..b6613b8aad 100644 --- a/codex-rs/core/src/tools/handlers/view_image.rs +++ b/codex-rs/core/src/tools/handlers/view_image.rs @@ -2,16 +2,16 @@ use async_trait::async_trait; use codex_protocol::models::ContentItem; use codex_protocol::models::FunctionCallOutputContentItem; use codex_protocol::models::ImageDetail; -use codex_protocol::models::local_image_content_items_with_label_number; +use codex_protocol::models::local_image_content_items_from_bytes_with_label_number; use codex_protocol::openai_models::InputModality; use codex_utils_image::PromptImageMode; use serde::Deserialize; -use tokio::fs; use crate::function_tool::FunctionCallError; use crate::original_image_detail::can_request_original_image_detail; use crate::protocol::EventMsg; use crate::protocol::ViewImageToolCallEvent; +use crate::sandboxed_fs; use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; @@ -88,20 +88,6 @@ impl ToolHandler for ViewImageHandler { }; let abs_path = turn.resolve_path(Some(args.path)); - - let metadata = fs::metadata(&abs_path).await.map_err(|error| { - FunctionCallError::RespondToModel(format!( - "unable to locate image at `{}`: {error}", - abs_path.display() - )) - })?; - - if !metadata.is_file() { - return Err(FunctionCallError::RespondToModel(format!( - "image path `{}` is not a file", - abs_path.display() - ))); - } let event_path = abs_path.clone(); let can_request_original_detail = @@ -114,9 +100,17 @@ impl ToolHandler for ViewImageHandler { PromptImageMode::ResizeToFit }; let image_detail = use_original_detail.then_some(ImageDetail::Original); + let image_bytes = sandboxed_fs::read_bytes(&session, &turn, &abs_path) + .await + .map_err(|error| { + FunctionCallError::RespondToModel(render_view_image_read_error(&abs_path, &error)) + })?; - let content = local_image_content_items_with_label_number( - &abs_path, /*label_number*/ None, image_mode, + let content = local_image_content_items_from_bytes_with_label_number( + &abs_path, + image_bytes, + /*label_number*/ None, + image_mode, ) .into_iter() .map(|item| match item { @@ -142,3 +136,30 @@ impl ToolHandler for ViewImageHandler { Ok(FunctionToolOutput::from_content(content, Some(true))) } } + +fn render_view_image_read_error( + path: &std::path::Path, + error: &sandboxed_fs::SandboxedFsError, +) -> String { + let operation_message = error + .operation_error_message() + .map(str::to_owned) + .unwrap_or_else(|| error.to_string()); + match error.operation_error_kind() { + Some(codex_fs_ops::FsErrorKind::IsADirectory) => { + format!("image path `{}` is not a file", path.display()) + } + Some(codex_fs_ops::FsErrorKind::NotFound) => { + format!( + "unable to locate image at `{}`: {operation_message}", + path.display() + ) + } + Some(_) | None => { + format!( + "unable to read image at `{}`: {operation_message}", + path.display() + ) + } + } +} diff --git a/codex-rs/core/tests/suite/view_image.rs b/codex-rs/core/tests/suite/view_image.rs index a3e341f192..79ad594c3b 100644 --- a/codex-rs/core/tests/suite/view_image.rs +++ b/codex-rs/core/tests/suite/view_image.rs @@ -3,6 +3,7 @@ use base64::Engine; use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use codex_core::CodexAuth; +use codex_core::config::Constrained; use codex_core::features::Feature; use codex_protocol::config_types::ReasoningSummary; use codex_protocol::openai_models::ConfigShellToolType; @@ -13,9 +14,15 @@ use codex_protocol::openai_models::ModelsResponse; use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::openai_models::ReasoningEffortPreset; use codex_protocol::openai_models::TruncationPolicyConfig; +use codex_protocol::permissions::FileSystemAccessMode; +use codex_protocol::permissions::FileSystemPath; +use codex_protocol::permissions::FileSystemSandboxEntry; +use codex_protocol::permissions::FileSystemSandboxPolicy; +use codex_protocol::permissions::FileSystemSpecialPath; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; +use codex_protocol::protocol::ReadOnlyAccess; use codex_protocol::protocol::SandboxPolicy; use codex_protocol::user_input::UserInput; use core_test_support::responses; @@ -1244,6 +1251,110 @@ async fn view_image_tool_errors_when_file_missing() -> anyhow::Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn view_image_tool_respects_filesystem_sandbox() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let sandbox_policy_for_config = SandboxPolicy::ReadOnly { + access: ReadOnlyAccess::Restricted { + include_platform_defaults: true, + readable_roots: Vec::new(), + }, + network_access: false, + }; + let mut builder = test_codex().with_config({ + let sandbox_policy_for_config = sandbox_policy_for_config.clone(); + move |config| { + config.permissions.sandbox_policy = + Constrained::allow_any(sandbox_policy_for_config); + config.permissions.file_system_sandbox_policy = + FileSystemSandboxPolicy::restricted(vec![ + FileSystemSandboxEntry { + path: FileSystemPath::Special { + value: FileSystemSpecialPath::Minimal, + }, + access: FileSystemAccessMode::Read, + }, + FileSystemSandboxEntry { + path: FileSystemPath::Special { + value: FileSystemSpecialPath::CurrentWorkingDirectory, + }, + access: FileSystemAccessMode::Read, + }, + ]); + } + }); + let TestCodex { + codex, + config, + cwd, + session_configured, + .. + } = builder.build(&server).await?; + + let outside_dir = tempfile::tempdir()?; + let abs_path = outside_dir.path().join("blocked.png"); + let image = ImageBuffer::from_pixel(256, 128, Rgba([10u8, 20, 30, 255])); + image.save(&abs_path)?; + + let call_id = "view-image-sandbox-denied"; + let arguments = serde_json::json!({ "path": abs_path }).to_string(); + + let first_response = sse(vec![ + ev_response_created("resp-1"), + ev_function_call(call_id, "view_image", &arguments), + ev_completed("resp-1"), + ]); + responses::mount_sse_once(&server, first_response).await; + + let second_response = sse(vec![ + ev_assistant_message("msg-1", "done"), + ev_completed("resp-2"), + ]); + let mock = responses::mount_sse_once(&server, second_response).await; + + let session_model = session_configured.model.clone(); + + codex + .submit(Op::UserTurn { + items: vec![UserInput::Text { + text: "please attach the outside image".into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + cwd: cwd.path().to_path_buf(), + approval_policy: AskForApproval::Never, + sandbox_policy: config.permissions.sandbox_policy.get().clone(), + model: session_model, + effort: None, + summary: None, + service_tier: None, + collaboration_mode: None, + personality: None, + }) + .await?; + + wait_for_event(&codex, |event| matches!(event, EventMsg::TurnComplete(_))).await; + + let request = mock.single_request(); + assert!( + request.inputs_of_type("input_image").is_empty(), + "sandbox-denied image should not produce an input_image message" + ); + let output_text = request + .function_call_output_content_and_success(call_id) + .and_then(|(content, _)| content) + .expect("output text present"); + let expected_prefix = format!("unable to read image at `{}`:", abs_path.display()); + assert!( + output_text.starts_with(&expected_prefix), + "expected sandbox denial prefix `{expected_prefix}` but got `{output_text}`" + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn view_image_tool_returns_unsupported_message_for_text_only_model() -> anyhow::Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/fs-ops/BUILD.bazel b/codex-rs/fs-ops/BUILD.bazel new file mode 100644 index 0000000000..87fe4d5517 --- /dev/null +++ b/codex-rs/fs-ops/BUILD.bazel @@ -0,0 +1,6 @@ +load("//:defs.bzl", "codex_rust_crate") + +codex_rust_crate( + name = "fs-ops", + crate_name = "codex_fs_ops", +) diff --git a/codex-rs/fs-ops/Cargo.toml b/codex-rs/fs-ops/Cargo.toml new file mode 100644 index 0000000000..2b7495f0d2 --- /dev/null +++ b/codex-rs/fs-ops/Cargo.toml @@ -0,0 +1,22 @@ +[package] +name = "codex-fs-ops" +version.workspace = true +edition.workspace = true +license.workspace = true + +[lib] +name = "codex_fs_ops" +path = "src/lib.rs" + +[lints] +workspace = true + +[dependencies] +anyhow = { workspace = true } +base64 = { workspace = true } +serde = { workspace = true, features = ["derive"] } +serde_json = { workspace = true } + +[dev-dependencies] +pretty_assertions = { workspace = true } +tempfile = { workspace = true } diff --git a/codex-rs/fs-ops/src/lib.rs b/codex-rs/fs-ops/src/lib.rs new file mode 100644 index 0000000000..32d5c2cf6f --- /dev/null +++ b/codex-rs/fs-ops/src/lib.rs @@ -0,0 +1,229 @@ +use anyhow::Context; +use anyhow::Result; +use base64::Engine; +use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; +use serde::Deserialize; +use serde::Serialize; +use std::ffi::OsString; +use std::io::ErrorKind; +use std::io::Write; +use std::path::PathBuf; + +/// Special argv[1] flag used when the Codex executable self-invokes to run the +/// internal sandbox-backed filesystem helper path. +pub const CODEX_CORE_FS_OPS_ARG1: &str = "--codex-run-as-fs-ops"; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum FsCommand { + ReadBytes { path: PathBuf }, + ReadText { path: PathBuf }, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum FsErrorKind { + NotFound, + PermissionDenied, + IsADirectory, + InvalidData, + Other, +} + +impl From for FsErrorKind { + fn from(value: ErrorKind) -> Self { + match value { + ErrorKind::NotFound => Self::NotFound, + ErrorKind::PermissionDenied => Self::PermissionDenied, + ErrorKind::IsADirectory => Self::IsADirectory, + ErrorKind::InvalidData => Self::InvalidData, + _ => Self::Other, + } + } +} + +impl FsErrorKind { + pub fn to_io_error_kind(&self) -> ErrorKind { + match self { + Self::NotFound => ErrorKind::NotFound, + Self::PermissionDenied => ErrorKind::PermissionDenied, + Self::IsADirectory => ErrorKind::IsADirectory, + Self::InvalidData => ErrorKind::InvalidData, + Self::Other => ErrorKind::Other, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct FsError { + pub kind: FsErrorKind, + pub message: String, + pub raw_os_error: Option, +} + +impl std::fmt::Display for FsError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(&self.message) + } +} + +impl FsError { + pub fn to_io_error(&self) -> std::io::Error { + if let Some(raw_os_error) = self.raw_os_error { + std::io::Error::from_raw_os_error(raw_os_error) + } else { + std::io::Error::new(self.kind.to_io_error_kind(), self.message.clone()) + } + } +} + +impl From for FsError { + fn from(error: std::io::Error) -> Self { + Self { + kind: error.kind().into(), + message: error.to_string(), + raw_os_error: error.raw_os_error(), + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(tag = "kind", rename_all = "snake_case")] +pub enum FsPayload { + Bytes { base64: String }, + Text { text: String }, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(tag = "status", rename_all = "snake_case")] +pub enum FsResponse { + Success { payload: FsPayload }, + Error { error: FsError }, +} + +pub fn parse_command_from_args( + mut args: impl Iterator, +) -> Result { + let Some(operation) = args.next() else { + return Err("missing operation".to_string()); + }; + let Some(operation) = operation.to_str() else { + return Err("operation must be valid UTF-8".to_string()); + }; + let Some(path) = args.next() else { + return Err(format!("missing path for operation `{operation}`")); + }; + if args.next().is_some() { + return Err(format!( + "unexpected extra arguments for operation `{operation}`" + )); + } + + let path = PathBuf::from(path); + match operation { + "read_bytes" => Ok(FsCommand::ReadBytes { path }), + "read_text" => Ok(FsCommand::ReadText { path }), + _ => Err(format!( + "unsupported filesystem operation `{operation}`; expected one of `read_bytes`, `read_text`" + )), + } +} + +pub fn execute(command: FsCommand) -> FsResponse { + match command { + FsCommand::ReadBytes { path } => match std::fs::read(&path) { + Ok(bytes) => FsResponse::Success { + payload: FsPayload::Bytes { + base64: BASE64_STANDARD.encode(bytes), + }, + }, + Err(error) => FsResponse::Error { + error: error.into(), + }, + }, + FsCommand::ReadText { path } => match std::fs::read_to_string(&path) { + Ok(text) => FsResponse::Success { + payload: FsPayload::Text { text }, + }, + Err(error) => FsResponse::Error { + error: error.into(), + }, + }, + } +} + +pub fn write_response(stdout: &mut impl Write, response: &FsResponse) -> Result<()> { + serde_json::to_writer(&mut *stdout, response).context("failed to serialize fs response")?; + writeln!(stdout).context("failed to terminate fs response with newline")?; + Ok(()) +} + +pub fn run_from_args(args: impl Iterator) -> Result<()> { + let command = parse_command_from_args(args).map_err(anyhow::Error::msg)?; + let response = execute(command); + let mut stdout = std::io::stdout().lock(); + write_response(&mut stdout, &response) +} + +#[cfg(test)] +mod tests { + use super::FsCommand; + use super::FsErrorKind; + use super::FsPayload; + use super::FsResponse; + use super::execute; + use super::parse_command_from_args; + use pretty_assertions::assert_eq; + use tempfile::tempdir; + + #[test] + fn parse_read_bytes_command() { + let command = parse_command_from_args( + ["read_bytes", "/tmp/example.png"] + .into_iter() + .map(Into::into), + ) + .expect("command should parse"); + + assert_eq!( + command, + FsCommand::ReadBytes { + path: "/tmp/example.png".into(), + } + ); + } + + #[test] + fn read_text_returns_text_payload() { + let tempdir = tempdir().expect("tempdir"); + let path = tempdir.path().join("note.txt"); + std::fs::write(&path, "hello").expect("write test file"); + + let response = execute(FsCommand::ReadText { path }); + + assert_eq!( + response, + FsResponse::Success { + payload: FsPayload::Text { + text: "hello".to_string(), + }, + } + ); + } + + #[test] + fn read_bytes_reports_directory_error() { + let tempdir = tempdir().expect("tempdir"); + + let response = execute(FsCommand::ReadBytes { + path: tempdir.path().to_path_buf(), + }); + + let FsResponse::Error { error } = response else { + panic!("expected error response"); + }; + #[cfg(target_os = "windows")] + assert_eq!(error.kind, FsErrorKind::PermissionDenied); + #[cfg(not(target_os = "windows"))] + assert_eq!(error.kind, FsErrorKind::IsADirectory); + } +} diff --git a/codex-rs/protocol/src/models.rs b/codex-rs/protocol/src/models.rs index 1e5a7445e3..f60f431132 100644 --- a/codex-rs/protocol/src/models.rs +++ b/codex-rs/protocol/src/models.rs @@ -3,6 +3,7 @@ use std::path::Path; use codex_utils_image::PromptImageMode; use codex_utils_image::load_for_prompt; +use codex_utils_image::load_for_prompt_from_bytes; use serde::Deserialize; use serde::Deserializer; use serde::Serialize; @@ -947,7 +948,28 @@ pub fn local_image_content_items_with_label_number( label_number: Option, mode: PromptImageMode, ) -> Vec { - match load_for_prompt(path, mode) { + local_image_content_items_for_load_result(path, label_number, load_for_prompt(path, mode)) +} + +pub fn local_image_content_items_from_bytes_with_label_number( + path: &std::path::Path, + bytes: Vec, + label_number: Option, + mode: PromptImageMode, +) -> Vec { + local_image_content_items_for_load_result( + path, + label_number, + load_for_prompt_from_bytes(path, bytes, mode), + ) +} + +fn local_image_content_items_for_load_result( + path: &std::path::Path, + label_number: Option, + image_result: Result, +) -> Vec { + match image_result { Ok(image) => { let mut items = Vec::with_capacity(3); if let Some(label_number) = label_number { diff --git a/codex-rs/utils/image/src/lib.rs b/codex-rs/utils/image/src/lib.rs index d0aba2f608..89547fa3a7 100644 --- a/codex-rs/utils/image/src/lib.rs +++ b/codex-rs/utils/image/src/lib.rs @@ -62,9 +62,16 @@ pub fn load_for_prompt( mode: PromptImageMode, ) -> Result { let path_buf = path.to_path_buf(); - let file_bytes = read_file_bytes(path, &path_buf)?; + load_for_prompt_from_bytes(path, file_bytes, mode) +} +pub fn load_for_prompt_from_bytes( + path: &Path, + file_bytes: Vec, + mode: PromptImageMode, +) -> Result { + let path_buf = path.to_path_buf(); let key = ImageCacheKey { digest: sha1_digest(&file_bytes), mode,