From df000da917952024edd14cb624af8d1e1740a2b5 Mon Sep 17 00:00:00 2001 From: Rasmus Rygaard Date: Wed, 4 Feb 2026 06:01:47 -0800 Subject: [PATCH 1/8] Add a codex.rate_limits event for websockets (#10324) When communicating over websockets, we can't rely on headers to deliver rate limit information. This PR adds a `codex.rate_limits` event that the server can pass to the client to inform them about rate limit usage. The client parses this data the same way we parse rate limit headers in HTTP mode. This PR also wires up the etag and reasoning headers for websockets --- .../src/endpoint/responses_websocket.rs | 29 ++++++- codex-rs/codex-api/src/rate_limits.rs | 66 ++++++++++++++ codex-rs/codex-api/src/sse/responses.rs | 6 ++ .../core/tests/suite/client_websockets.rs | 86 +++++++++++++++++++ 4 files changed, 183 insertions(+), 4 deletions(-) diff --git a/codex-rs/codex-api/src/endpoint/responses_websocket.rs b/codex-rs/codex-api/src/endpoint/responses_websocket.rs index cac686dd02..20c6067b17 100644 --- a/codex-rs/codex-api/src/endpoint/responses_websocket.rs +++ b/codex-rs/codex-api/src/endpoint/responses_websocket.rs @@ -5,6 +5,7 @@ use crate::common::ResponseStream; use crate::common::ResponsesWsRequest; use crate::error::ApiError; use crate::provider::Provider; +use crate::rate_limits::parse_rate_limit_event; use crate::sse::responses::ResponsesStreamEvent; use crate::sse::responses::process_responses_event; use crate::telemetry::WebsocketTelemetry; @@ -33,6 +34,7 @@ use url::Url; type WsStream = WebSocketStream>; const X_CODEX_TURN_STATE_HEADER: &str = "x-codex-turn-state"; +const X_MODELS_ETAG_HEADER: &str = "x-models-etag"; const X_REASONING_INCLUDED_HEADER: &str = "x-reasoning-included"; pub struct ResponsesWebsocketConnection { @@ -40,6 +42,7 @@ pub struct ResponsesWebsocketConnection { // TODO (pakrym): is this the right place for timeout? idle_timeout: Duration, server_reasoning_included: bool, + models_etag: Option, telemetry: Option>, } @@ -48,12 +51,14 @@ impl ResponsesWebsocketConnection { stream: WsStream, idle_timeout: Duration, server_reasoning_included: bool, + models_etag: Option, telemetry: Option>, ) -> Self { Self { stream: Arc::new(Mutex::new(Some(stream))), idle_timeout, server_reasoning_included, + models_etag, telemetry, } } @@ -71,12 +76,16 @@ impl ResponsesWebsocketConnection { let stream = Arc::clone(&self.stream); let idle_timeout = self.idle_timeout; let server_reasoning_included = self.server_reasoning_included; + let models_etag = self.models_etag.clone(); let telemetry = self.telemetry.clone(); let request_body = serde_json::to_value(&request).map_err(|err| { ApiError::Stream(format!("failed to encode websocket request: {err}")) })?; tokio::spawn(async move { + if let Some(etag) = models_etag { + let _ = tx_event.send(Ok(ResponseEvent::ModelsEtag(etag))).await; + } if server_reasoning_included { let _ = tx_event .send(Ok(ResponseEvent::ServerReasoningIncluded(true))) @@ -136,12 +145,13 @@ impl ResponsesWebsocketClient { headers.extend(extra_headers); add_auth_headers_to_header_map(&self.auth, &mut headers); - let (stream, server_reasoning_included) = - connect_websocket(ws_url, headers, turn_state).await?; + let (stream, server_reasoning_included, models_etag) = + connect_websocket(ws_url, headers, turn_state.clone()).await?; Ok(ResponsesWebsocketConnection::new( stream, self.provider.stream_idle_timeout, server_reasoning_included, + models_etag, telemetry, )) } @@ -151,7 +161,7 @@ async fn connect_websocket( url: Url, headers: HeaderMap, turn_state: Option>>, -) -> Result<(WsStream, bool), ApiError> { +) -> Result<(WsStream, bool, Option), ApiError> { info!("connecting to websocket: {url}"); let mut request = url @@ -177,6 +187,11 @@ async fn connect_websocket( }; let reasoning_included = response.headers().contains_key(X_REASONING_INCLUDED_HEADER); + let models_etag = response + .headers() + .get(X_MODELS_ETAG_HEADER) + .and_then(|value| value.to_str().ok()) + .map(ToString::to_string); if let Some(turn_state) = turn_state && let Some(header_value) = response .headers() @@ -185,7 +200,7 @@ async fn connect_websocket( { let _ = turn_state.set(header_value.to_string()); } - Ok((stream, reasoning_included)) + Ok((stream, reasoning_included, models_etag)) } fn map_ws_error(err: WsError, url: &Url) -> ApiError { @@ -273,6 +288,12 @@ async fn run_websocket_response_stream( continue; } }; + if event.kind() == "codex.rate_limits" { + if let Some(snapshot) = parse_rate_limit_event(&text) { + let _ = tx_event.send(Ok(ResponseEvent::RateLimits(snapshot))).await; + } + continue; + } match process_responses_event(event) { Ok(Some(event)) => { let is_completed = matches!(event, ResponseEvent::Completed { .. }); diff --git a/codex-rs/codex-api/src/rate_limits.rs b/codex-rs/codex-api/src/rate_limits.rs index c29aab21f8..047c5934bd 100644 --- a/codex-rs/codex-api/src/rate_limits.rs +++ b/codex-rs/codex-api/src/rate_limits.rs @@ -1,7 +1,9 @@ +use codex_protocol::account::PlanType; use codex_protocol::protocol::CreditsSnapshot; use codex_protocol::protocol::RateLimitSnapshot; use codex_protocol::protocol::RateLimitWindow; use http::HeaderMap; +use serde::Deserialize; use std::fmt::Display; #[derive(Debug)] @@ -41,6 +43,70 @@ pub fn parse_rate_limit(headers: &HeaderMap) -> Option { }) } +#[derive(Debug, Deserialize)] +struct RateLimitEventWindow { + used_percent: f64, + window_minutes: Option, + reset_at: Option, +} + +#[derive(Debug, Deserialize)] +struct RateLimitEventDetails { + primary: Option, + secondary: Option, +} + +#[derive(Debug, Deserialize)] +struct RateLimitEventCredits { + has_credits: bool, + unlimited: bool, + balance: Option, +} + +#[derive(Debug, Deserialize)] +struct RateLimitEvent { + #[serde(rename = "type")] + kind: String, + plan_type: Option, + rate_limits: Option, + credits: Option, +} + +pub fn parse_rate_limit_event(payload: &str) -> Option { + let event: RateLimitEvent = serde_json::from_str(payload).ok()?; + if event.kind != "codex.rate_limits" { + return None; + } + let (primary, secondary) = if let Some(details) = event.rate_limits.as_ref() { + ( + map_event_window(details.primary.as_ref()), + map_event_window(details.secondary.as_ref()), + ) + } else { + (None, None) + }; + let credits = event.credits.map(|credits| CreditsSnapshot { + has_credits: credits.has_credits, + unlimited: credits.unlimited, + balance: credits.balance, + }); + Some(RateLimitSnapshot { + primary, + secondary, + credits, + plan_type: event.plan_type, + }) +} + +fn map_event_window(window: Option<&RateLimitEventWindow>) -> Option { + let window = window?; + Some(RateLimitWindow { + used_percent: window.used_percent, + window_minutes: window.window_minutes, + resets_at: window.reset_at, + }) +} + /// Parses the bespoke Codex rate-limit headers into a `RateLimitSnapshot`. pub fn parse_promo_message(headers: &HeaderMap) -> Option { parse_header_str(headers, "x-codex-promo-message") diff --git a/codex-rs/codex-api/src/sse/responses.rs b/codex-rs/codex-api/src/sse/responses.rs index 2c91104727..95bda35fb2 100644 --- a/codex-rs/codex-api/src/sse/responses.rs +++ b/codex-rs/codex-api/src/sse/responses.rs @@ -165,6 +165,12 @@ pub struct ResponsesStreamEvent { content_index: Option, } +impl ResponsesStreamEvent { + pub fn kind(&self) -> &str { + &self.kind + } +} + #[derive(Debug)] pub enum ResponsesEventError { Api(ApiError), diff --git a/codex-rs/core/tests/suite/client_websockets.rs b/codex-rs/core/tests/suite/client_websockets.rs index 9f66283442..15005caa48 100644 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/codex-rs/core/tests/suite/client_websockets.rs @@ -18,6 +18,7 @@ use codex_otel::OtelManager; use codex_otel::metrics::MetricsClient; use codex_otel::metrics::MetricsConfig; use codex_protocol::ThreadId; +use codex_protocol::account::PlanType; use codex_protocol::config_types::ReasoningSummary; use core_test_support::load_default_config_for_test; use core_test_support::responses::WebSocketConnectionConfig; @@ -30,6 +31,7 @@ use core_test_support::skip_if_no_network; use futures::StreamExt; use opentelemetry_sdk::metrics::InMemoryMetricExporter; use pretty_assertions::assert_eq; +use serde_json::json; use std::sync::Arc; use std::time::Duration; use tempfile::TempDir; @@ -203,6 +205,90 @@ async fn responses_websocket_emits_reasoning_included_event() { server.shutdown().await; } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn responses_websocket_emits_rate_limit_events() { + skip_if_no_network!(); + + let rate_limit_event = json!({ + "type": "codex.rate_limits", + "plan_type": "plus", + "rate_limits": { + "allowed": true, + "limit_reached": false, + "primary": { + "used_percent": 42, + "window_minutes": 60, + "reset_at": 1700000000 + }, + "secondary": null + }, + "code_review_rate_limits": null, + "credits": { + "has_credits": true, + "unlimited": false, + "balance": "123" + }, + "promo": null + }); + + let server = start_websocket_server_with_headers(vec![WebSocketConnectionConfig { + requests: vec![vec![ + rate_limit_event, + ev_response_created("resp-1"), + ev_completed("resp-1"), + ]], + response_headers: vec![ + ("X-Models-Etag".to_string(), "etag-123".to_string()), + ("X-Reasoning-Included".to_string(), "true".to_string()), + ], + }]) + .await; + + let harness = websocket_harness(&server).await; + let mut session = harness.client.new_session(None); + let prompt = prompt_with_input(vec![message_item("hello")]); + + let mut stream = session + .stream(&prompt) + .await + .expect("websocket stream failed"); + + let mut saw_rate_limits = None; + let mut saw_models_etag = None; + let mut saw_reasoning_included = false; + + while let Some(event) = stream.next().await { + match event.expect("event") { + ResponseEvent::RateLimits(snapshot) => { + saw_rate_limits = Some(snapshot); + } + ResponseEvent::ModelsEtag(etag) => { + saw_models_etag = Some(etag); + } + ResponseEvent::ServerReasoningIncluded(true) => { + saw_reasoning_included = true; + } + ResponseEvent::Completed { .. } => break, + _ => {} + } + } + + let rate_limits = saw_rate_limits.expect("missing rate limits"); + let primary = rate_limits.primary.expect("missing primary window"); + assert_eq!(primary.used_percent, 42.0); + assert_eq!(primary.window_minutes, Some(60)); + assert_eq!(primary.resets_at, Some(1_700_000_000)); + assert_eq!(rate_limits.plan_type, Some(PlanType::Plus)); + let credits = rate_limits.credits.expect("missing credits"); + assert!(credits.has_credits); + assert!(!credits.unlimited); + assert_eq!(credits.balance.as_deref(), Some("123")); + assert_eq!(saw_models_etag.as_deref(), Some("etag-123")); + assert!(saw_reasoning_included); + + server.shutdown().await; +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn responses_websocket_appends_on_prefix() { skip_if_no_network!(); From 583e5d4f413dc10457e685d26672529ba7d2f70e Mon Sep 17 00:00:00 2001 From: jif-oai Date: Wed, 4 Feb 2026 14:31:12 +0000 Subject: [PATCH 2/8] Migrate state DB path helpers to versioned filename (#10623) Summary - add versioned state sqlite filename helpers and re-export them from the state crate - remove legacy state files when initializing the runtime and update consumers/tests to use the new helpers - tweak logs client description and database resolution to match the new path --- codex-rs/core/src/state_db.rs | 7 +- codex-rs/core/tests/suite/sqlite_state.rs | 7 +- codex-rs/state/src/bin/logs_client.rs | 5 +- codex-rs/state/src/lib.rs | 3 + codex-rs/state/src/runtime.rs | 171 +++++++++++++++++++++- 5 files changed, 180 insertions(+), 13 deletions(-) diff --git a/codex-rs/core/src/state_db.rs b/codex-rs/core/src/state_db.rs index 7a0894a327..1cfb8f1566 100644 --- a/codex-rs/core/src/state_db.rs +++ b/codex-rs/core/src/state_db.rs @@ -14,7 +14,6 @@ use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::SessionSource; use codex_state::DB_METRIC_COMPARE_ERROR; pub use codex_state::LogEntry; -use codex_state::STATE_DB_FILENAME; use codex_state::ThreadMetadataBuilder; use serde_json::Value; use std::path::Path; @@ -32,7 +31,7 @@ pub(crate) async fn init_if_enabled( config: &Config, otel: Option<&OtelManager>, ) -> Option { - let state_path = config.codex_home.join(STATE_DB_FILENAME); + let state_path = codex_state::state_db_path(config.codex_home.as_path()); if !config.features.enabled(Feature::Sqlite) { return None; } @@ -74,7 +73,7 @@ pub(crate) async fn init_if_enabled( /// Get the DB if the feature is enabled and the DB exists. pub async fn get_state_db(config: &Config, otel: Option<&OtelManager>) -> Option { - let state_path = config.codex_home.join(STATE_DB_FILENAME); + let state_path = codex_state::state_db_path(config.codex_home.as_path()); if !config.features.enabled(Feature::Sqlite) || !tokio::fs::try_exists(&state_path).await.unwrap_or(false) { @@ -93,7 +92,7 @@ pub async fn get_state_db(config: &Config, otel: Option<&OtelManager>) -> Option /// /// This is used for parity checks during the SQLite migration phase. pub async fn open_if_present(codex_home: &Path, default_provider: &str) -> Option { - let db_path = codex_home.join(STATE_DB_FILENAME); + let db_path = codex_state::state_db_path(codex_home); if !tokio::fs::try_exists(&db_path).await.unwrap_or(false) { return None; } diff --git a/codex-rs/core/tests/suite/sqlite_state.rs b/codex-rs/core/tests/suite/sqlite_state.rs index 218da34825..ac71c14d36 100644 --- a/codex-rs/core/tests/suite/sqlite_state.rs +++ b/codex-rs/core/tests/suite/sqlite_state.rs @@ -9,7 +9,6 @@ use codex_protocol::protocol::SessionMeta; use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::UserMessageEvent; -use codex_state::STATE_DB_FILENAME; use core_test_support::load_sse_fixture_with_id; use core_test_support::responses; use core_test_support::responses::ev_completed; @@ -39,7 +38,7 @@ async fn new_thread_is_recorded_in_state_db() -> Result<()> { let thread_id = test.session_configured.session_id; let rollout_path = test.codex.rollout_path().expect("rollout path"); - let db_path = test.config.codex_home.join(STATE_DB_FILENAME); + let db_path = codex_state::state_db_path(test.config.codex_home.as_path()); for _ in 0..100 { if tokio::fs::try_exists(&db_path).await.unwrap_or(false) { @@ -149,7 +148,7 @@ async fn backfill_scans_existing_rollouts() -> Result<()> { let test = builder.build(&server).await?; - let db_path = test.config.codex_home.join(STATE_DB_FILENAME); + let db_path = codex_state::state_db_path(test.config.codex_home.as_path()); let rollout_path = test.config.codex_home.join(&rollout_rel_path); let default_provider = test.config.model_provider_id.clone(); @@ -205,7 +204,7 @@ async fn user_messages_persist_in_state_db() -> Result<()> { }); let test = builder.build(&server).await?; - let db_path = test.config.codex_home.join(STATE_DB_FILENAME); + let db_path = codex_state::state_db_path(test.config.codex_home.as_path()); for _ in 0..100 { if tokio::fs::try_exists(&db_path).await.unwrap_or(false) { break; diff --git a/codex-rs/state/src/bin/logs_client.rs b/codex-rs/state/src/bin/logs_client.rs index d1329a1a89..796a968cb8 100644 --- a/codex-rs/state/src/bin/logs_client.rs +++ b/codex-rs/state/src/bin/logs_client.rs @@ -6,14 +6,13 @@ use chrono::DateTime; use clap::Parser; use codex_state::LogQuery; use codex_state::LogRow; -use codex_state::STATE_DB_FILENAME; use codex_state::StateRuntime; use dirs::home_dir; use owo_colors::OwoColorize; #[derive(Debug, Parser)] #[command(name = "codex-state-logs")] -#[command(about = "Tail Codex logs from state.sqlite with simple filters")] +#[command(about = "Tail Codex logs from the state SQLite DB with simple filters")] struct Args { /// Path to CODEX_HOME. Defaults to $CODEX_HOME or ~/.codex. #[arg(long, env = "CODEX_HOME")] @@ -104,7 +103,7 @@ fn resolve_db_path(args: &Args) -> anyhow::Result { } let codex_home = args.codex_home.clone().unwrap_or_else(default_codex_home); - Ok(codex_home.join(STATE_DB_FILENAME)) + Ok(codex_state::state_db_path(codex_home.as_path())) } fn default_codex_home() -> PathBuf { diff --git a/codex-rs/state/src/lib.rs b/codex-rs/state/src/lib.rs index c08c76a1cf..2d37ecee96 100644 --- a/codex-rs/state/src/lib.rs +++ b/codex-rs/state/src/lib.rs @@ -29,6 +29,9 @@ pub use model::ThreadMetadata; pub use model::ThreadMetadataBuilder; pub use model::ThreadsPage; pub use runtime::STATE_DB_FILENAME; +pub use runtime::STATE_DB_VERSION; +pub use runtime::state_db_filename; +pub use runtime::state_db_path; /// Errors encountered during DB operations. Tags: [stage] pub const DB_ERROR_METRIC: &str = "codex.db.error"; diff --git a/codex-rs/state/src/runtime.rs b/codex-rs/state/src/runtime.rs index c64b7d56c6..f4c2d76e8a 100644 --- a/codex-rs/state/src/runtime.rs +++ b/codex-rs/state/src/runtime.rs @@ -35,7 +35,8 @@ use std::sync::Arc; use std::time::Duration; use tracing::warn; -pub const STATE_DB_FILENAME: &str = "state.sqlite"; +pub const STATE_DB_FILENAME: &str = "state"; +pub const STATE_DB_VERSION: u32 = 2; const METRIC_DB_INIT: &str = "codex.db.init"; @@ -56,7 +57,8 @@ impl StateRuntime { otel: Option, ) -> anyhow::Result> { tokio::fs::create_dir_all(&codex_home).await?; - let state_path = codex_home.join(STATE_DB_FILENAME); + remove_legacy_state_files(&codex_home).await; + let state_path = state_db_path(codex_home.as_path()); let existed = tokio::fs::try_exists(&state_path).await.unwrap_or(false); let pool = match open_sqlite(&state_path).await { Ok(db) => Arc::new(db), @@ -624,6 +626,77 @@ async fn open_sqlite(path: &Path) -> anyhow::Result { Ok(pool) } +pub fn state_db_filename() -> String { + format!("{STATE_DB_FILENAME}_{STATE_DB_VERSION}.sqlite") +} + +pub fn state_db_path(codex_home: &Path) -> PathBuf { + codex_home.join(state_db_filename()) +} + +async fn remove_legacy_state_files(codex_home: &Path) { + let current_name = state_db_filename(); + let mut entries = match tokio::fs::read_dir(codex_home).await { + Ok(entries) => entries, + Err(err) => { + warn!( + "failed to read codex_home for state db cleanup {}: {err}", + codex_home.display() + ); + return; + } + }; + while let Ok(Some(entry)) = entries.next_entry().await { + if !entry + .file_type() + .await + .map(|file_type| file_type.is_file()) + .unwrap_or(false) + { + continue; + } + let file_name = entry.file_name(); + let file_name = file_name.to_string_lossy(); + if !should_remove_state_file(file_name.as_ref(), current_name.as_str()) { + continue; + } + + let legacy_path = entry.path(); + if let Err(err) = tokio::fs::remove_file(&legacy_path).await { + warn!( + "failed to remove legacy state db file {}: {err}", + legacy_path.display() + ); + } + } +} + +fn should_remove_state_file(file_name: &str, current_name: &str) -> bool { + let mut base_name = file_name; + for suffix in ["-wal", "-shm", "-journal"] { + if let Some(stripped) = file_name.strip_suffix(suffix) { + base_name = stripped; + break; + } + } + if base_name == current_name { + return false; + } + let unversioned_name = format!("{STATE_DB_FILENAME}.sqlite"); + if base_name == unversioned_name { + return true; + } + + let Some(version_with_extension) = base_name.strip_prefix(&format!("{STATE_DB_FILENAME}_")) + else { + return false; + }; + let Some(version_suffix) = version_with_extension.strip_suffix(".sqlite") else { + return false; + }; + !version_suffix.is_empty() && version_suffix.chars().all(|ch| ch.is_ascii_digit()) +} + fn push_thread_filters<'a>( builder: &mut QueryBuilder<'a, Sqlite>, archived_only: bool, @@ -692,3 +765,97 @@ fn push_thread_order_and_limit( builder.push(" LIMIT "); builder.push_bind(limit as i64); } + +#[cfg(test)] +mod tests { + use super::STATE_DB_FILENAME; + use super::STATE_DB_VERSION; + use super::StateRuntime; + use super::state_db_filename; + use pretty_assertions::assert_eq; + use std::path::PathBuf; + use std::time::SystemTime; + use std::time::UNIX_EPOCH; + + fn unique_temp_dir() -> PathBuf { + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0, |duration| duration.as_nanos()); + std::env::temp_dir().join(format!("codex-state-runtime-test-{nanos}")) + } + + #[tokio::test] + async fn init_removes_legacy_state_db_files() { + let codex_home = unique_temp_dir(); + tokio::fs::create_dir_all(&codex_home) + .await + .expect("create codex_home"); + + let current_name = state_db_filename(); + let previous_version = STATE_DB_VERSION.saturating_sub(1); + let unversioned_name = format!("{STATE_DB_FILENAME}.sqlite"); + for suffix in ["", "-wal", "-shm", "-journal"] { + let path = codex_home.join(format!("{unversioned_name}{suffix}")); + tokio::fs::write(path, b"legacy") + .await + .expect("write legacy"); + let old_version_path = codex_home.join(format!( + "{STATE_DB_FILENAME}_{previous_version}.sqlite{suffix}" + )); + tokio::fs::write(old_version_path, b"old_version") + .await + .expect("write old version"); + } + let unrelated_path = codex_home.join("state.sqlite_backup"); + tokio::fs::write(&unrelated_path, b"keep") + .await + .expect("write unrelated"); + let numeric_path = codex_home.join("123"); + tokio::fs::write(&numeric_path, b"keep") + .await + .expect("write numeric"); + + let _runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string(), None) + .await + .expect("initialize runtime"); + + for suffix in ["", "-wal", "-shm", "-journal"] { + let legacy_path = codex_home.join(format!("{unversioned_name}{suffix}")); + assert_eq!( + tokio::fs::try_exists(&legacy_path) + .await + .expect("check legacy path"), + false + ); + let old_version_path = codex_home.join(format!( + "{STATE_DB_FILENAME}_{previous_version}.sqlite{suffix}" + )); + assert_eq!( + tokio::fs::try_exists(&old_version_path) + .await + .expect("check old version path"), + false + ); + } + assert_eq!( + tokio::fs::try_exists(codex_home.join(current_name)) + .await + .expect("check new db path"), + true + ); + assert_eq!( + tokio::fs::try_exists(&unrelated_path) + .await + .expect("check unrelated path"), + true + ); + assert_eq!( + tokio::fs::try_exists(&numeric_path) + .await + .expect("check numeric path"), + true + ); + + let _ = tokio::fs::remove_dir_all(codex_home).await; + } +} From 0efd33f7f4291d6a056dc39a9654ce04ebe0400c Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Wed, 4 Feb 2026 08:38:06 -0800 Subject: [PATCH 3/8] Update tests to stop using sse_completed fixture (#10638) Summary: - replace the `sse_completed` fixture and related JSON template with direct `responses::ev_completed` payload builders - cascade the new SSE helpers through all affected core tests for consistency and clarity - remove legacy fixtures that were no longer needed once the helpers are in place Testing: - Not run (not requested) --- codex-rs/core/tests/common/lib.rs | 35 ----- .../tests/fixtures/completed_template.json | 16 --- codex-rs/core/tests/suite/client.rs | 124 ++++++++++++++---- .../tests/suite/collaboration_instructions.rs | 100 +++++++++++--- codex-rs/core/tests/suite/fork_thread.rs | 10 +- .../core/tests/suite/hierarchical_agents.rs | 20 ++- codex-rs/core/tests/suite/model_tools.rs | 10 +- .../core/tests/suite/permissions_messages.rs | 82 +++++++++--- codex-rs/core/tests/suite/personality.rs | 48 +++++-- codex-rs/core/tests/suite/prompt_caching.rs | 99 +++++++++++--- codex-rs/core/tests/suite/sqlite_state.rs | 10 +- .../suite/stream_error_allows_next_turn.rs | 16 ++- .../core/tests/suite/stream_no_completed.rs | 16 ++- codex-rs/core/tests/suite/web_search.rs | 36 +++-- 14 files changed, 428 insertions(+), 194 deletions(-) delete mode 100644 codex-rs/core/tests/fixtures/completed_template.json diff --git a/codex-rs/core/tests/common/lib.rs b/codex-rs/core/tests/common/lib.rs index 5df2a74edd..3feadb0722 100644 --- a/codex-rs/core/tests/common/lib.rs +++ b/codex-rs/core/tests/common/lib.rs @@ -1,7 +1,6 @@ #![expect(clippy::expect_used)] use codex_utils_cargo_bin::CargoBinError; -use codex_utils_cargo_bin::find_resource; use tempfile::TempDir; use codex_core::CodexThread; @@ -147,40 +146,6 @@ pub fn load_sse_fixture_with_id_from_str(raw: &str, id: &str) -> String { .collect() } -/// Same as [`load_sse_fixture`], but replaces the placeholder `__ID__` in the -/// fixture template with the supplied identifier before parsing. This lets a -/// single JSON template be reused by multiple tests that each need a unique -/// `response_id`. -pub fn load_sse_fixture_with_id(path: impl AsRef, id: &str) -> String { - let p = path.as_ref(); - let full_path = match find_resource!(p) { - Ok(p) => p, - Err(err) => panic!( - "failed to find fixture template at {:?}: {err}", - path.as_ref() - ), - }; - - let raw = std::fs::read_to_string(full_path).expect("read fixture template"); - let replaced = raw.replace("__ID__", id); - let events: Vec = - serde_json::from_str(&replaced).expect("parse JSON fixture"); - events - .into_iter() - .map(|e| { - let kind = e - .get("type") - .and_then(|v| v.as_str()) - .expect("fixture event missing type"); - if e.as_object().map(|o| o.len() == 1).unwrap_or(false) { - format!("event: {kind}\n\n") - } else { - format!("event: {kind}\ndata: {e}\n\n") - } - }) - .collect() -} - pub async fn wait_for_event(codex: &CodexThread, predicate: F) -> codex_core::protocol::EventMsg where F: FnMut(&codex_core::protocol::EventMsg) -> bool, diff --git a/codex-rs/core/tests/fixtures/completed_template.json b/codex-rs/core/tests/fixtures/completed_template.json deleted file mode 100644 index 1774dc5e84..0000000000 --- a/codex-rs/core/tests/fixtures/completed_template.json +++ /dev/null @@ -1,16 +0,0 @@ -[ - { - "type": "response.completed", - "response": { - "id": "__ID__", - "usage": { - "input_tokens": 0, - "input_tokens_details": null, - "output_tokens": 0, - "output_tokens_details": null, - "total_tokens": 0 - }, - "output": [] - } - } -] diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index bd50708a2c..62135127f2 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -36,8 +36,9 @@ use codex_protocol::models::WebSearchAction; use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::user_input::UserInput; use core_test_support::load_default_config_for_test; -use core_test_support::load_sse_fixture_with_id; +use core_test_support::responses::ev_completed; use core_test_support::responses::ev_completed_with_tokens; +use core_test_support::responses::ev_response_created; use core_test_support::responses::mount_sse_once; use core_test_support::responses::mount_sse_once_match; use core_test_support::responses::mount_sse_sequence; @@ -64,11 +65,6 @@ use wiremock::matchers::method; use wiremock::matchers::path; use wiremock::matchers::query_param; -/// Build minimal SSE stream with completed marker using the JSON fixture. -fn sse_completed(id: &str) -> String { - load_sse_fixture_with_id("../fixtures/completed_template.json", id) -} - #[expect(clippy::unwrap_used)] fn assert_message_role(request_body: &serde_json::Value, role: &str) { assert_eq!(request_body["role"].as_str().unwrap(), role); @@ -259,7 +255,11 @@ async fn resume_includes_initial_messages_and_sends_prior_items() { // Mock server that will receive the resumed request let server = MockServer::start().await; - let resp_mock = mount_sse_once(&server, sse_completed("resp1")).await; + let resp_mock = mount_sse_once( + &server, + sse(vec![ev_response_created("resp1"), ev_completed("resp1")]), + ) + .await; // Configure Codex to resume from our file let codex_home = Arc::new(TempDir::new().unwrap()); @@ -377,7 +377,11 @@ async fn includes_conversation_id_and_model_headers_in_request() { // Mock server let server = MockServer::start().await; - let resp_mock = mount_sse_once(&server, sse_completed("resp1")).await; + let resp_mock = mount_sse_once( + &server, + sse(vec![ev_response_created("resp1"), ev_completed("resp1")]), + ) + .await; let mut builder = test_codex().with_auth(CodexAuth::from_api_key("Test API Key")); let test = builder @@ -418,7 +422,11 @@ async fn includes_base_instructions_override_in_request() { skip_if_no_network!(); // Mock server let server = MockServer::start().await; - let resp_mock = mount_sse_once(&server, sse_completed("resp1")).await; + let resp_mock = mount_sse_once( + &server, + sse(vec![ev_response_created("resp1"), ev_completed("resp1")]), + ) + .await; let mut builder = test_codex() .with_auth(CodexAuth::from_api_key("Test API Key")) @@ -462,7 +470,11 @@ async fn chatgpt_auth_sends_correct_request() { // Mock server let server = MockServer::start().await; - let resp_mock = mount_sse_once(&server, sse_completed("resp1")).await; + let resp_mock = mount_sse_once( + &server, + sse(vec![ev_response_created("resp1"), ev_completed("resp1")]), + ) + .await; let mut model_provider = built_in_model_providers()["openai"].clone(); model_provider.base_url = Some(format!("{}/api/codex", server.uri())); @@ -524,7 +536,10 @@ async fn prefers_apikey_when_config_prefers_apikey_even_with_chatgpt_tokens() { let first = ResponseTemplate::new(200) .insert_header("content-type", "text/event-stream") - .set_body_raw(sse_completed("resp1"), "text/event-stream"); + .set_body_raw( + sse(vec![ev_response_created("resp1"), ev_completed("resp1")]), + "text/event-stream", + ); // Expect API key header, no ChatGPT account header required. Mock::given(method("POST")) @@ -590,7 +605,11 @@ async fn includes_user_instructions_message_in_request() { skip_if_no_network!(); let server = MockServer::start().await; - let resp_mock = mount_sse_once(&server, sse_completed("resp1")).await; + let resp_mock = mount_sse_once( + &server, + sse(vec![ev_response_created("resp1"), ev_completed("resp1")]), + ) + .await; let mut builder = test_codex() .with_auth(CodexAuth::from_api_key("Test API Key")) @@ -652,7 +671,11 @@ async fn skills_append_to_instructions() { skip_if_no_network!(); let server = MockServer::start().await; - let resp_mock = mount_sse_once(&server, sse_completed("resp1")).await; + let resp_mock = mount_sse_once( + &server, + sse(vec![ev_response_created("resp1"), ev_completed("resp1")]), + ) + .await; let codex_home = Arc::new(TempDir::new().unwrap()); let skill_dir = codex_home.path().join("skills/demo"); @@ -720,7 +743,11 @@ async fn includes_configured_effort_in_request() -> anyhow::Result<()> { skip_if_no_network!(Ok(())); let server = MockServer::start().await; - let resp_mock = mount_sse_once(&server, sse_completed("resp1")).await; + let resp_mock = mount_sse_once( + &server, + sse(vec![ev_response_created("resp1"), ev_completed("resp1")]), + ) + .await; let TestCodex { codex, .. } = test_codex() .with_model("gpt-5.1-codex") .with_config(|config| { @@ -761,7 +788,11 @@ async fn includes_no_effort_in_request() -> anyhow::Result<()> { skip_if_no_network!(Ok(())); let server = MockServer::start().await; - let resp_mock = mount_sse_once(&server, sse_completed("resp1")).await; + let resp_mock = mount_sse_once( + &server, + sse(vec![ev_response_created("resp1"), ev_completed("resp1")]), + ) + .await; let TestCodex { codex, .. } = test_codex() .with_model("gpt-5.1-codex") .build(&server) @@ -800,7 +831,11 @@ async fn includes_default_reasoning_effort_in_request_when_defined_by_model_info skip_if_no_network!(Ok(())); let server = MockServer::start().await; - let resp_mock = mount_sse_once(&server, sse_completed("resp1")).await; + let resp_mock = mount_sse_once( + &server, + sse(vec![ev_response_created("resp1"), ev_completed("resp1")]), + ) + .await; let TestCodex { codex, .. } = test_codex().with_model("gpt-5.1").build(&server).await?; codex @@ -835,7 +870,11 @@ async fn user_turn_collaboration_mode_overrides_model_and_effort() -> anyhow::Re skip_if_no_network!(Ok(())); let server = MockServer::start().await; - let resp_mock = mount_sse_once(&server, sse_completed("resp1")).await; + let resp_mock = mount_sse_once( + &server, + sse(vec![ev_response_created("resp1"), ev_completed("resp1")]), + ) + .await; let TestCodex { codex, config, @@ -893,7 +932,11 @@ async fn configured_reasoning_summary_is_sent() -> anyhow::Result<()> { skip_if_no_network!(Ok(())); let server = MockServer::start().await; - let resp_mock = mount_sse_once(&server, sse_completed("resp1")).await; + let resp_mock = mount_sse_once( + &server, + sse(vec![ev_response_created("resp1"), ev_completed("resp1")]), + ) + .await; let TestCodex { codex, .. } = test_codex() .with_config(|config| { config.model_reasoning_summary = ReasoningSummary::Concise; @@ -933,7 +976,11 @@ async fn reasoning_summary_is_omitted_when_disabled() -> anyhow::Result<()> { skip_if_no_network!(Ok(())); let server = MockServer::start().await; - let resp_mock = mount_sse_once(&server, sse_completed("resp1")).await; + let resp_mock = mount_sse_once( + &server, + sse(vec![ev_response_created("resp1"), ev_completed("resp1")]), + ) + .await; let TestCodex { codex, .. } = test_codex() .with_config(|config| { config.model_reasoning_summary = ReasoningSummary::None; @@ -972,7 +1019,11 @@ async fn includes_default_verbosity_in_request() -> anyhow::Result<()> { skip_if_no_network!(Ok(())); let server = MockServer::start().await; - let resp_mock = mount_sse_once(&server, sse_completed("resp1")).await; + let resp_mock = mount_sse_once( + &server, + sse(vec![ev_response_created("resp1"), ev_completed("resp1")]), + ) + .await; let TestCodex { codex, .. } = test_codex().with_model("gpt-5.1").build(&server).await?; codex @@ -1007,7 +1058,11 @@ async fn configured_verbosity_not_sent_for_models_without_support() -> anyhow::R skip_if_no_network!(Ok(())); let server = MockServer::start().await; - let resp_mock = mount_sse_once(&server, sse_completed("resp1")).await; + let resp_mock = mount_sse_once( + &server, + sse(vec![ev_response_created("resp1"), ev_completed("resp1")]), + ) + .await; let TestCodex { codex, .. } = test_codex() .with_model("gpt-5.1-codex") .with_config(|config| { @@ -1047,7 +1102,11 @@ async fn configured_verbosity_is_sent() -> anyhow::Result<()> { skip_if_no_network!(Ok(())); let server = MockServer::start().await; - let resp_mock = mount_sse_once(&server, sse_completed("resp1")).await; + let resp_mock = mount_sse_once( + &server, + sse(vec![ev_response_created("resp1"), ev_completed("resp1")]), + ) + .await; let TestCodex { codex, .. } = test_codex() .with_model("gpt-5.1") .with_config(|config| { @@ -1088,7 +1147,11 @@ async fn includes_developer_instructions_message_in_request() { skip_if_no_network!(); let server = MockServer::start().await; - let resp_mock = mount_sse_once(&server, sse_completed("resp1")).await; + let resp_mock = mount_sse_once( + &server, + sse(vec![ev_response_created("resp1"), ev_completed("resp1")]), + ) + .await; let mut builder = test_codex() .with_auth(CodexAuth::from_api_key("Test API Key")) .with_config(|config| { @@ -1570,7 +1633,10 @@ async fn context_window_error_sets_total_tokens_to_model_window() -> anyhow::Res mount_sse_once_match( &server, body_string_contains("seed turn"), - sse_completed("resp_seed"), + sse(vec![ + ev_response_created("resp_seed"), + ev_completed("resp_seed"), + ]), ) .await; @@ -1656,7 +1722,10 @@ async fn azure_overrides_assign_properties_used_for_responses_url() { // First request – must NOT include `previous_response_id`. let first = ResponseTemplate::new(200) .insert_header("content-type", "text/event-stream") - .set_body_raw(sse_completed("resp1"), "text/event-stream"); + .set_body_raw( + sse(vec![ev_response_created("resp1"), ev_completed("resp1")]), + "text/event-stream", + ); // Expect POST to /openai/responses with api-version query param Mock::given(method("POST")) @@ -1737,7 +1806,10 @@ async fn env_var_overrides_loaded_auth() { // First request – must NOT include `previous_response_id`. let first = ResponseTemplate::new(200) .insert_header("content-type", "text/event-stream") - .set_body_raw(sse_completed("resp1"), "text/event-stream"); + .set_body_raw( + sse(vec![ev_response_created("resp1"), ev_completed("resp1")]), + "text/event-stream", + ); // Expect POST to /openai/responses with api-version query param Mock::given(method("POST")) diff --git a/codex-rs/core/tests/suite/collaboration_instructions.rs b/codex-rs/core/tests/suite/collaboration_instructions.rs index 21a0b27613..dc9905379e 100644 --- a/codex-rs/core/tests/suite/collaboration_instructions.rs +++ b/codex-rs/core/tests/suite/collaboration_instructions.rs @@ -18,10 +18,6 @@ use core_test_support::wait_for_event; use pretty_assertions::assert_eq; use serde_json::Value; -fn sse_completed(id: &str) -> String { - sse(vec![ev_response_created(id), ev_completed(id)]) -} - fn collab_mode_with_mode_and_instructions( mode: ModeKind, instructions: Option<&str>, @@ -72,7 +68,11 @@ async fn no_collaboration_instructions_by_default() -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; - let req = mount_sse_once(&server, sse_completed("resp-1")).await; + let req = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; let test = test_codex().build(&server).await?; @@ -100,7 +100,11 @@ async fn user_input_includes_collaboration_instructions_after_override() -> Resu skip_if_no_network!(Ok(())); let server = start_mock_server().await; - let req = mount_sse_once(&server, sse_completed("resp-1")).await; + let req = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; let test = test_codex().build(&server).await?; @@ -144,7 +148,11 @@ async fn collaboration_instructions_added_on_user_turn() -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; - let req = mount_sse_once(&server, sse_completed("resp-1")).await; + let req = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; let test = test_codex().build(&server).await?; let collab_text = "turn instructions"; @@ -182,7 +190,11 @@ async fn override_then_next_turn_uses_updated_collaboration_instructions() -> Re skip_if_no_network!(Ok(())); let server = start_mock_server().await; - let req = mount_sse_once(&server, sse_completed("resp-1")).await; + let req = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; let test = test_codex().build(&server).await?; let collab_text = "override instructions"; @@ -226,7 +238,11 @@ async fn user_turn_overrides_collaboration_instructions_after_override() -> Resu skip_if_no_network!(Ok(())); let server = start_mock_server().await; - let req = mount_sse_once(&server, sse_completed("resp-1")).await; + let req = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; let test = test_codex().build(&server).await?; let base_text = "base instructions"; @@ -282,8 +298,16 @@ async fn collaboration_mode_update_emits_new_instruction_message() -> Result<()> skip_if_no_network!(Ok(())); let server = start_mock_server().await; - let _req1 = mount_sse_once(&server, sse_completed("resp-1")).await; - let req2 = mount_sse_once(&server, sse_completed("resp-2")).await; + let _req1 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; + let req2 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + ) + .await; let test = test_codex().build(&server).await?; let first_text = "first instructions"; @@ -354,8 +378,16 @@ async fn collaboration_mode_update_noop_does_not_append() -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; - let _req1 = mount_sse_once(&server, sse_completed("resp-1")).await; - let req2 = mount_sse_once(&server, sse_completed("resp-2")).await; + let _req1 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; + let req2 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + ) + .await; let test = test_codex().build(&server).await?; let collab_text = "same instructions"; @@ -423,8 +455,16 @@ async fn collaboration_mode_update_emits_new_instruction_message_when_mode_chang skip_if_no_network!(Ok(())); let server = start_mock_server().await; - let _req1 = mount_sse_once(&server, sse_completed("resp-1")).await; - let req2 = mount_sse_once(&server, sse_completed("resp-2")).await; + let _req1 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; + let req2 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + ) + .await; let test = test_codex().build(&server).await?; let default_text = "default mode instructions"; @@ -501,8 +541,16 @@ async fn collaboration_mode_update_noop_does_not_append_when_mode_is_unchanged() skip_if_no_network!(Ok(())); let server = start_mock_server().await; - let _req1 = mount_sse_once(&server, sse_completed("resp-1")).await; - let req2 = mount_sse_once(&server, sse_completed("resp-2")).await; + let _req1 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; + let req2 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + ) + .await; let test = test_codex().build(&server).await?; let collab_text = "mode-stable instructions"; @@ -576,8 +624,16 @@ async fn resume_replays_collaboration_instructions() -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; - let _req1 = mount_sse_once(&server, sse_completed("resp-1")).await; - let req2 = mount_sse_once(&server, sse_completed("resp-2")).await; + let _req1 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; + let req2 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + ) + .await; let mut builder = test_codex(); let initial = builder.build(&server).await?; @@ -642,7 +698,11 @@ async fn empty_collaboration_instructions_are_ignored() -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; - let req = mount_sse_once(&server, sse_completed("resp-1")).await; + let req = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; let test = test_codex().build(&server).await?; diff --git a/codex-rs/core/tests/suite/fork_thread.rs b/codex-rs/core/tests/suite/fork_thread.rs index 55474a68aa..d5363cc0a9 100644 --- a/codex-rs/core/tests/suite/fork_thread.rs +++ b/codex-rs/core/tests/suite/fork_thread.rs @@ -6,6 +6,9 @@ use codex_core::protocol::RolloutItem; use codex_core::protocol::RolloutLine; use codex_protocol::items::TurnItem; use codex_protocol::user_input::UserInput; +use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_response_created; +use core_test_support::responses::sse; use core_test_support::skip_if_no_network; use core_test_support::test_codex::test_codex; use core_test_support::wait_for_event; @@ -15,18 +18,13 @@ use wiremock::ResponseTemplate; use wiremock::matchers::method; use wiremock::matchers::path; -/// Build minimal SSE stream with completed marker using the JSON fixture. -fn sse_completed(id: &str) -> String { - core_test_support::load_sse_fixture_with_id("../fixtures/completed_template.json", id) -} - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn fork_thread_twice_drops_to_first_message() { skip_if_no_network!(); // Start a mock server that completes three turns. let server = MockServer::start().await; - let sse = sse_completed("resp"); + let sse = sse(vec![ev_response_created("resp"), ev_completed("resp")]); let first = ResponseTemplate::new(200) .insert_header("content-type", "text/event-stream") .set_body_raw(sse.clone(), "text/event-stream"); diff --git a/codex-rs/core/tests/suite/hierarchical_agents.rs b/codex-rs/core/tests/suite/hierarchical_agents.rs index 5dfa636489..5f1f84340d 100644 --- a/codex-rs/core/tests/suite/hierarchical_agents.rs +++ b/codex-rs/core/tests/suite/hierarchical_agents.rs @@ -1,20 +1,22 @@ use codex_core::features::Feature; -use core_test_support::load_sse_fixture_with_id; +use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_response_created; use core_test_support::responses::mount_sse_once; +use core_test_support::responses::sse; use core_test_support::responses::start_mock_server; use core_test_support::test_codex::test_codex; const HIERARCHICAL_AGENTS_SNIPPET: &str = "Files called AGENTS.md commonly appear in many places inside a container"; -fn sse_completed(id: &str) -> String { - load_sse_fixture_with_id("../fixtures/completed_template.json", id) -} - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn hierarchical_agents_appends_to_project_doc_in_user_instructions() { let server = start_mock_server().await; - let resp_mock = mount_sse_once(&server, sse_completed("resp1")).await; + let resp_mock = mount_sse_once( + &server, + sse(vec![ev_response_created("resp1"), ev_completed("resp1")]), + ) + .await; let mut builder = test_codex().with_config(|config| { config.features.enable(Feature::ChildAgentsMd); @@ -49,7 +51,11 @@ async fn hierarchical_agents_appends_to_project_doc_in_user_instructions() { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn hierarchical_agents_emits_when_no_project_doc() { let server = start_mock_server().await; - let resp_mock = mount_sse_once(&server, sse_completed("resp1")).await; + let resp_mock = mount_sse_once( + &server, + sse(vec![ev_response_created("resp1"), ev_completed("resp1")]), + ) + .await; let mut builder = test_codex().with_config(|config| { config.features.enable(Feature::ChildAgentsMd); diff --git a/codex-rs/core/tests/suite/model_tools.rs b/codex-rs/core/tests/suite/model_tools.rs index b08b5281c8..c687545f26 100644 --- a/codex-rs/core/tests/suite/model_tools.rs +++ b/codex-rs/core/tests/suite/model_tools.rs @@ -2,16 +2,11 @@ use codex_core::features::Feature; use codex_protocol::config_types::WebSearchMode; -use core_test_support::load_sse_fixture_with_id; use core_test_support::responses; use core_test_support::responses::start_mock_server; use core_test_support::skip_if_no_network; use core_test_support::test_codex::test_codex; -fn sse_completed(id: &str) -> String { - load_sse_fixture_with_id("../fixtures/completed_template.json", id) -} - #[allow(clippy::expect_used)] fn tool_identifiers(body: &serde_json::Value) -> Vec { body["tools"] @@ -31,7 +26,10 @@ fn tool_identifiers(body: &serde_json::Value) -> Vec { #[allow(clippy::expect_used)] async fn collect_tool_identifiers_for_model(model: &str) -> Vec { let server = start_mock_server().await; - let sse = sse_completed(model); + let sse = responses::sse(vec![ + responses::ev_response_created(model), + responses::ev_completed(model), + ]); let resp_mock = responses::mount_sse_once(&server, sse).await; let mut builder = test_codex() diff --git a/codex-rs/core/tests/suite/permissions_messages.rs b/codex-rs/core/tests/suite/permissions_messages.rs index c54203b4f1..fc1f0fa0b6 100644 --- a/codex-rs/core/tests/suite/permissions_messages.rs +++ b/codex-rs/core/tests/suite/permissions_messages.rs @@ -43,16 +43,16 @@ fn permissions_texts(input: &[serde_json::Value]) -> Vec { .collect() } -fn sse_completed(id: &str) -> String { - sse(vec![ev_response_created(id), ev_completed(id)]) -} - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn permissions_message_sent_once_on_start() -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; - let req = mount_sse_once(&server, sse_completed("resp-1")).await; + let req = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; let mut builder = test_codex().with_config(move |config| { config.approval_policy = Constrained::allow_any(AskForApproval::OnRequest); @@ -84,8 +84,16 @@ async fn permissions_message_added_on_override_change() -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; - let req1 = mount_sse_once(&server, sse_completed("resp-1")).await; - let req2 = mount_sse_once(&server, sse_completed("resp-2")).await; + let req1 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; + let req2 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + ) + .await; let mut builder = test_codex().with_config(move |config| { config.approval_policy = Constrained::allow_any(AskForApproval::OnRequest); @@ -148,8 +156,16 @@ async fn permissions_message_not_added_when_no_change() -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; - let req1 = mount_sse_once(&server, sse_completed("resp-1")).await; - let req2 = mount_sse_once(&server, sse_completed("resp-2")).await; + let req1 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; + let req2 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + ) + .await; let mut builder = test_codex().with_config(move |config| { config.approval_policy = Constrained::allow_any(AskForApproval::OnRequest); @@ -197,9 +213,21 @@ async fn resume_replays_permissions_messages() -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; - let _req1 = mount_sse_once(&server, sse_completed("resp-1")).await; - let _req2 = mount_sse_once(&server, sse_completed("resp-2")).await; - let req3 = mount_sse_once(&server, sse_completed("resp-3")).await; + let _req1 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; + let _req2 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + ) + .await; + let req3 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-3"), ev_completed("resp-3")]), + ) + .await; let mut builder = test_codex().with_config(|config| { config.approval_policy = Constrained::allow_any(AskForApproval::OnRequest); @@ -279,10 +307,26 @@ async fn resume_and_fork_append_permissions_messages() -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; - let _req1 = mount_sse_once(&server, sse_completed("resp-1")).await; - let req2 = mount_sse_once(&server, sse_completed("resp-2")).await; - let req3 = mount_sse_once(&server, sse_completed("resp-3")).await; - let req4 = mount_sse_once(&server, sse_completed("resp-4")).await; + let _req1 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; + let req2 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + ) + .await; + let req3 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-3"), ev_completed("resp-3")]), + ) + .await; + let req4 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-4"), ev_completed("resp-4")]), + ) + .await; let mut builder = test_codex().with_config(|config| { config.approval_policy = Constrained::allow_any(AskForApproval::OnRequest); @@ -404,7 +448,11 @@ async fn permissions_message_includes_writable_roots() -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; - let req = mount_sse_once(&server, sse_completed("resp-1")).await; + let req = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; let writable = TempDir::new()?; let writable_root = AbsolutePathBuf::try_from(writable.path())?; let sandbox_policy = SandboxPolicy::WorkspaceWrite { diff --git a/codex-rs/core/tests/suite/personality.rs b/codex-rs/core/tests/suite/personality.rs index 87978ceb5a..919ce852ed 100644 --- a/codex-rs/core/tests/suite/personality.rs +++ b/codex-rs/core/tests/suite/personality.rs @@ -42,10 +42,6 @@ const LOCAL_FRIENDLY_TEMPLATE: &str = "You optimize for team morale and being a supportive teammate as much as code quality."; const LOCAL_PRAGMATIC_TEMPLATE: &str = "You are a deeply pragmatic, effective software engineer."; -fn sse_completed(id: &str) -> String { - sse(vec![ev_response_created(id), ev_completed(id)]) -} - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn personality_does_not_mutate_base_instructions_without_template() { let codex_home = TempDir::new().expect("create temp dir"); @@ -82,7 +78,11 @@ async fn user_turn_personality_none_does_not_add_update_message() -> anyhow::Res skip_if_no_network!(Ok(())); let server = start_mock_server().await; - let resp_mock = mount_sse_once(&server, sse_completed("resp-1")).await; + let resp_mock = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; let mut builder = test_codex() .with_model("gpt-5.2-codex") .with_config(|config| { @@ -128,7 +128,11 @@ async fn config_personality_some_sets_instructions_template() -> anyhow::Result< skip_if_no_network!(Ok(())); let server = start_mock_server().await; - let resp_mock = mount_sse_once(&server, sse_completed("resp-1")).await; + let resp_mock = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; let mut builder = test_codex() .with_model("gpt-5.2-codex") .with_config(|config| { @@ -184,7 +188,10 @@ async fn user_turn_personality_some_adds_update_message() -> anyhow::Result<()> let server = start_mock_server().await; let resp_mock = mount_sse_sequence( &server, - vec![sse_completed("resp-1"), sse_completed("resp-2")], + vec![ + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + ], ) .await; let mut builder = test_codex() @@ -280,7 +287,10 @@ async fn user_turn_personality_same_value_does_not_add_update_message() -> anyho let server = start_mock_server().await; let resp_mock = mount_sse_sequence( &server, - vec![sse_completed("resp-1"), sse_completed("resp-2")], + vec![ + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + ], ) .await; let mut builder = test_codex() @@ -387,7 +397,10 @@ async fn user_turn_personality_skips_if_feature_disabled() -> anyhow::Result<()> let server = start_mock_server().await; let resp_mock = mount_sse_sequence( &server, - vec![sse_completed("resp-1"), sse_completed("resp-2")], + vec![ + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + ], ) .await; let mut builder = test_codex() @@ -524,7 +537,11 @@ async fn ignores_remote_personality_if_remote_models_disabled() -> anyhow::Resul ) .await; - let resp_mock = mount_sse_once(&server, sse_completed("resp-1")).await; + let resp_mock = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; let mut builder = test_codex() .with_auth(codex_core::CodexAuth::create_dummy_chatgpt_auth_for_testing()) @@ -640,7 +657,11 @@ async fn remote_model_friendly_personality_instructions_with_feature() -> anyhow ) .await; - let resp_mock = mount_sse_once(&server, sse_completed("resp-1")).await; + let resp_mock = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; let mut builder = test_codex() .with_auth(codex_core::CodexAuth::create_dummy_chatgpt_auth_for_testing()) @@ -753,7 +774,10 @@ async fn user_turn_personality_remote_model_template_includes_update_message() - let resp_mock = mount_sse_sequence( &server, - vec![sse_completed("resp-1"), sse_completed("resp-2")], + vec![ + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + ], ) .await; diff --git a/codex-rs/core/tests/suite/prompt_caching.rs b/codex-rs/core/tests/suite/prompt_caching.rs index bf1819cbf1..810cd52497 100644 --- a/codex-rs/core/tests/suite/prompt_caching.rs +++ b/codex-rs/core/tests/suite/prompt_caching.rs @@ -18,8 +18,10 @@ use codex_protocol::config_types::WebSearchMode; use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::user_input::UserInput; use codex_utils_absolute_path::AbsolutePathBuf; -use core_test_support::load_sse_fixture_with_id; +use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_response_created; use core_test_support::responses::mount_sse_once; +use core_test_support::responses::sse; use core_test_support::responses::start_mock_server; use core_test_support::skip_if_no_network; use core_test_support::test_codex::TestCodex; @@ -46,11 +48,6 @@ fn default_env_context_str(cwd: &str, shell: &Shell) -> String { ) } -/// Build minimal SSE stream with completed marker using the JSON fixture. -fn sse_completed(id: &str) -> String { - load_sse_fixture_with_id("../fixtures/completed_template.json", id) -} - fn assert_tool_names(body: &serde_json::Value, expected_names: &[&str]) { assert_eq!( body["tools"] @@ -79,8 +76,16 @@ async fn prompt_tools_are_consistent_across_requests() -> anyhow::Result<()> { use pretty_assertions::assert_eq; let server = start_mock_server().await; - let req1 = mount_sse_once(&server, sse_completed("resp-1")).await; - let req2 = mount_sse_once(&server, sse_completed("resp-2")).await; + let req1 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; + let req2 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + ) + .await; let TestCodex { codex, @@ -172,8 +177,16 @@ async fn codex_mini_latest_tools() -> anyhow::Result<()> { use pretty_assertions::assert_eq; let server = start_mock_server().await; - let req1 = mount_sse_once(&server, sse_completed("resp-1")).await; - let req2 = mount_sse_once(&server, sse_completed("resp-2")).await; + let req1 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; + let req2 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + ) + .await; let TestCodex { codex, .. } = test_codex() .with_config(|config| { @@ -238,8 +251,16 @@ async fn prefixes_context_and_instructions_once_and_consistently_across_requests use pretty_assertions::assert_eq; let server = start_mock_server().await; - let req1 = mount_sse_once(&server, sse_completed("resp-1")).await; - let req2 = mount_sse_once(&server, sse_completed("resp-2")).await; + let req1 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; + let req2 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + ) + .await; let TestCodex { codex, config, .. } = test_codex() .with_config(|config| { @@ -315,8 +336,16 @@ async fn overrides_turn_context_but_keeps_cached_prefix_and_key_constant() -> an use pretty_assertions::assert_eq; let server = start_mock_server().await; - let req1 = mount_sse_once(&server, sse_completed("resp-1")).await; - let req2 = mount_sse_once(&server, sse_completed("resp-2")).await; + let req1 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; + let req2 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + ) + .await; let TestCodex { codex, .. } = test_codex() .with_config(|config| { @@ -407,7 +436,11 @@ async fn override_before_first_turn_emits_environment_context() -> anyhow::Resul skip_if_no_network!(Ok(())); let server = start_mock_server().await; - let req = mount_sse_once(&server, sse_completed("resp-1")).await; + let req = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; let TestCodex { codex, .. } = test_codex().build(&server).await?; @@ -542,8 +575,16 @@ async fn per_turn_overrides_keep_cached_prefix_and_key_constant() -> anyhow::Res use pretty_assertions::assert_eq; let server = start_mock_server().await; - let req1 = mount_sse_once(&server, sse_completed("resp-1")).await; - let req2 = mount_sse_once(&server, sse_completed("resp-2")).await; + let req1 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; + let req2 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + ) + .await; let TestCodex { codex, .. } = test_codex() .with_config(|config| { @@ -646,8 +687,16 @@ async fn send_user_turn_with_no_changes_does_not_send_environment_context() -> a use pretty_assertions::assert_eq; let server = start_mock_server().await; - let req1 = mount_sse_once(&server, sse_completed("resp-1")).await; - let req2 = mount_sse_once(&server, sse_completed("resp-2")).await; + let req1 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; + let req2 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + ) + .await; let TestCodex { codex, @@ -747,8 +796,16 @@ async fn send_user_turn_with_changes_sends_environment_context() -> anyhow::Resu let server = start_mock_server().await; - let req1 = mount_sse_once(&server, sse_completed("resp-1")).await; - let req2 = mount_sse_once(&server, sse_completed("resp-2")).await; + let req1 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; + let req2 = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + ) + .await; let TestCodex { codex, config, diff --git a/codex-rs/core/tests/suite/sqlite_state.rs b/codex-rs/core/tests/suite/sqlite_state.rs index ac71c14d36..07268f4fdb 100644 --- a/codex-rs/core/tests/suite/sqlite_state.rs +++ b/codex-rs/core/tests/suite/sqlite_state.rs @@ -9,7 +9,6 @@ use codex_protocol::protocol::SessionMeta; use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::UserMessageEvent; -use core_test_support::load_sse_fixture_with_id; use core_test_support::responses; use core_test_support::responses::ev_completed; use core_test_support::responses::ev_function_call; @@ -24,10 +23,6 @@ use tokio::time::Duration; use tracing_subscriber::prelude::*; use uuid::Uuid; -fn sse_completed(id: &str) -> String { - load_sse_fixture_with_id("../fixtures/completed_template.json", id) -} - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn new_thread_is_recorded_in_state_db() -> Result<()> { let server = start_mock_server().await; @@ -195,7 +190,10 @@ async fn user_messages_persist_in_state_db() -> Result<()> { let server = start_mock_server().await; mount_sse_sequence( &server, - vec![sse_completed("resp-1"), sse_completed("resp-2")], + vec![ + responses::sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + responses::sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + ], ) .await; 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 ec18816235..3562ec2500 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 @@ -3,7 +3,9 @@ use codex_core::WireApi; use codex_core::protocol::EventMsg; use codex_core::protocol::Op; use codex_protocol::user_input::UserInput; -use core_test_support::load_sse_fixture_with_id; +use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_response_created; +use core_test_support::responses::sse; use core_test_support::skip_if_no_network; use core_test_support::test_codex::TestCodex; use core_test_support::test_codex::test_codex; @@ -15,10 +17,6 @@ use wiremock::matchers::body_string_contains; use wiremock::matchers::method; use wiremock::matchers::path; -fn sse_completed(id: &str) -> String { - load_sse_fixture_with_id("../fixtures/completed_template.json", id) -} - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn continue_after_stream_error() { skip_if_no_network!(); @@ -46,7 +44,13 @@ async fn continue_after_stream_error() { let ok = ResponseTemplate::new(200) .insert_header("content-type", "text/event-stream") - .set_body_raw(sse_completed("resp_ok2"), "text/event-stream"); + .set_body_raw( + sse(vec![ + ev_response_created("resp_ok2"), + ev_completed("resp_ok2"), + ]), + "text/event-stream", + ); Mock::given(method("POST")) .and(path("/v1/responses")) diff --git a/codex-rs/core/tests/suite/stream_no_completed.rs b/codex-rs/core/tests/suite/stream_no_completed.rs index 68b7763eca..e8b130970e 100644 --- a/codex-rs/core/tests/suite/stream_no_completed.rs +++ b/codex-rs/core/tests/suite/stream_no_completed.rs @@ -7,7 +7,9 @@ use codex_core::protocol::EventMsg; use codex_core::protocol::Op; use codex_protocol::user_input::UserInput; use core_test_support::load_sse_fixture; -use core_test_support::load_sse_fixture_with_id; +use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_response_created; +use core_test_support::responses::sse; use core_test_support::skip_if_no_network; use core_test_support::test_codex::TestCodex; use core_test_support::test_codex::test_codex; @@ -24,10 +26,6 @@ fn sse_incomplete() -> String { load_sse_fixture("tests/fixtures/incomplete_sse.json") } -fn sse_completed(id: &str) -> String { - load_sse_fixture_with_id("../fixtures/completed_template.json", id) -} - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn retries_on_early_close() { skip_if_no_network!(); @@ -48,7 +46,13 @@ async fn retries_on_early_close() { } else { ResponseTemplate::new(200) .insert_header("content-type", "text/event-stream") - .set_body_raw(sse_completed("resp_ok"), "text/event-stream") + .set_body_raw( + sse(vec![ + ev_response_created("resp_ok"), + ev_completed("resp_ok"), + ]), + "text/event-stream", + ) } } } diff --git a/codex-rs/core/tests/suite/web_search.rs b/codex-rs/core/tests/suite/web_search.rs index 9397e88edb..8b1cc2a4e7 100644 --- a/codex-rs/core/tests/suite/web_search.rs +++ b/codex-rs/core/tests/suite/web_search.rs @@ -5,7 +5,6 @@ use codex_core::built_in_model_providers; use codex_core::features::Feature; use codex_core::protocol::SandboxPolicy; use codex_protocol::config_types::WebSearchMode; -use core_test_support::load_sse_fixture_with_id; use core_test_support::responses; use core_test_support::responses::start_mock_server; use core_test_support::skip_if_no_network; @@ -13,10 +12,6 @@ use core_test_support::test_codex::test_codex; use pretty_assertions::assert_eq; use serde_json::Value; -fn sse_completed(id: &str) -> String { - load_sse_fixture_with_id("../fixtures/completed_template.json", id) -} - #[allow(clippy::expect_used)] fn find_web_search_tool(body: &Value) -> &Value { body["tools"] @@ -41,7 +36,10 @@ async fn web_search_mode_cached_sets_external_web_access_false() { skip_if_no_network!(); let server = start_mock_server().await; - let sse = sse_completed("resp-1"); + let sse = responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_completed("resp-1"), + ]); let resp_mock = responses::mount_sse_once(&server, sse).await; let mut builder = test_codex() @@ -72,7 +70,10 @@ async fn web_search_mode_takes_precedence_over_legacy_flags() { skip_if_no_network!(); let server = start_mock_server().await; - let sse = sse_completed("resp-1"); + let sse = responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_completed("resp-1"), + ]); let resp_mock = responses::mount_sse_once(&server, sse).await; let mut builder = test_codex() @@ -104,7 +105,10 @@ async fn web_search_mode_defaults_to_cached_when_unset() { skip_if_no_network!(); let server = start_mock_server().await; - let sse = sse_completed("resp-1"); + let sse = responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_completed("resp-1"), + ]); let resp_mock = responses::mount_sse_once(&server, sse).await; let mut builder = test_codex() @@ -139,7 +143,16 @@ async fn web_search_mode_updates_between_turns_with_sandbox_policy() { let server = start_mock_server().await; let resp_mock = responses::mount_sse_sequence( &server, - vec![sse_completed("resp-1"), sse_completed("resp-2")], + vec![ + responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_completed("resp-1"), + ]), + responses::sse(vec![ + responses::ev_response_created("resp-2"), + responses::ev_completed("resp-2"), + ]), + ], ) .await; @@ -191,7 +204,10 @@ async fn web_search_mode_defaults_to_disabled_for_azure_responses() { skip_if_no_network!(); let server = start_mock_server().await; - let sse = sse_completed("resp-1"); + let sse = responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_completed("resp-1"), + ]); let resp_mock = responses::mount_sse_once(&server, sse).await; let mut builder = test_codex() From 49dd67a2606c73eed7316f882c8db027cbf71bbe Mon Sep 17 00:00:00 2001 From: jif-oai Date: Wed, 4 Feb 2026 16:39:41 +0000 Subject: [PATCH 4/8] feat: land unified_exec (#10641) Land `unified_exec` for all non-windows OS --- codex-rs/core/src/config/mod.rs | 8 +- codex-rs/core/src/features.rs | 16 +-- codex-rs/core/src/tools/spec.rs | 40 +++++-- codex-rs/core/tests/suite/model_tools.rs | 126 ++++++++++++-------- codex-rs/core/tests/suite/prompt_caching.rs | 10 +- 5 files changed, 118 insertions(+), 82 deletions(-) diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index d4547db19f..e22cfe65a1 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -3862,7 +3862,7 @@ model_verbosity = "high" forced_login_method: None, include_apply_patch_tool: false, web_search_mode: None, - use_experimental_unified_exec_tool: false, + use_experimental_unified_exec_tool: !cfg!(windows), ghost_snapshot: GhostSnapshotConfig::default(), features: Features::with_defaults(), suppress_unstable_features_warning: false, @@ -3947,7 +3947,7 @@ model_verbosity = "high" forced_login_method: None, include_apply_patch_tool: false, web_search_mode: None, - use_experimental_unified_exec_tool: false, + use_experimental_unified_exec_tool: !cfg!(windows), ghost_snapshot: GhostSnapshotConfig::default(), features: Features::with_defaults(), suppress_unstable_features_warning: false, @@ -4047,7 +4047,7 @@ model_verbosity = "high" forced_login_method: None, include_apply_patch_tool: false, web_search_mode: None, - use_experimental_unified_exec_tool: false, + use_experimental_unified_exec_tool: !cfg!(windows), ghost_snapshot: GhostSnapshotConfig::default(), features: Features::with_defaults(), suppress_unstable_features_warning: false, @@ -4133,7 +4133,7 @@ model_verbosity = "high" forced_login_method: None, include_apply_patch_tool: false, web_search_mode: None, - use_experimental_unified_exec_tool: false, + use_experimental_unified_exec_tool: !cfg!(windows), ghost_snapshot: GhostSnapshotConfig::default(), features: Features::with_defaults(), suppress_unstable_features_warning: false, diff --git a/codex-rs/core/src/features.rs b/codex-rs/core/src/features.rs index 1eb2a2dea5..a44477905e 100644 --- a/codex-rs/core/src/features.rs +++ b/codex-rs/core/src/features.rs @@ -406,6 +406,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::Stable, default_enabled: true, }, + FeatureSpec { + id: Feature::UnifiedExec, + key: "unified_exec", + stage: Stage::Stable, + default_enabled: !cfg!(windows), + }, FeatureSpec { id: Feature::WebSearchRequest, key: "web_search_request", @@ -419,16 +425,6 @@ pub const FEATURES: &[FeatureSpec] = &[ default_enabled: false, }, // Experimental program. Rendered in the `/experimental` menu for users. - FeatureSpec { - id: Feature::UnifiedExec, - key: "unified_exec", - stage: Stage::Experimental { - name: "Background terminal", - menu_description: "Run long-running terminal commands in the background.", - announcement: "NEW! Try Background terminals for long-running commands. Enable in /experimental!", - }, - default_enabled: false, - }, FeatureSpec { id: Feature::ShellSnapshot, key: "shell_snapshot", diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index 8851a157a2..ddac191a9b 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -1687,6 +1687,22 @@ mod tests { assert_eq!(&tool_names, &expected_tools,); } + fn assert_default_model_tools( + model_slug: &str, + features: &Features, + web_search_mode: Option, + shell_tool: &'static str, + expected_tail: &[&str], + ) { + let mut expected = if features.enabled(Feature::UnifiedExec) { + vec!["exec_command", "write_stdin"] + } else { + vec![shell_tool] + }; + expected.extend(expected_tail); + assert_model_tools(model_slug, features, web_search_mode, &expected); + } + #[test] fn web_search_mode_cached_sets_external_web_access_false() { let config = test_config(); @@ -1735,12 +1751,12 @@ mod tests { fn test_build_specs_gpt5_codex_default() { let mut features = Features::with_defaults(); features.enable(Feature::CollaborationModes); - assert_model_tools( + assert_default_model_tools( "gpt-5-codex", &features, Some(WebSearchMode::Cached), + "shell_command", &[ - "shell_command", "list_mcp_resources", "list_mcp_resource_templates", "read_mcp_resource", @@ -1757,12 +1773,12 @@ mod tests { fn test_build_specs_gpt51_codex_default() { let mut features = Features::with_defaults(); features.enable(Feature::CollaborationModes); - assert_model_tools( + assert_default_model_tools( "gpt-5.1-codex", &features, Some(WebSearchMode::Cached), + "shell_command", &[ - "shell_command", "list_mcp_resources", "list_mcp_resource_templates", "read_mcp_resource", @@ -1827,12 +1843,12 @@ mod tests { fn test_codex_mini_defaults() { let mut features = Features::with_defaults(); features.enable(Feature::CollaborationModes); - assert_model_tools( + assert_default_model_tools( "codex-mini-latest", &features, Some(WebSearchMode::Cached), + "local_shell", &[ - "local_shell", "list_mcp_resources", "list_mcp_resource_templates", "read_mcp_resource", @@ -1848,12 +1864,12 @@ mod tests { fn test_codex_5_1_mini_defaults() { let mut features = Features::with_defaults(); features.enable(Feature::CollaborationModes); - assert_model_tools( + assert_default_model_tools( "gpt-5.1-codex-mini", &features, Some(WebSearchMode::Cached), + "shell_command", &[ - "shell_command", "list_mcp_resources", "list_mcp_resource_templates", "read_mcp_resource", @@ -1870,12 +1886,12 @@ mod tests { fn test_gpt_5_defaults() { let mut features = Features::with_defaults(); features.enable(Feature::CollaborationModes); - assert_model_tools( + assert_default_model_tools( "gpt-5", &features, Some(WebSearchMode::Cached), + "shell", &[ - "shell", "list_mcp_resources", "list_mcp_resource_templates", "read_mcp_resource", @@ -1891,12 +1907,12 @@ mod tests { fn test_gpt_5_1_defaults() { let mut features = Features::with_defaults(); features.enable(Feature::CollaborationModes); - assert_model_tools( + assert_default_model_tools( "gpt-5.1", &features, Some(WebSearchMode::Cached), + "shell_command", &[ - "shell_command", "list_mcp_resources", "list_mcp_resource_templates", "read_mcp_resource", diff --git a/codex-rs/core/tests/suite/model_tools.rs b/codex-rs/core/tests/suite/model_tools.rs index c687545f26..d74ab0d7be 100644 --- a/codex-rs/core/tests/suite/model_tools.rs +++ b/codex-rs/core/tests/suite/model_tools.rs @@ -50,6 +50,16 @@ async fn collect_tool_identifiers_for_model(model: &str) -> Vec { tool_identifiers(&body) } +fn expected_default_tools(shell_tool: &str, tail: &[&str]) -> Vec { + let mut tools = if cfg!(windows) { + vec![shell_tool.to_string()] + } else { + vec!["exec_command".to_string(), "write_stdin".to_string()] + }; + tools.extend(tail.iter().map(|tool| (*tool).to_string())); + tools +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn model_selects_expected_tools() { skip_if_no_network!(); @@ -58,83 +68,93 @@ async fn model_selects_expected_tools() { let codex_tools = collect_tool_identifiers_for_model("codex-mini-latest").await; assert_eq!( codex_tools, - vec![ - "local_shell".to_string(), - "list_mcp_resources".to_string(), - "list_mcp_resource_templates".to_string(), - "read_mcp_resource".to_string(), - "update_plan".to_string(), - "request_user_input".to_string(), - "web_search".to_string(), - "view_image".to_string() - ], + expected_default_tools( + "local_shell", + &[ + "list_mcp_resources", + "list_mcp_resource_templates", + "read_mcp_resource", + "update_plan", + "request_user_input", + "web_search", + "view_image", + ], + ), "codex-mini-latest should expose the local shell tool", ); let gpt5_codex_tools = collect_tool_identifiers_for_model("gpt-5-codex").await; assert_eq!( gpt5_codex_tools, - vec![ - "shell_command".to_string(), - "list_mcp_resources".to_string(), - "list_mcp_resource_templates".to_string(), - "read_mcp_resource".to_string(), - "update_plan".to_string(), - "request_user_input".to_string(), - "apply_patch".to_string(), - "web_search".to_string(), - "view_image".to_string() - ], + expected_default_tools( + "shell_command", + &[ + "list_mcp_resources", + "list_mcp_resource_templates", + "read_mcp_resource", + "update_plan", + "request_user_input", + "apply_patch", + "web_search", + "view_image", + ], + ), "gpt-5-codex should expose the apply_patch tool", ); let gpt51_codex_tools = collect_tool_identifiers_for_model("gpt-5.1-codex").await; assert_eq!( gpt51_codex_tools, - vec![ - "shell_command".to_string(), - "list_mcp_resources".to_string(), - "list_mcp_resource_templates".to_string(), - "read_mcp_resource".to_string(), - "update_plan".to_string(), - "request_user_input".to_string(), - "apply_patch".to_string(), - "web_search".to_string(), - "view_image".to_string() - ], + expected_default_tools( + "shell_command", + &[ + "list_mcp_resources", + "list_mcp_resource_templates", + "read_mcp_resource", + "update_plan", + "request_user_input", + "apply_patch", + "web_search", + "view_image", + ], + ), "gpt-5.1-codex should expose the apply_patch tool", ); let gpt5_tools = collect_tool_identifiers_for_model("gpt-5").await; assert_eq!( gpt5_tools, - vec![ - "shell".to_string(), - "list_mcp_resources".to_string(), - "list_mcp_resource_templates".to_string(), - "read_mcp_resource".to_string(), - "update_plan".to_string(), - "request_user_input".to_string(), - "web_search".to_string(), - "view_image".to_string() - ], + expected_default_tools( + "shell", + &[ + "list_mcp_resources", + "list_mcp_resource_templates", + "read_mcp_resource", + "update_plan", + "request_user_input", + "web_search", + "view_image", + ], + ), "gpt-5 should expose the apply_patch tool", ); let gpt51_tools = collect_tool_identifiers_for_model("gpt-5.1").await; assert_eq!( gpt51_tools, - vec![ - "shell_command".to_string(), - "list_mcp_resources".to_string(), - "list_mcp_resource_templates".to_string(), - "read_mcp_resource".to_string(), - "update_plan".to_string(), - "request_user_input".to_string(), - "apply_patch".to_string(), - "web_search".to_string(), - "view_image".to_string() - ], + expected_default_tools( + "shell_command", + &[ + "list_mcp_resources", + "list_mcp_resource_templates", + "read_mcp_resource", + "update_plan", + "request_user_input", + "apply_patch", + "web_search", + "view_image", + ], + ), "gpt-5.1 should expose the apply_patch tool", ); let exp_tools = collect_tool_identifiers_for_model("exp-5.1").await; diff --git a/codex-rs/core/tests/suite/prompt_caching.rs b/codex-rs/core/tests/suite/prompt_caching.rs index 810cd52497..acff97381a 100644 --- a/codex-rs/core/tests/suite/prompt_caching.rs +++ b/codex-rs/core/tests/suite/prompt_caching.rs @@ -136,8 +136,12 @@ async fn prompt_tools_are_consistent_across_requests() -> anyhow::Result<()> { .await?; wait_for_event(&codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await; - let expected_tools_names = vec![ - "shell_command", + let mut expected_tools_names = if cfg!(windows) { + vec!["shell_command"] + } else { + vec!["exec_command", "write_stdin"] + }; + expected_tools_names.extend([ "list_mcp_resources", "list_mcp_resource_templates", "read_mcp_resource", @@ -146,7 +150,7 @@ async fn prompt_tools_are_consistent_across_requests() -> anyhow::Result<()> { "apply_patch", "web_search", "view_image", - ]; + ]); let body0 = req1.single_request().body_json(); let expected_instructions = if expected_tools_names.contains(&"apply_patch") { From 7f203576116de4002fcda5942a1e4248e18692ff Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Wed, 4 Feb 2026 09:05:37 -0800 Subject: [PATCH 5/8] Stop client from being state carrier (#10595) I'd like to make client session wide. This requires shedding all random state it has to carry. --- codex-rs/core/src/client.rs | 55 -------- codex-rs/core/src/codex.rs | 133 ++++++++++++------- codex-rs/core/src/compact.rs | 10 +- codex-rs/core/src/compact_remote.rs | 4 +- codex-rs/core/src/context_manager/history.rs | 6 +- codex-rs/core/src/mcp/skill_dependencies.rs | 2 +- codex-rs/core/src/mcp_tool_call.rs | 6 +- codex-rs/core/src/stream_events_utils.rs | 3 +- codex-rs/core/src/tasks/compact.rs | 5 +- codex-rs/core/src/tasks/mod.rs | 3 +- codex-rs/core/src/tasks/review.rs | 4 +- codex-rs/core/src/tasks/user_shell.rs | 2 +- codex-rs/core/src/tools/handlers/collab.rs | 66 +++------ codex-rs/core/src/tools/orchestrator.rs | 2 +- codex-rs/core/src/tools/registry.rs | 2 +- 15 files changed, 127 insertions(+), 176 deletions(-) diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 1d0dc7fb65..dd92c5da01 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -153,61 +153,6 @@ impl ModelClient { } impl ModelClient { - pub fn get_model_context_window(&self) -> Option { - let model_info = &self.state.model_info; - let effective_context_window_percent = model_info.effective_context_window_percent; - model_info.context_window.map(|context_window| { - context_window.saturating_mul(effective_context_window_percent) / 100 - }) - } - - pub fn config(&self) -> Arc { - Arc::clone(&self.state.config) - } - - pub fn provider(&self) -> &ModelProviderInfo { - &self.state.provider - } - - pub fn get_provider(&self) -> ModelProviderInfo { - self.state.provider.clone() - } - - pub fn get_otel_manager(&self) -> OtelManager { - self.state.otel_manager.clone() - } - - pub fn get_session_source(&self) -> SessionSource { - self.state.session_source.clone() - } - - pub(crate) fn transport_manager(&self) -> TransportManager { - self.state.transport_manager.clone() - } - - /// Returns the currently configured model slug. - pub fn get_model(&self) -> String { - self.state.model_info.slug.clone() - } - - pub fn get_model_info(&self) -> ModelInfo { - self.state.model_info.clone() - } - - /// Returns the current reasoning effort setting. - pub fn get_reasoning_effort(&self) -> Option { - self.state.effort - } - - /// Returns the current reasoning summary setting. - pub fn get_reasoning_summary(&self) -> ReasoningSummaryConfig { - self.state.summary - } - - pub fn get_auth_manager(&self) -> Option> { - self.state.auth_manager.clone() - } - /// Compacts the current conversation history using the Compact endpoint. /// /// This is a unary call (no streaming) that returns a new list of diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 7fe3b8c810..979f2a5f62 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -212,6 +212,7 @@ use codex_protocol::models::ContentItem; use codex_protocol::models::DeveloperInstructions; use codex_protocol::models::ResponseInputItem; use codex_protocol::models::ResponseItem; +use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig; use codex_protocol::protocol::CodexErrorInfo; use codex_protocol::protocol::InitialHistory; use codex_protocol::user_input::UserInput; @@ -484,6 +485,15 @@ pub(crate) struct Session { pub(crate) struct TurnContext { pub(crate) sub_id: String, pub(crate) client: ModelClient, + pub(crate) config: Arc, + pub(crate) auth_manager: Option>, + pub(crate) model_info: ModelInfo, + pub(crate) otel_manager: OtelManager, + pub(crate) provider: ModelProviderInfo, + pub(crate) reasoning_effort: Option, + pub(crate) reasoning_summary: ReasoningSummaryConfig, + pub(crate) session_source: SessionSource, + pub(crate) transport_manager: TransportManager, /// The session's current working directory. All relative paths provided by /// the model as well as sandbox policies are resolved against this path /// instead of `std::env::current_dir()`. @@ -507,6 +517,13 @@ pub(crate) struct TurnContext { turn_metadata_header: OnceCell>, } impl TurnContext { + pub(crate) fn model_context_window(&self) -> Option { + let effective_context_window_percent = self.model_info.effective_context_window_percent; + self.model_info.context_window.map(|context_window| { + context_window.saturating_mul(effective_context_window_percent) / 100 + }) + } + pub(crate) fn resolve_path(&self, path: Option) -> PathBuf { path.as_ref() .map(PathBuf::from) @@ -694,10 +711,17 @@ impl Session { sub_id: String, transport_manager: TransportManager, ) -> TurnContext { + let reasoning_effort = session_configuration.collaboration_mode.reasoning_effort(); + let reasoning_summary = session_configuration.model_reasoning_summary; let otel_manager = otel_manager.clone().with_model( session_configuration.collaboration_mode.model(), model_info.slug.as_str(), ); + let session_source = session_configuration.session_source.clone(); + let auth_manager_for_context = auth_manager.clone(); + let provider_for_context = provider.clone(); + let transport_manager_for_context = transport_manager.clone(); + let otel_manager_for_context = otel_manager.clone(); let per_turn_config = Arc::new(per_turn_config); let client = ModelClient::new( per_turn_config.clone(), @@ -705,10 +729,10 @@ impl Session { model_info.clone(), otel_manager, provider, - session_configuration.collaboration_mode.reasoning_effort(), - session_configuration.model_reasoning_summary, + reasoning_effort, + reasoning_summary, conversation_id, - session_configuration.session_source.clone(), + session_source.clone(), transport_manager, ); @@ -722,6 +746,15 @@ impl Session { TurnContext { sub_id, client, + config: per_turn_config.clone(), + auth_manager: auth_manager_for_context, + model_info: model_info.clone(), + otel_manager: otel_manager_for_context, + provider: provider_for_context, + reasoning_effort, + reasoning_summary, + session_source, + transport_manager: transport_manager_for_context, cwd, developer_instructions: session_configuration.developer_instructions.clone(), compact_prompt: session_configuration.compact_prompt.clone(), @@ -1100,7 +1133,7 @@ impl Session { None } }) { - let curr = turn_context.client.get_model(); + let curr = turn_context.model_info.slug.as_str(); if prev != curr { warn!( "resuming session with different model: previous={prev}, current={curr}" @@ -1378,8 +1411,8 @@ impl Session { if let Some(personality) = next.personality && next.personality != previous.personality { - let model_info = next.client.get_model_info(); - let personality_message = Self::personality_message_for(&model_info, personality); + let model_info = &next.model_info; + let personality_message = Self::personality_message_for(model_info, personality); personality_message.map(|personality_message| { DeveloperInstructions::personality_spec_message(personality_message).into() }) @@ -1943,7 +1976,7 @@ impl Session { if self.features.enabled(Feature::Personality) && let Some(personality) = turn_context.personality { - let model_info = turn_context.client.get_model_info(); + let model_info = turn_context.model_info.clone(); let has_baked_personality = model_info.supports_personality() && base_instructions == model_info.get_model_instructions(Some(personality)); if !has_baked_personality @@ -1996,10 +2029,8 @@ impl Session { { let mut state = self.state.lock().await; if let Some(token_usage) = token_usage { - state.update_token_info_from_usage( - token_usage, - turn_context.client.get_model_context_window(), - ); + state + .update_token_info_from_usage(token_usage, turn_context.model_context_window()); } } self.send_token_count_event(turn_context).await; @@ -2030,7 +2061,7 @@ impl Session { }; if info.model_context_window.is_none() { - info.model_context_window = turn_context.client.get_model_context_window(); + info.model_context_window = turn_context.model_context_window(); } state.set_token_info(Some(info)); @@ -2088,7 +2119,7 @@ impl Session { } pub(crate) async fn set_total_tokens_full(&self, turn_context: &TurnContext) { - if let Some(context_window) = turn_context.client.get_model_context_window() { + if let Some(context_window) = turn_context.model_context_window() { let mut state = self.state.lock().await; state.set_token_usage_full(context_window); } @@ -2702,10 +2733,7 @@ mod handlers { // new_turn_with_sub_id already emits the error event. return; }; - current_context - .client - .get_otel_manager() - .user_prompt(&items); + current_context.otel_manager.user_prompt(&items); // Attempt to inject input into current task if let Err(items) = sess.inject_input(items).await { @@ -3247,7 +3275,7 @@ async fn spawn_review_thread( let model = config .review_model .clone() - .unwrap_or_else(|| parent_turn_context.client.get_model()); + .unwrap_or_else(|| parent_turn_context.model_info.slug.clone()); let review_model_info = sess .services .models_manager @@ -3266,8 +3294,8 @@ async fn spawn_review_thread( }); let review_prompt = resolved.prompt.clone(); - let provider = parent_turn_context.client.get_provider(); - let auth_manager = parent_turn_context.client.get_auth_manager(); + let provider = parent_turn_context.provider.clone(); + let auth_manager = parent_turn_context.auth_manager.clone(); let model_info = review_model_info.clone(); // Build per‑turn client with the requested model/family. @@ -3277,9 +3305,16 @@ async fn spawn_review_thread( per_turn_config.web_search_mode = Some(review_web_search_mode); let otel_manager = parent_turn_context - .client - .get_otel_manager() + .otel_manager + .clone() .with_model(model.as_str(), review_model_info.slug.as_str()); + let auth_manager_for_context = auth_manager.clone(); + let provider_for_context = provider.clone(); + let otel_manager_for_context = otel_manager.clone(); + let reasoning_effort = per_turn_config.model_reasoning_effort; + let reasoning_summary = per_turn_config.model_reasoning_summary; + let session_source = parent_turn_context.session_source.clone(); + let transport_manager = parent_turn_context.transport_manager.clone(); let per_turn_config = Arc::new(per_turn_config); let client = ModelClient::new( @@ -3288,16 +3323,25 @@ async fn spawn_review_thread( model_info.clone(), otel_manager, provider, - per_turn_config.model_reasoning_effort, - per_turn_config.model_reasoning_summary, + reasoning_effort, + reasoning_summary, sess.conversation_id, - parent_turn_context.client.get_session_source(), - parent_turn_context.client.transport_manager(), + session_source.clone(), + transport_manager.clone(), ); let review_turn_context = TurnContext { sub_id: sub_id.to_string(), client, + config: per_turn_config, + auth_manager: auth_manager_for_context, + model_info: model_info.clone(), + otel_manager: otel_manager_for_context, + provider: provider_for_context, + reasoning_effort, + reasoning_summary, + session_source, + transport_manager, tools_config, ghost_snapshot: parent_turn_context.ghost_snapshot.clone(), developer_instructions: None, @@ -3414,12 +3458,12 @@ pub(crate) async fn run_turn( return None; } - let model_info = turn_context.client.get_model_info(); + let model_info = turn_context.model_info.clone(); let auto_compact_limit = model_info.auto_compact_token_limit().unwrap_or(i64::MAX); let total_usage_tokens = sess.get_total_token_usage().await; let event = EventMsg::TurnStarted(TurnStartedEvent { - model_context_window: turn_context.client.get_model_context_window(), + model_context_window: turn_context.model_context_window(), collaboration_mode_kind: turn_context.collaboration_mode.mode, }); sess.send_event(&turn_context, event).await; @@ -3438,7 +3482,7 @@ pub(crate) async fn run_turn( || (HashMap::new(), HashMap::new()), |outcome| build_skill_name_counts(&outcome.skills, &outcome.disabled_paths), ); - let connector_slug_counts = if turn_context.client.config().features.enabled(Feature::Apps) { + let connector_slug_counts = if turn_context.config.features.enabled(Feature::Apps) { let mcp_tools = match sess .services .mcp_connection_manager @@ -3467,7 +3511,7 @@ pub(crate) async fn run_turn( }); let explicit_app_paths = collect_explicit_app_paths(&input); - let config = turn_context.client.config(); + let config = turn_context.config.clone(); if config .features .enabled(Feature::SkillEnvVarDependencyPrompt) @@ -3484,9 +3528,9 @@ pub(crate) async fn run_turn( ) .await; - let otel_manager = turn_context.client.get_otel_manager(); + let otel_manager = turn_context.otel_manager.clone(); let thread_id = sess.conversation_id.to_string(); - let tracking = build_track_events_context(turn_context.client.get_model(), thread_id); + let tracking = build_track_events_context(turn_context.model_info.slug.clone(), thread_id); let SkillInjections { items: skill_items, warnings: skill_warnings, @@ -3639,7 +3683,7 @@ pub(crate) async fn run_turn( } async fn run_auto_compact(sess: &Arc, turn_context: &Arc) { - if should_use_remote_compact_task(sess.as_ref(), &turn_context.client.get_provider()) { + if should_use_remote_compact_task(sess.as_ref(), &turn_context.provider) { run_inline_remote_auto_compact_task(Arc::clone(sess), Arc::clone(turn_context)).await; } else { run_inline_auto_compact_task(Arc::clone(sess), Arc::clone(turn_context)).await; @@ -3748,7 +3792,7 @@ struct SamplingRequestToolSelection<'a> { skip_all, fields( turn_id = %turn_context.sub_id, - model = %turn_context.client.get_model(), + model = %turn_context.model_info.slug, cwd = %turn_context.cwd.display() ) )] @@ -3769,7 +3813,7 @@ async fn run_sampling_request( .list_all_tools() .or_cancel(&cancellation_token) .await?; - let connectors_for_tools = if turn_context.client.config().features.enabled(Feature::Apps) { + let connectors_for_tools = if turn_context.config.features.enabled(Feature::Apps) { let connectors = connectors::accessible_connectors_from_mcp_tools(&mcp_tools); Some(filter_connectors_for_input( connectors, @@ -3794,10 +3838,7 @@ async fn run_sampling_request( turn_context.dynamic_tools.as_slice(), )); - let model_supports_parallel = turn_context - .client - .get_model_info() - .supports_parallel_tool_calls; + let model_supports_parallel = turn_context.model_info.supports_parallel_tool_calls; let base_instructions = sess.get_base_instructions().await; @@ -3845,7 +3886,7 @@ async fn run_sampling_request( } // Use the configured provider-specific stream retry budget. - let max_retries = turn_context.client.get_provider().stream_max_retries(); + let max_retries = turn_context.provider.stream_max_retries(); if retries >= max_retries && client_session.try_switch_fallback_transport() { sess.send_event( &turn_context, @@ -4291,7 +4332,7 @@ async fn drain_in_flight( skip_all, fields( turn_id = %turn_context.sub_id, - model = %turn_context.client.get_model() + model = %turn_context.model_info.slug ) )] async fn try_run_sampling_request( @@ -4308,11 +4349,11 @@ async fn try_run_sampling_request( cwd: turn_context.cwd.clone(), approval_policy: turn_context.approval_policy, sandbox_policy: turn_context.sandbox_policy.clone(), - model: turn_context.client.get_model(), + model: turn_context.model_info.slug.clone(), personality: turn_context.personality, collaboration_mode: Some(collaboration_mode), - effort: turn_context.client.get_reasoning_effort(), - summary: turn_context.client.get_reasoning_summary(), + effort: turn_context.reasoning_effort, + summary: turn_context.reasoning_summary, user_instructions: turn_context.user_instructions.clone(), developer_instructions: turn_context.developer_instructions.clone(), final_output_json_schema: turn_context.final_output_json_schema.clone(), @@ -4320,10 +4361,10 @@ async fn try_run_sampling_request( }); feedback_tags!( - model = turn_context.client.get_model(), + model = turn_context.model_info.slug.clone(), approval_policy = turn_context.approval_policy, sandbox_policy = turn_context.sandbox_policy, - effort = turn_context.client.get_reasoning_effort(), + effort = turn_context.reasoning_effort, auth_mode = sess.services.auth_manager.get_auth_mode(), features = sess.features.enabled_features(), ); diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index 7ae773884c..c8aba56613 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -60,7 +60,7 @@ pub(crate) async fn run_compact_task( input: Vec, ) { let start_event = EventMsg::TurnStarted(TurnStartedEvent { - model_context_window: turn_context.client.get_model_context_window(), + model_context_window: turn_context.model_context_window(), collaboration_mode_kind: turn_context.collaboration_mode.mode, }); sess.send_event(&turn_context, start_event).await; @@ -85,7 +85,7 @@ async fn run_compact_task_inner( let mut truncated_count = 0usize; - let max_retries = turn_context.client.get_provider().stream_max_retries(); + let max_retries = turn_context.provider.stream_max_retries(); let mut retries = 0; // TODO: If we need to guarantee the persisted mode always matches the prompt used for this @@ -97,11 +97,11 @@ async fn run_compact_task_inner( cwd: turn_context.cwd.clone(), approval_policy: turn_context.approval_policy, sandbox_policy: turn_context.sandbox_policy.clone(), - model: turn_context.client.get_model(), + model: turn_context.model_info.slug.clone(), personality: turn_context.personality, collaboration_mode: Some(collaboration_mode), - effort: turn_context.client.get_reasoning_effort(), - summary: turn_context.client.get_reasoning_summary(), + effort: turn_context.reasoning_effort, + summary: turn_context.reasoning_summary, user_instructions: turn_context.user_instructions.clone(), developer_instructions: turn_context.developer_instructions.clone(), final_output_json_schema: turn_context.final_output_json_schema.clone(), diff --git a/codex-rs/core/src/compact_remote.rs b/codex-rs/core/src/compact_remote.rs index 12bc769ce2..fbb6863c37 100644 --- a/codex-rs/core/src/compact_remote.rs +++ b/codex-rs/core/src/compact_remote.rs @@ -24,7 +24,7 @@ pub(crate) async fn run_inline_remote_auto_compact_task( pub(crate) async fn run_remote_compact_task(sess: Arc, turn_context: Arc) { let start_event = EventMsg::TurnStarted(TurnStartedEvent { - model_context_window: turn_context.client.get_model_context_window(), + model_context_window: turn_context.model_context_window(), collaboration_mode_kind: turn_context.collaboration_mode.mode, }); sess.send_event(&turn_context, start_event).await; @@ -104,7 +104,7 @@ fn trim_function_call_history_to_fit_context_window( turn_context: &TurnContext, ) -> usize { let mut deleted_items = 0usize; - let Some(context_window) = turn_context.client.get_model_context_window() else { + let Some(context_window) = turn_context.model_context_window() else { return deleted_items; }; diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index a29f7df7e0..a68c4daafe 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -85,10 +85,8 @@ impl ContextManager { // Estimate token usage using byte-based heuristics from the truncation helpers. // This is a coarse lower bound, not a tokenizer-accurate count. pub(crate) fn estimate_token_count(&self, turn_context: &TurnContext) -> Option { - let model_info = turn_context.client.get_model_info(); - let personality = turn_context - .personality - .or(turn_context.client.config().personality); + let model_info = &turn_context.model_info; + let personality = turn_context.personality.or(turn_context.config.personality); let base_instructions = model_info.get_model_instructions(personality); let base_tokens = i64::try_from(approx_token_count(&base_instructions)).unwrap_or(i64::MAX); diff --git a/codex-rs/core/src/mcp/skill_dependencies.rs b/codex-rs/core/src/mcp/skill_dependencies.rs index 37620d73db..14b71fcf19 100644 --- a/codex-rs/core/src/mcp/skill_dependencies.rs +++ b/codex-rs/core/src/mcp/skill_dependencies.rs @@ -144,7 +144,7 @@ pub(crate) async fn maybe_prompt_and_install_mcp_dependencies( return; } - let config = turn_context.client.config(); + let config = turn_context.config.clone(); if mentioned_skills.is_empty() || !config.features.enabled(Feature::SkillMcpDependencyInstall) { return; } diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index 75248f34cc..d1d009470d 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -120,8 +120,7 @@ pub(crate) async fn handle_mcp_tool_call( let status = if result.is_ok() { "ok" } else { "error" }; turn_context - .client - .get_otel_manager() + .otel_manager .counter("codex.mcp.call", 1, &[("status", status)]); return ResponseInputItem::McpToolCallOutput { call_id, result }; @@ -153,8 +152,7 @@ pub(crate) async fn handle_mcp_tool_call( let status = if result.is_ok() { "ok" } else { "error" }; turn_context - .client - .get_otel_manager() + .otel_manager .counter("codex.mcp.call", 1, &[("status", status)]); ResponseInputItem::McpToolCallOutput { call_id, result } diff --git a/codex-rs/core/src/stream_events_utils.rs b/codex-rs/core/src/stream_events_utils.rs index 02d9822510..31b4f18372 100644 --- a/codex-rs/core/src/stream_events_utils.rs +++ b/codex-rs/core/src/stream_events_utils.rs @@ -100,8 +100,7 @@ pub(crate) async fn handle_output_item_done( Err(FunctionCallError::MissingLocalShellCallId) => { let msg = "LocalShellCall without call_id or id"; ctx.turn_context - .client - .get_otel_manager() + .otel_manager .log_tool_failed("local_shell", msg); tracing::error!(msg); diff --git a/codex-rs/core/src/tasks/compact.rs b/codex-rs/core/src/tasks/compact.rs index 00d882f607..e1fdbfd91f 100644 --- a/codex-rs/core/src/tasks/compact.rs +++ b/codex-rs/core/src/tasks/compact.rs @@ -25,10 +25,7 @@ impl SessionTask for CompactTask { _cancellation_token: CancellationToken, ) -> Option { let session = session.clone_session(); - if crate::compact::should_use_remote_compact_task( - session.as_ref(), - &ctx.client.get_provider(), - ) { + if crate::compact::should_use_remote_compact_task(session.as_ref(), &ctx.provider) { let _ = session.services.otel_manager.counter( "codex.task.compact", 1, diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index d5d6f60582..0ab6e2f49b 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -158,8 +158,7 @@ impl Session { }; let timer = turn_context - .client - .get_otel_manager() + .otel_manager .start_timer("codex.turn.e2e_duration_ms", &[]) .ok(); diff --git a/codex-rs/core/src/tasks/review.rs b/codex-rs/core/src/tasks/review.rs index cfc554f856..87d6c41933 100644 --- a/codex-rs/core/src/tasks/review.rs +++ b/codex-rs/core/src/tasks/review.rs @@ -82,7 +82,7 @@ async fn start_review_conversation( input: Vec, cancellation_token: CancellationToken, ) -> Option> { - let config = ctx.client.config(); + let config = ctx.config.clone(); let mut sub_agent_config = config.as_ref().clone(); // Carry over review-only feature restrictions so the delegate cannot // re-enable blocked tools (web search, view image). @@ -94,7 +94,7 @@ async fn start_review_conversation( let model = config .review_model .clone() - .unwrap_or_else(|| ctx.client.get_model()); + .unwrap_or_else(|| ctx.model_info.slug.clone()); sub_agent_config.model = Some(model); (run_codex_thread_one_shot( sub_agent_config, diff --git a/codex-rs/core/src/tasks/user_shell.rs b/codex-rs/core/src/tasks/user_shell.rs index a3d38afd2c..6c37cb150e 100644 --- a/codex-rs/core/src/tasks/user_shell.rs +++ b/codex-rs/core/src/tasks/user_shell.rs @@ -66,7 +66,7 @@ impl SessionTask for UserShellCommandTask { .counter("codex.task.user_shell", 1, &[]); let event = EventMsg::TurnStarted(TurnStartedEvent { - model_context_window: turn_context.client.get_model_context_window(), + model_context_window: turn_context.model_context_window(), collaboration_mode_kind: turn_context.collaboration_mode.mode, }); let session = session.clone_session(); diff --git a/codex-rs/core/src/tools/handlers/collab.rs b/codex-rs/core/src/tools/handlers/collab.rs index 61ccc1932e..7ef19506d5 100644 --- a/codex-rs/core/src/tools/handlers/collab.rs +++ b/codex-rs/core/src/tools/handlers/collab.rs @@ -114,7 +114,7 @@ mod spawn { "Empty message can't be sent to an agent".to_string(), )); } - let session_source = turn.client.get_session_source(); + let session_source = turn.session_source.clone(); let child_depth = next_thread_spawn_depth(&session_source); if exceeds_thread_spawn_depth_limit(child_depth) { return Err(FunctionCallError::RespondToModel( @@ -593,13 +593,13 @@ fn build_agent_spawn_config( turn: &TurnContext, child_depth: i32, ) -> Result { - let base_config = turn.client.config(); + let base_config = turn.config.clone(); let mut config = (*base_config).clone(); config.base_instructions = Some(base_instructions.text.clone()); - config.model = Some(turn.client.get_model()); - config.model_provider = turn.client.get_provider(); - config.model_reasoning_effort = turn.client.get_reasoning_effort(); - config.model_reasoning_summary = turn.client.get_reasoning_summary(); + config.model = Some(turn.model_info.slug.clone()); + config.model_provider = turn.provider.clone(); + config.model_reasoning_effort = turn.reasoning_effort; + config.model_reasoning_summary = turn.reasoning_summary; config.developer_instructions = turn.developer_instructions.clone(); config.compact_prompt = turn.compact_prompt.clone(); config.shell_environment_policy = turn.shell_environment_policy.clone(); @@ -633,7 +633,6 @@ mod tests { use crate::ThreadManager; use crate::agent::MAX_THREAD_SPAWN_DEPTH; use crate::built_in_model_providers; - use crate::client::ModelClient; use crate::codex::make_session_and_context; use crate::config::types::ShellEnvironmentPolicy; use crate::function_tool::FunctionCallError; @@ -767,22 +766,10 @@ mod tests { let manager = thread_manager(); session.services.agent_control = manager.agent_control(); - let session_source = SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + turn.session_source = SessionSource::SubAgent(SubAgentSource::ThreadSpawn { parent_thread_id: session.conversation_id, depth: MAX_THREAD_SPAWN_DEPTH, }); - turn.client = ModelClient::new( - turn.client.config(), - Some(session.services.auth_manager.clone()), - turn.client.get_model_info(), - turn.client.get_otel_manager(), - turn.client.get_provider(), - turn.client.get_reasoning_effort(), - turn.client.get_reasoning_summary(), - session.conversation_id, - session_source, - session.services.transport_manager.clone(), - ); let invocation = invocation( Arc::new(session), @@ -865,7 +852,7 @@ mod tests { let (mut session, turn) = make_session_and_context().await; let manager = thread_manager(); session.services.agent_control = manager.agent_control(); - let config = turn.client.config().as_ref().clone(); + let config = turn.config.as_ref().clone(); let thread = manager.start_thread(config).await.expect("start thread"); let agent_id = thread.thread_id; let invocation = invocation( @@ -1008,7 +995,7 @@ mod tests { let (mut session, turn) = make_session_and_context().await; let manager = thread_manager(); session.services.agent_control = manager.agent_control(); - let config = turn.client.config().as_ref().clone(); + let config = turn.config.as_ref().clone(); let thread = manager.start_thread(config).await.expect("start thread"); let agent_id = thread.thread_id; let invocation = invocation( @@ -1053,7 +1040,7 @@ mod tests { let (mut session, turn) = make_session_and_context().await; let manager = thread_manager(); session.services.agent_control = manager.agent_control(); - let config = turn.client.config().as_ref().clone(); + let config = turn.config.as_ref().clone(); let thread = manager.start_thread(config).await.expect("start thread"); let agent_id = thread.thread_id; let invocation = invocation( @@ -1084,7 +1071,7 @@ mod tests { let (mut session, turn) = make_session_and_context().await; let manager = thread_manager(); session.services.agent_control = manager.agent_control(); - let config = turn.client.config().as_ref().clone(); + let config = turn.config.as_ref().clone(); let thread = manager.start_thread(config).await.expect("start thread"); let agent_id = thread.thread_id; let mut status_rx = manager @@ -1138,7 +1125,7 @@ mod tests { let (mut session, turn) = make_session_and_context().await; let manager = thread_manager(); session.services.agent_control = manager.agent_control(); - let config = turn.client.config().as_ref().clone(); + let config = turn.config.as_ref().clone(); let thread = manager.start_thread(config).await.expect("start thread"); let agent_id = thread.thread_id; let status_before = manager.agent_control().get_status(agent_id).await; @@ -1193,12 +1180,12 @@ mod tests { turn.sandbox_policy = SandboxPolicy::DangerFullAccess; let config = build_agent_spawn_config(&base_instructions, &turn, 0).expect("spawn config"); - let mut expected = (*turn.client.config()).clone(); + let mut expected = (*turn.config).clone(); expected.base_instructions = Some(base_instructions.text); - expected.model = Some(turn.client.get_model()); - expected.model_provider = turn.client.get_provider(); - expected.model_reasoning_effort = turn.client.get_reasoning_effort(); - expected.model_reasoning_summary = turn.client.get_reasoning_summary(); + expected.model = Some(turn.model_info.slug.clone()); + expected.model_provider = turn.provider.clone(); + expected.model_reasoning_effort = turn.reasoning_effort; + expected.model_reasoning_summary = turn.reasoning_summary; expected.developer_instructions = turn.developer_instructions.clone(); expected.compact_prompt = turn.compact_prompt.clone(); expected.shell_environment_policy = turn.shell_environment_policy.clone(); @@ -1217,24 +1204,11 @@ mod tests { #[tokio::test] async fn build_agent_spawn_config_preserves_base_user_instructions() { - let (session, mut turn) = make_session_and_context().await; - let session_source = turn.client.get_session_source(); - let mut base_config = (*turn.client.config()).clone(); + let (_session, mut turn) = make_session_and_context().await; + let mut base_config = (*turn.config).clone(); base_config.user_instructions = Some("base-user".to_string()); turn.user_instructions = Some("resolved-user".to_string()); - let transport_manager = turn.client.transport_manager(); - turn.client = ModelClient::new( - Arc::new(base_config.clone()), - Some(session.services.auth_manager.clone()), - turn.client.get_model_info(), - turn.client.get_otel_manager(), - turn.client.get_provider(), - turn.client.get_reasoning_effort(), - turn.client.get_reasoning_summary(), - session.conversation_id, - session_source, - transport_manager, - ); + turn.config = Arc::new(base_config.clone()); let base_instructions = BaseInstructions { text: "base".to_string(), }; diff --git a/codex-rs/core/src/tools/orchestrator.rs b/codex-rs/core/src/tools/orchestrator.rs index e9fdd6208b..fb2150953b 100644 --- a/codex-rs/core/src/tools/orchestrator.rs +++ b/codex-rs/core/src/tools/orchestrator.rs @@ -43,7 +43,7 @@ impl ToolOrchestrator { where T: ToolRuntime, { - let otel = turn_ctx.client.get_otel_manager(); + let otel = turn_ctx.otel_manager.clone(); let otel_tn = &tool_ctx.tool_name; let otel_ci = &tool_ctx.call_id; let otel_user = ToolDecisionSource::User; diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index 72b848e315..f67db63943 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -70,7 +70,7 @@ impl ToolRegistry { ) -> Result { let tool_name = invocation.tool_name.clone(); let call_id_owned = invocation.call_id.clone(); - let otel = invocation.turn.client.get_otel_manager(); + let otel = invocation.turn.otel_manager.clone(); let payload_for_response = invocation.payload.clone(); let log_payload = payload_for_response.log_payload(); From 282f42c0ce3c067a1a00c7590f76ef7bec47e207 Mon Sep 17 00:00:00 2001 From: canvrno-oai Date: Wed, 4 Feb 2026 09:38:41 -0800 Subject: [PATCH 6/8] Add option to approve and remember MCP/Apps tool usage (#10584) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This PR adds a new approval option for app/MCP tool calls: “Allow and remember” (session-scoped). When selected, Codex stores a temporary approval and auto-approves matching future calls for the rest of the session. Added a session-scoped approval key (`server`, `connector_id`, `tool_name`) and persisted it in `tool_approvals` as `ApprovedForSession`. On subsequent matching calls, approval is skipped and treated as accepted. - Updated the approval question options to conditionally include: - Accept - Allow and remember (conditional) - Decline - Cancel The new “Allow and remember” option is only shown when all of these are true: 1. The call is routed through the Codex Apps MCP server (codex_apps). 2. The tool requires approval based on annotations: - read_only_hint == false, and - destructive_hint == true or open_world_hint == true. 3. The tool includes a connector_id in metadata (used to build the remembered approval key). If no `connector_id` is present, the prompt still appears (when approval is required), but only with the existing choices (Accept / Decline / Cancel). Approval prompting in this path has an explicit early return unless server == `codex_apps`. --- codex-rs/core/src/mcp_tool_call.rs | 93 ++++++++++++++++++++++++------ 1 file changed, 75 insertions(+), 18 deletions(-) diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index d1d009470d..0c9a4a0033 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -14,12 +14,14 @@ use codex_protocol::mcp::CallToolResult; use codex_protocol::models::FunctionCallOutputPayload; use codex_protocol::models::ResponseInputItem; use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::ReviewDecision; use codex_protocol::protocol::SandboxPolicy; use codex_protocol::request_user_input::RequestUserInputArgs; use codex_protocol::request_user_input::RequestUserInputQuestion; use codex_protocol::request_user_input::RequestUserInputQuestionOption; use codex_protocol::request_user_input::RequestUserInputResponse; use rmcp::model::ToolAnnotations; +use serde::Serialize; use std::sync::Arc; /// Handles the specified tool call dispatches the appropriate @@ -64,7 +66,7 @@ pub(crate) async fn handle_mcp_tool_call( .await { let result = match decision { - McpToolApprovalDecision::Accept => { + McpToolApprovalDecision::Accept | McpToolApprovalDecision::AcceptAndRemember => { let tool_call_begin_event = EventMsg::McpToolCallBegin(McpToolCallBeginEvent { call_id: call_id.clone(), invocation: invocation.clone(), @@ -165,21 +167,31 @@ async fn notify_mcp_tool_call_event(sess: &Session, turn_context: &TurnContext, #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum McpToolApprovalDecision { Accept, + AcceptAndRemember, Decline, Cancel, } struct McpToolApprovalMetadata { annotations: ToolAnnotations, + connector_id: Option, connector_name: Option, tool_title: Option, } const MCP_TOOL_APPROVAL_QUESTION_ID_PREFIX: &str = "mcp_tool_call_approval"; -const MCP_TOOL_APPROVAL_ACCEPT: &str = "Accept"; -const MCP_TOOL_APPROVAL_DECLINE: &str = "Decline"; +const MCP_TOOL_APPROVAL_ACCEPT: &str = "Approve Once"; +const MCP_TOOL_APPROVAL_ACCEPT_AND_REMEMBER: &str = "Approve this Session"; +const MCP_TOOL_APPROVAL_DECLINE: &str = "Deny"; const MCP_TOOL_APPROVAL_CANCEL: &str = "Cancel"; +#[derive(Debug, Serialize)] +struct McpToolApprovalKey { + server: String, + connector_id: String, + tool_name: String, +} + async fn maybe_request_mcp_tool_approval( sess: &Session, turn_context: &TurnContext, @@ -198,6 +210,19 @@ async fn maybe_request_mcp_tool_approval( if !requires_mcp_tool_approval(&metadata.annotations) { return None; } + let approval_key = metadata + .connector_id + .as_deref() + .map(|connector_id| McpToolApprovalKey { + server: server.to_string(), + connector_id: connector_id.to_string(), + tool_name: tool_name.to_string(), + }); + if let Some(key) = approval_key.as_ref() + && mcp_tool_approval_is_remembered(sess, key).await + { + return Some(McpToolApprovalDecision::Accept); + } let question_id = format!("{MCP_TOOL_APPROVAL_QUESTION_ID_PREFIX}_{call_id}"); let question = build_mcp_tool_approval_question( @@ -206,6 +231,7 @@ async fn maybe_request_mcp_tool_approval( metadata.tool_title.as_deref(), metadata.connector_name.as_deref(), &metadata.annotations, + approval_key.is_some(), ); let args = RequestUserInputArgs { questions: vec![question], @@ -213,7 +239,13 @@ async fn maybe_request_mcp_tool_approval( let response = sess .request_user_input(turn_context, call_id.to_string(), args) .await; - Some(parse_mcp_tool_approval_response(response, &question_id)) + let decision = parse_mcp_tool_approval_response(response, &question_id); + if matches!(decision, McpToolApprovalDecision::AcceptAndRemember) + && let Some(key) = approval_key + { + remember_mcp_tool_approval(sess, key).await; + } + Some(decision) } fn is_full_access_mode(turn_context: &TurnContext) -> bool { @@ -244,6 +276,7 @@ async fn lookup_mcp_tool_metadata( .annotations .map(|annotations| McpToolApprovalMetadata { annotations, + connector_id: tool_info.connector_id, connector_name: tool_info.connector_name, tool_title: tool_info.tool.title, }) @@ -259,6 +292,7 @@ fn build_mcp_tool_approval_question( tool_title: Option<&str>, connector_name: Option<&str>, annotations: &ToolAnnotations, + allow_remember_option: bool, ) -> RequestUserInputQuestion { let destructive = annotations.destructive_hint == Some(true); let open_world = annotations.open_world_hint == Some(true); @@ -277,26 +311,34 @@ fn build_mcp_tool_approval_question( "{app_label} wants to run the tool \"{tool_label}\", which {reason}. Allow this action?" ); + let mut options = vec![RequestUserInputQuestionOption { + label: MCP_TOOL_APPROVAL_ACCEPT.to_string(), + description: "Run the tool and continue.".to_string(), + }]; + if allow_remember_option { + options.push(RequestUserInputQuestionOption { + label: MCP_TOOL_APPROVAL_ACCEPT_AND_REMEMBER.to_string(), + description: "Run the tool and remember this choice for this session.".to_string(), + }); + } + options.extend([ + RequestUserInputQuestionOption { + label: MCP_TOOL_APPROVAL_DECLINE.to_string(), + description: "Decline this tool call and continue.".to_string(), + }, + RequestUserInputQuestionOption { + label: MCP_TOOL_APPROVAL_CANCEL.to_string(), + description: "Cancel this tool call".to_string(), + }, + ]); + RequestUserInputQuestion { id: question_id, header: "Approve app tool call?".to_string(), question, is_other: false, is_secret: false, - options: Some(vec![ - RequestUserInputQuestionOption { - label: MCP_TOOL_APPROVAL_ACCEPT.to_string(), - description: "Run the tool and continue.".to_string(), - }, - RequestUserInputQuestionOption { - label: MCP_TOOL_APPROVAL_DECLINE.to_string(), - description: "Decline this tool call and continue.".to_string(), - }, - RequestUserInputQuestionOption { - label: MCP_TOOL_APPROVAL_CANCEL.to_string(), - description: "Cancel this tool call".to_string(), - }, - ]), + options: Some(options), } } @@ -315,6 +357,11 @@ fn parse_mcp_tool_approval_response( return McpToolApprovalDecision::Cancel; }; if answers + .iter() + .any(|answer| answer == MCP_TOOL_APPROVAL_ACCEPT_AND_REMEMBER) + { + McpToolApprovalDecision::AcceptAndRemember + } else if answers .iter() .any(|answer| answer == MCP_TOOL_APPROVAL_ACCEPT) { @@ -329,6 +376,16 @@ fn parse_mcp_tool_approval_response( } } +async fn mcp_tool_approval_is_remembered(sess: &Session, key: &McpToolApprovalKey) -> bool { + let store = sess.services.tool_approvals.lock().await; + matches!(store.get(key), Some(ReviewDecision::ApprovedForSession)) +} + +async fn remember_mcp_tool_approval(sess: &Session, key: McpToolApprovalKey) { + let mut store = sess.services.tool_approvals.lock().await; + store.put(key, ReviewDecision::ApprovedForSession); +} + fn requires_mcp_tool_approval(annotations: &ToolAnnotations) -> bool { annotations.read_only_hint == Some(false) && (annotations.destructive_hint == Some(true) || annotations.open_world_hint == Some(true)) From 71e63f8d10f789b697548bfaa1ef714f049fe226 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Wed, 4 Feb 2026 17:59:22 +0000 Subject: [PATCH 7/8] fix: flaky test (#10644) --- codex-rs/core/tests/suite/tool_parallelism.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/codex-rs/core/tests/suite/tool_parallelism.rs b/codex-rs/core/tests/suite/tool_parallelism.rs index 955b2f7ec3..0df8f8cbc9 100644 --- a/codex-rs/core/tests/suite/tool_parallelism.rs +++ b/codex-rs/core/tests/suite/tool_parallelism.rs @@ -149,6 +149,8 @@ async fn shell_tools_run_in_parallel() -> anyhow::Result<()> { let shell_args = json!({ "command": "sleep 0.3", + // Avoid user-specific shell startup cost (e.g. zsh profile scripts) in timing assertions. + "login": false, "timeout_ms": 1_000, }); let args_one = serde_json::to_string(&shell_args)?; From e9335374b9ef43710e76a411caca5e3a0bdf139b Mon Sep 17 00:00:00 2001 From: jif-oai Date: Wed, 4 Feb 2026 17:59:36 +0000 Subject: [PATCH 8/8] feat: add phase 1 mem client (#10629) Adding a client on top of https://github.com/openai/openai/pull/672176 --- codex-rs/codex-api/README.md | 9 + codex-rs/codex-api/src/common.rs | 28 ++ codex-rs/codex-api/src/endpoint/memories.rs | 108 ++++++++ codex-rs/codex-api/src/endpoint/mod.rs | 1 + codex-rs/codex-api/src/lib.rs | 5 + codex-rs/core/src/client.rs | 58 +++- codex-rs/core/src/lib.rs | 1 + codex-rs/core/src/memory_trace.rs | 292 ++++++++++++++++++++ 8 files changed, 498 insertions(+), 4 deletions(-) create mode 100644 codex-rs/codex-api/src/endpoint/memories.rs create mode 100644 codex-rs/core/src/memory_trace.rs diff --git a/codex-rs/codex-api/README.md b/codex-rs/codex-api/README.md index c1f7d230c0..0570cf7570 100644 --- a/codex-rs/codex-api/README.md +++ b/codex-rs/codex-api/README.md @@ -29,4 +29,13 @@ The public interface of this crate is intentionally small and uniform: - Output: `Vec`. - `CompactClient::compact_input(&CompactionInput, extra_headers)` wraps the JSON encoding and retry/telemetry wiring. +- **Memory trace summarize endpoint** + - Input: `MemoryTraceSummarizeInput` (re-exported as `codex_api::MemoryTraceSummarizeInput`): + - `model: String`. + - `traces: Vec`. + - `MemoryTrace` includes `id`, `metadata.source_path`, and normalized `items`. + - `reasoning: Option`. + - Output: `Vec`. + - `MemoriesClient::trace_summarize_input(&MemoryTraceSummarizeInput, extra_headers)` wraps JSON encoding and retry/telemetry wiring. + All HTTP details (URLs, headers, retry/backoff policies, SSE framing) are encapsulated in `codex-api` and `codex-client`. Callers construct prompts/inputs using protocol types and work with typed streams of `ResponseEvent` or compacted `ResponseItem` values. diff --git a/codex-rs/codex-api/src/common.rs b/codex-rs/codex-api/src/common.rs index a9127644f1..f1a996f12c 100644 --- a/codex-rs/codex-api/src/common.rs +++ b/codex-rs/codex-api/src/common.rs @@ -6,6 +6,7 @@ use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig; use codex_protocol::protocol::RateLimitSnapshot; use codex_protocol::protocol::TokenUsage; use futures::Stream; +use serde::Deserialize; use serde::Serialize; use serde_json::Value; use std::pin::Pin; @@ -37,6 +38,33 @@ pub struct CompactionInput<'a> { pub instructions: &'a str, } +/// Canonical input payload for the memory trace summarize endpoint. +#[derive(Debug, Clone, Serialize)] +pub struct MemoryTraceSummarizeInput { + pub model: String, + pub traces: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub reasoning: Option, +} + +#[derive(Debug, Clone, Serialize)] +pub struct MemoryTrace { + pub id: String, + pub metadata: MemoryTraceMetadata, + pub items: Vec, +} + +#[derive(Debug, Clone, Serialize)] +pub struct MemoryTraceMetadata { + pub source_path: String, +} + +#[derive(Debug, Clone, Deserialize, PartialEq, Eq)] +pub struct MemoryTraceSummaryOutput { + pub trace_summary: String, + pub memory_summary: String, +} + #[derive(Debug)] pub enum ResponseEvent { Created, diff --git a/codex-rs/codex-api/src/endpoint/memories.rs b/codex-rs/codex-api/src/endpoint/memories.rs new file mode 100644 index 0000000000..c8f35d7e16 --- /dev/null +++ b/codex-rs/codex-api/src/endpoint/memories.rs @@ -0,0 +1,108 @@ +use crate::auth::AuthProvider; +use crate::common::MemoryTraceSummarizeInput; +use crate::common::MemoryTraceSummaryOutput; +use crate::endpoint::session::EndpointSession; +use crate::error::ApiError; +use crate::provider::Provider; +use codex_client::HttpTransport; +use codex_client::RequestTelemetry; +use http::HeaderMap; +use http::Method; +use serde::Deserialize; +use serde_json::to_value; +use std::sync::Arc; + +pub struct MemoriesClient { + session: EndpointSession, +} + +impl MemoriesClient { + pub fn new(transport: T, provider: Provider, auth: A) -> Self { + Self { + session: EndpointSession::new(transport, provider, auth), + } + } + + pub fn with_telemetry(self, request: Option>) -> Self { + Self { + session: self.session.with_request_telemetry(request), + } + } + + fn path() -> &'static str { + "memories/trace_summarize" + } + + pub async fn trace_summarize( + &self, + body: serde_json::Value, + extra_headers: HeaderMap, + ) -> Result, ApiError> { + let resp = self + .session + .execute(Method::POST, Self::path(), extra_headers, Some(body)) + .await?; + let parsed: TraceSummarizeResponse = + serde_json::from_slice(&resp.body).map_err(|e| ApiError::Stream(e.to_string()))?; + Ok(parsed.output) + } + + pub async fn trace_summarize_input( + &self, + input: &MemoryTraceSummarizeInput, + extra_headers: HeaderMap, + ) -> Result, ApiError> { + let body = to_value(input).map_err(|e| { + ApiError::Stream(format!( + "failed to encode memory trace summarize input: {e}" + )) + })?; + self.trace_summarize(body, extra_headers).await + } +} + +#[derive(Debug, Deserialize)] +struct TraceSummarizeResponse { + output: Vec, +} + +#[cfg(test)] +mod tests { + use super::*; + use async_trait::async_trait; + use codex_client::Request; + use codex_client::Response; + use codex_client::StreamResponse; + use codex_client::TransportError; + + #[derive(Clone, Default)] + struct DummyTransport; + + #[async_trait] + impl HttpTransport for DummyTransport { + async fn execute(&self, _req: Request) -> Result { + Err(TransportError::Build("execute should not run".to_string())) + } + + async fn stream(&self, _req: Request) -> Result { + Err(TransportError::Build("stream should not run".to_string())) + } + } + + #[derive(Clone, Default)] + struct DummyAuth; + + impl AuthProvider for DummyAuth { + fn bearer_token(&self) -> Option { + None + } + } + + #[test] + fn path_is_memories_trace_summarize() { + assert_eq!( + MemoriesClient::::path(), + "memories/trace_summarize" + ); + } +} diff --git a/codex-rs/codex-api/src/endpoint/mod.rs b/codex-rs/codex-api/src/endpoint/mod.rs index 23579ffcf1..0dede138e8 100644 --- a/codex-rs/codex-api/src/endpoint/mod.rs +++ b/codex-rs/codex-api/src/endpoint/mod.rs @@ -1,5 +1,6 @@ pub mod aggregate; pub mod compact; +pub mod memories; pub mod models; pub mod responses; pub mod responses_websocket; diff --git a/codex-rs/codex-api/src/lib.rs b/codex-rs/codex-api/src/lib.rs index b0c70084d4..70652d2d78 100644 --- a/codex-rs/codex-api/src/lib.rs +++ b/codex-rs/codex-api/src/lib.rs @@ -15,6 +15,10 @@ pub use codex_client::TransportError; pub use crate::auth::AuthProvider; pub use crate::common::CompactionInput; +pub use crate::common::MemoryTrace; +pub use crate::common::MemoryTraceMetadata; +pub use crate::common::MemoryTraceSummarizeInput; +pub use crate::common::MemoryTraceSummaryOutput; pub use crate::common::Prompt; pub use crate::common::ResponseAppendWsRequest; pub use crate::common::ResponseCreateWsRequest; @@ -24,6 +28,7 @@ pub use crate::common::ResponsesApiRequest; pub use crate::common::create_text_param_for_request; pub use crate::endpoint::aggregate::AggregateStreamExt; pub use crate::endpoint::compact::CompactClient; +pub use crate::endpoint::memories::MemoriesClient; pub use crate::endpoint::models::ModelsClient; pub use crate::endpoint::responses::ResponsesClient; pub use crate::endpoint::responses::ResponsesOptions; diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index dd92c5da01..0a0f702b17 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -7,6 +7,10 @@ use crate::api_bridge::map_api_error; use crate::auth::UnauthorizedRecovery; use codex_api::CompactClient as ApiCompactClient; use codex_api::CompactionInput as ApiCompactionInput; +use codex_api::MemoriesClient as ApiMemoriesClient; +use codex_api::MemoryTrace as ApiMemoryTrace; +use codex_api::MemoryTraceSummarizeInput as ApiMemoryTraceSummarizeInput; +use codex_api::MemoryTraceSummaryOutput as ApiMemoryTraceSummaryOutput; use codex_api::Prompt as ApiPrompt; use codex_api::RequestTelemetry; use codex_api::ReqwestTransport; @@ -183,6 +187,55 @@ impl ModelClient { instructions: &instructions, }; + let extra_headers = self.build_subagent_headers(); + client + .compact_input(&payload, extra_headers) + .await + .map_err(map_api_error) + } + + /// Builds memory summaries for each provided normalized trace. + /// + /// This is a unary call (no streaming) to `/v1/memories/trace_summarize`. + pub async fn summarize_memory_traces( + &self, + traces: Vec, + ) -> Result> { + if traces.is_empty() { + return Ok(Vec::new()); + } + + let auth_manager = self.state.auth_manager.clone(); + let auth = match auth_manager.as_ref() { + Some(manager) => manager.auth().await, + None => None, + }; + let api_provider = self + .state + .provider + .to_api_provider(auth.as_ref().map(CodexAuth::internal_auth_mode))?; + let api_auth = auth_provider_from_auth(auth, &self.state.provider)?; + let transport = ReqwestTransport::new(build_reqwest_client()); + let request_telemetry = self.build_request_telemetry(); + let client = ApiMemoriesClient::new(transport, api_provider, api_auth) + .with_telemetry(Some(request_telemetry)); + + let payload = ApiMemoryTraceSummarizeInput { + model: self.state.model_info.slug.clone(), + traces, + reasoning: self.state.effort.map(|effort| Reasoning { + effort: Some(effort), + summary: None, + }), + }; + + client + .trace_summarize_input(&payload, self.build_subagent_headers()) + .await + .map_err(map_api_error) + } + + fn build_subagent_headers(&self) -> ApiHeaderMap { let mut extra_headers = ApiHeaderMap::new(); if let SessionSource::SubAgent(sub) = &self.state.session_source { let subagent = match sub { @@ -195,10 +248,7 @@ impl ModelClient { extra_headers.insert("x-openai-subagent", val); } } - client - .compact_input(&payload, extra_headers) - .await - .map_err(map_api_error) + extra_headers } } diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 9bd4c87253..36fdbdc9a6 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -161,4 +161,5 @@ pub use codex_protocol::models::ResponseItem; pub use compact::content_items_to_text; pub use event_mapping::parse_turn_item; pub mod compact; +pub mod memory_trace; pub mod otel_init; diff --git a/codex-rs/core/src/memory_trace.rs b/codex-rs/core/src/memory_trace.rs new file mode 100644 index 0000000000..807ff64df1 --- /dev/null +++ b/codex-rs/core/src/memory_trace.rs @@ -0,0 +1,292 @@ +use std::path::Path; +use std::path::PathBuf; + +use crate::ModelClient; +use crate::error::CodexErr; +use crate::error::Result; +use codex_api::MemoryTrace as ApiMemoryTrace; +use codex_api::MemoryTraceMetadata as ApiMemoryTraceMetadata; +use serde_json::Map; +use serde_json::Value; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct BuiltTraceMemory { + pub trace_id: String, + pub source_path: PathBuf, + pub trace_summary: String, + pub memory_summary: String, +} + +struct PreparedTrace { + trace_id: String, + source_path: PathBuf, + payload: ApiMemoryTrace, +} + +/// Loads raw trace files, normalizes trace items, and builds memory summaries. +/// +/// The request/response wiring mirrors the memory trace summarize E2E flow: +/// `/v1/memories/trace_summarize` with one output object per input trace. +pub async fn build_memories_from_trace_files( + client: &ModelClient, + trace_paths: &[PathBuf], +) -> Result> { + if trace_paths.is_empty() { + return Ok(Vec::new()); + } + + let mut prepared = Vec::with_capacity(trace_paths.len()); + for (index, path) in trace_paths.iter().enumerate() { + prepared.push(prepare_trace(index + 1, path).await?); + } + + let traces = prepared.iter().map(|trace| trace.payload.clone()).collect(); + let output = client.summarize_memory_traces(traces).await?; + if output.len() != prepared.len() { + return Err(CodexErr::InvalidRequest(format!( + "unexpected memory summarize output length: expected {}, got {}", + prepared.len(), + output.len() + ))); + } + + Ok(prepared + .into_iter() + .zip(output) + .map(|(trace, summary)| BuiltTraceMemory { + trace_id: trace.trace_id, + source_path: trace.source_path, + trace_summary: summary.trace_summary, + memory_summary: summary.memory_summary, + }) + .collect()) +} + +async fn prepare_trace(index: usize, path: &Path) -> Result { + let text = load_trace_text(path).await?; + let items = load_trace_items(path, &text)?; + let trace_id = build_trace_id(index, path); + let source_path = path.to_path_buf(); + + Ok(PreparedTrace { + trace_id: trace_id.clone(), + source_path: source_path.clone(), + payload: ApiMemoryTrace { + id: trace_id, + metadata: ApiMemoryTraceMetadata { + source_path: source_path.display().to_string(), + }, + items, + }, + }) +} + +async fn load_trace_text(path: &Path) -> Result { + let raw = tokio::fs::read(path).await?; + Ok(decode_trace_bytes(&raw)) +} + +fn decode_trace_bytes(raw: &[u8]) -> String { + if let Some(without_bom) = raw.strip_prefix(&[0xEF, 0xBB, 0xBF]) + && let Ok(text) = String::from_utf8(without_bom.to_vec()) + { + return text; + } + if let Ok(text) = String::from_utf8(raw.to_vec()) { + return text; + } + raw.iter().map(|b| char::from(*b)).collect() +} + +fn load_trace_items(path: &Path, text: &str) -> Result> { + if let Ok(Value::Array(items)) = serde_json::from_str::(text) { + let dict_items = items + .into_iter() + .filter(serde_json::Value::is_object) + .collect::>(); + if dict_items.is_empty() { + return Err(CodexErr::InvalidRequest(format!( + "no object items found in trace file: {}", + path.display() + ))); + } + return normalize_trace_items(dict_items, path); + } + + let mut parsed_items = Vec::new(); + for line in text.lines() { + let line = line.trim(); + if line.is_empty() || (!line.starts_with('{') && !line.starts_with('[')) { + continue; + } + + let Ok(obj) = serde_json::from_str::(line) else { + continue; + }; + + match obj { + Value::Object(_) => parsed_items.push(obj), + Value::Array(inner) => { + parsed_items.extend(inner.into_iter().filter(serde_json::Value::is_object)) + } + _ => {} + } + } + + if parsed_items.is_empty() { + return Err(CodexErr::InvalidRequest(format!( + "no JSON items parsed from trace file: {}", + path.display() + ))); + } + + normalize_trace_items(parsed_items, path) +} + +fn normalize_trace_items(items: Vec, path: &Path) -> Result> { + let mut normalized = Vec::new(); + + for item in items { + let Value::Object(obj) = item else { + continue; + }; + + if let Some(payload) = obj.get("payload") { + if obj.get("type").and_then(Value::as_str) != Some("response_item") { + continue; + } + + match payload { + Value::Object(payload_item) => { + if is_allowed_trace_item(payload_item) { + normalized.push(Value::Object(payload_item.clone())); + } + } + Value::Array(payload_items) => { + for payload_item in payload_items { + if let Value::Object(payload_item) = payload_item + && is_allowed_trace_item(payload_item) + { + normalized.push(Value::Object(payload_item.clone())); + } + } + } + _ => {} + } + continue; + } + + if is_allowed_trace_item(&obj) { + normalized.push(Value::Object(obj)); + } + } + + if normalized.is_empty() { + return Err(CodexErr::InvalidRequest(format!( + "no valid trace items after normalization: {}", + path.display() + ))); + } + Ok(normalized) +} + +fn is_allowed_trace_item(item: &Map) -> bool { + let Some(item_type) = item.get("type").and_then(Value::as_str) else { + return false; + }; + + if item_type == "message" { + return matches!( + item.get("role").and_then(Value::as_str), + Some("assistant" | "system" | "developer" | "user") + ); + } + + true +} + +fn build_trace_id(index: usize, path: &Path) -> String { + let stem = path + .file_stem() + .map(|stem| stem.to_string_lossy().into_owned()) + .filter(|stem| !stem.is_empty()) + .unwrap_or_else(|| "trace".to_string()); + format!("trace_{index}_{stem}") +} + +#[cfg(test)] +mod tests { + use super::*; + use pretty_assertions::assert_eq; + use tempfile::tempdir; + + #[test] + fn normalize_trace_items_handles_payload_wrapper_and_message_role_filtering() { + let items = vec![ + serde_json::json!({ + "type": "response_item", + "payload": {"type": "message", "role": "assistant", "content": []} + }), + serde_json::json!({ + "type": "response_item", + "payload": [ + {"type": "message", "role": "user", "content": []}, + {"type": "message", "role": "tool", "content": []}, + {"type": "function_call", "name": "shell", "arguments": "{}", "call_id": "c1"} + ] + }), + serde_json::json!({ + "type": "not_response_item", + "payload": {"type": "message", "role": "assistant", "content": []} + }), + serde_json::json!({ + "type": "message", + "role": "developer", + "content": [] + }), + ]; + + let normalized = normalize_trace_items(items, Path::new("trace.json")).expect("normalize"); + let expected = vec![ + serde_json::json!({"type": "message", "role": "assistant", "content": []}), + serde_json::json!({"type": "message", "role": "user", "content": []}), + serde_json::json!({"type": "function_call", "name": "shell", "arguments": "{}", "call_id": "c1"}), + serde_json::json!({"type": "message", "role": "developer", "content": []}), + ]; + assert_eq!(normalized, expected); + } + + #[test] + fn load_trace_items_supports_jsonl_arrays_and_objects() { + let text = r#" +{"type":"response_item","payload":{"type":"message","role":"assistant","content":[]}} +[{"type":"message","role":"user","content":[]},{"type":"message","role":"tool","content":[]}] +"#; + let loaded = load_trace_items(Path::new("trace.jsonl"), text).expect("load"); + let expected = vec![ + serde_json::json!({"type":"message","role":"assistant","content":[]}), + serde_json::json!({"type":"message","role":"user","content":[]}), + ]; + assert_eq!(loaded, expected); + } + + #[tokio::test] + async fn load_trace_text_decodes_utf8_sig() { + let dir = tempdir().expect("tempdir"); + let path = dir.path().join("trace.json"); + tokio::fs::write( + &path, + [ + 0xEF, 0xBB, 0xBF, b'[', b'{', b'"', b't', b'y', b'p', b'e', b'"', b':', b'"', b'm', + b'e', b's', b's', b'a', b'g', b'e', b'"', b',', b'"', b'r', b'o', b'l', b'e', b'"', + b':', b'"', b'u', b's', b'e', b'r', b'"', b',', b'"', b'c', b'o', b'n', b't', b'e', + b'n', b't', b'"', b':', b'[', b']', b'}', b']', + ], + ) + .await + .expect("write"); + + let text = load_trace_text(&path).await.expect("decode"); + assert!(text.starts_with('[')); + } +}