diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index d46a98ed31..9e4cf7df6f 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -508,6 +508,16 @@ impl ModelClient { responses_metadata.session_id.clone() } + // ChatGPT derives cache affinity from the Responses session-id header. Keep the + // actual session identity in turn metadata, hooks, and history/notes requests. + fn responses_session_id(&self, metadata: &CodexResponsesMetadata) -> String { + if self.state.session_source.is_non_root_agent() { + metadata.session_id.clone() + } else { + self.prompt_cache_key(metadata) + } + } + /// Creates a fresh turn-scoped streaming session. /// /// This constructor does not perform network I/O itself; the session opens a websocket lazily @@ -1162,7 +1172,7 @@ impl ModelClient { headers.insert("x-client-request-id", header_value); } headers.extend(build_session_headers( - Some(responses_metadata.session_id.to_string()), + Some(self.responses_session_id(responses_metadata)), Some(responses_metadata.thread_id.to_string()), )); headers.extend(self.build_responses_compatibility_headers(responses_metadata)); @@ -1217,7 +1227,7 @@ impl ModelClientSession { use_responses_lite: bool, ) -> ApiResponsesOptions { ApiResponsesOptions { - session_id: Some(responses_metadata.session_id.to_string()), + session_id: Some(self.client.responses_session_id(responses_metadata)), thread_id: Some(responses_metadata.thread_id.to_string()), session_source: Some(self.client.state.session_source.clone()), extra_headers: { diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 7f1fa2ef4e..8936f5afab 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -1276,9 +1276,13 @@ fn guardian_reviewer_uses_dedicated_endpoint_only_with_codex_backend_auth() { ); } +#[test_case::test_case(/*cache_key*/ None; "own_cache")] +#[test_case::test_case(Some("parent-session"); "inherited_cache")] #[tokio::test] -async fn websocket_handshake_includes_attestation_for_chatgpt_codex_responses() { - let (model_client, attestation_calls) = +async fn websocket_handshake_includes_attestation_for_chatgpt_codex_responses( + cache_key: Option<&str>, +) { + let (mut model_client, attestation_calls) = model_client_with_counting_attestation(/*include_attestation*/ true); let responses_metadata = test_responses_metadata_for_client( &model_client, @@ -1288,6 +1292,7 @@ async fn websocket_handshake_includes_attestation_for_chatgpt_codex_responses() TestCodexResponsesRequestKind::WebsocketConnection, ); + model_client.prompt_cache_key_override = cache_key.map(str::to_string); let headers = model_client .build_websocket_headers(&responses_metadata) .await; @@ -1299,6 +1304,11 @@ async fn websocket_handshake_includes_attestation_for_chatgpt_codex_responses() Some("v1.header-1"), ); assert_eq!(attestation_calls.load(Ordering::Relaxed), 1); + assert_eq!( + headers["session-id"], + cache_key.unwrap_or(&responses_metadata.session_id) + ); + assert_eq!(headers["thread-id"], responses_metadata.thread_id); } #[tokio::test] diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 101fdccb94..aebfd641e5 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -799,6 +799,22 @@ impl Session { )); } }; + // Ephemeral forks reuse cache routing, without sharing storage or lifecycle identity. + let fork_cache_key = match &initial_history { + InitialHistory::Forked(items) + if config.ephemeral + && !session_configuration.session_source.is_non_root_agent() => + { + items.iter().find_map(|item| match item { + RolloutItem::SessionMeta(meta) => Some(meta.meta.session_id.to_string()), + _ => None, + }) + } + InitialHistory::New + | InitialHistory::Cleared + | InitialHistory::Resumed(_) + | InitialHistory::Forked(_) => None, + }; let resumed_session_id = match &initial_history { InitialHistory::Resumed(resumed) => { resumed.history.iter().find_map(|item| match item { @@ -1580,7 +1596,8 @@ impl Session { crate::guardian::prompt_cache_key_override_for_review_session( &session_configuration.session_source, session_configuration.parent_thread_id, - ), + ) + .or(fork_cache_key), tx_event.clone(), ), executed_tool_calls: executed_tool_calls.clone(), diff --git a/codex-rs/core/tests/suite/prompt_cache_key.rs b/codex-rs/core/tests/suite/prompt_cache_key.rs index 5e258a936f..4b429a96e9 100644 --- a/codex-rs/core/tests/suite/prompt_cache_key.rs +++ b/codex-rs/core/tests/suite/prompt_cache_key.rs @@ -158,3 +158,63 @@ async fn api_key_subagent_uses_session_id_as_prompt_cache_key() -> Result<()> { Ok(()) } + +#[tokio::test] +async fn ephemeral_fork_shares_cache_routing_but_keeps_session_identity() -> Result<()> { + use codex_core::ForkSnapshot; + use codex_core::StartThreadOptions; + use core_test_support::responses::mount_sse_sequence; + + let server = start_mock_server().await; + let requests = mount_sse_sequence( + &server, + vec![ + sse(vec![ev_completed("parent")]), + sse(vec![ev_completed("fork")]), + ], + ) + .await; + let mut test = test_codex().build_with_auto_env(&server).await?; + test.submit_text_turn("parent").await?; + test.codex.flush_rollout().await?; + let parent_session = test.session_configured.session_id.to_string(); + let mut config = test.config.clone(); + config.ephemeral = true; + let fork = test + .thread_manager + .fork_thread( + ForkSnapshot::TruncateBeforeNthUserMessage(usize::MAX), + StartThreadOptions { + environments: Some(test.codex.environment_selections().await), + ..StartThreadOptions::new(config) + }, + test.codex.rollout_path().expect("parent rollout"), + ) + .await?; + let fork_session = fork.session_configured.session_id.to_string(); + assert_ne!(fork_session, parent_session); + test.codex = fork.thread; + test.submit_text_turn("side").await?; + let requests = requests.requests(); + let body = requests[1].body_json(); + let metadata: Value = serde_json::from_str( + body["client_metadata"]["x-codex-turn-metadata"] + .as_str() + .expect("turn metadata"), + )?; + assert_eq!( + json!({ + "cache": body["prompt_cache_key"], + "route": requests[1].header("session-id"), + "session": metadata["session_id"], + "thread": requests[1].header("thread-id"), + "tools": body["tools"], + }), + json!({ + "cache": parent_session, "route": parent_session, + "session": fork_session, "thread": fork.thread_id.to_string(), + "tools": requests[0].body_json()["tools"], + }), + ); + Ok(()) +} diff --git a/codex-rs/ext/goal/src/extension.rs b/codex-rs/ext/goal/src/extension.rs index e95280482e..f5dc85138a 100644 --- a/codex-rs/ext/goal/src/extension.rs +++ b/codex-rs/ext/goal/src/extension.rs @@ -15,8 +15,10 @@ use codex_extension_api::ThreadResumeInput; use codex_extension_api::ThreadStartInput; use codex_extension_api::ThreadStopInput; use codex_extension_api::TokenUsageContributor; +use codex_extension_api::ToolCall; use codex_extension_api::ToolCallOutcome; use codex_extension_api::ToolContributor; +use codex_extension_api::ToolExecutor; use codex_extension_api::ToolFinishInput; use codex_extension_api::ToolLifecycleContributor; use codex_extension_api::ToolLifecycleFuture; @@ -101,11 +103,12 @@ where Box::pin(async move { let config = (self.goal_config)(input.config); let enabled = config.enabled; - let tools_available_for_thread = input.persistent_thread_state_available - && !matches!( - input.session_source, - SessionSource::SubAgent(SubAgentSource::Review) - ); + let tools_visible_for_thread = !matches!( + input.session_source, + SessionSource::SubAgent(SubAgentSource::Review) + ); + let tools_available_for_thread = + input.persistent_thread_state_available && tools_visible_for_thread; input.thread_store.insert(config); let accounting_state = input .thread_store @@ -149,6 +152,7 @@ where analytics: self.analytics.clone(), enabled, tools_available_for_thread, + tools_visible_for_thread, root_accounting_state, }, ) @@ -538,16 +542,16 @@ where .get::() .and_then(|config| config.max_goal_token_budget); - vec![ - Arc::new(GoalToolExecutor::get( + let tools = [ + GoalToolExecutor::get( runtime.thread_id(), Arc::clone(&self.state_dbs), runtime.accounting_state(), self.analytics.clone(), self.event_emitter.clone(), self.metrics.clone(), - )), - Arc::new(GoalToolExecutor::create( + ), + GoalToolExecutor::create( runtime.thread_id(), Arc::clone(&self.state_dbs), runtime.accounting_state(), @@ -555,16 +559,23 @@ where self.event_emitter.clone(), self.metrics.clone(), max_goal_token_budget, - )), - Arc::new(GoalToolExecutor::update( + ), + GoalToolExecutor::update( runtime.thread_id(), Arc::clone(&self.state_dbs), runtime.accounting_state(), self.analytics.clone(), self.event_emitter.clone(), self.metrics.clone(), - )), - ] + ), + ]; + tools + .into_iter() + .map(|mut tool| { + tool.execution_allowed = runtime.tools_available(); + Arc::new(tool) as Arc ToolExecutor>> + }) + .collect() } } diff --git a/codex-rs/ext/goal/src/runtime.rs b/codex-rs/ext/goal/src/runtime.rs index 477290c751..7d826ac238 100644 --- a/codex-rs/ext/goal/src/runtime.rs +++ b/codex-rs/ext/goal/src/runtime.rs @@ -33,6 +33,7 @@ pub(crate) struct GoalRuntimeConfig { pub(crate) analytics: GoalAnalytics, pub(crate) enabled: bool, pub(crate) tools_available_for_thread: bool, + pub(crate) tools_visible_for_thread: bool, pub(crate) root_accounting_state: Option>, } @@ -54,6 +55,7 @@ struct GoalRuntimeInner { root_accounting_state: Option>, enabled: AtomicBool, tools_available_for_thread: bool, + tools_visible_for_thread: bool, goal_state_lock: Semaphore, } @@ -107,6 +109,7 @@ impl GoalRuntimeHandle { root_accounting_state: config.root_accounting_state, enabled: AtomicBool::new(config.enabled), tools_available_for_thread: config.tools_available_for_thread, + tools_visible_for_thread: config.tools_visible_for_thread, goal_state_lock: Semaphore::new(/*permits*/ 1), }), } @@ -121,6 +124,10 @@ impl GoalRuntimeHandle { } pub(crate) fn tools_visible(&self) -> bool { + self.is_enabled() && self.inner.tools_visible_for_thread + } + + pub(crate) fn tools_available(&self) -> bool { self.is_enabled() && self.inner.tools_available_for_thread } @@ -416,7 +423,7 @@ impl GoalRuntimeHandle { } pub(crate) async fn continue_if_idle(&self) -> Result<(), String> { - if !self.tools_visible() { + if !self.tools_available() { self.inner.accounting_state.clear_active_goal(); return Ok(()); } diff --git a/codex-rs/ext/goal/src/tool.rs b/codex-rs/ext/goal/src/tool.rs index 41bb312282..ff504060ee 100644 --- a/codex-rs/ext/goal/src/tool.rs +++ b/codex-rs/ext/goal/src/tool.rs @@ -29,6 +29,7 @@ use crate::spec::create_update_goal_tool; #[derive(Clone)] pub(crate) struct GoalToolExecutor { + pub(crate) execution_allowed: bool, kind: GoalToolKind, thread_id: ThreadId, state_db: Arc, @@ -84,6 +85,7 @@ impl GoalToolExecutor { ) -> Self { Self { kind: GoalToolKind::Get, + execution_allowed: true, thread_id, state_db, accounting_state, @@ -105,6 +107,7 @@ impl GoalToolExecutor { ) -> Self { Self { kind: GoalToolKind::Create, + execution_allowed: true, thread_id, state_db, accounting_state, @@ -125,6 +128,7 @@ impl GoalToolExecutor { ) -> Self { Self { kind: GoalToolKind::Update, + execution_allowed: true, thread_id, state_db, accounting_state, @@ -161,6 +165,11 @@ impl<'call> ToolExecutor> for GoalToolExecutor { 'call: 'a, { Box::pin(async move { + if !self.execution_allowed { + return Err(FunctionCallError::RespondToModel( + "Goal tools require a persistent thread.".to_string(), + )); + } match self.kind { GoalToolKind::Get => self.handle_get(invocation).await, GoalToolKind::Create => self.handle_create(invocation).await, diff --git a/codex-rs/ext/goal/tests/goal_extension_backend.rs b/codex-rs/ext/goal/tests/goal_extension_backend.rs index 7c1dbee345..124bfbd80f 100644 --- a/codex-rs/ext/goal/tests/goal_extension_backend.rs +++ b/codex-rs/ext/goal/tests/goal_extension_backend.rs @@ -148,18 +148,39 @@ async fn installed_goal_tools_apply_maximum_token_budget() -> anyhow::Result<()> } #[tokio::test] -async fn goal_tools_hidden_for_ephemeral_threads() -> anyhow::Result<()> { +async fn ephemeral_goal_tools_preserve_specs_but_reject_execution() -> anyhow::Result<()> { let runtime = test_runtime().await?; let thread_id = test_thread_id()?; let tools = installed_tools_with_start( - runtime, + runtime.clone(), thread_id, SessionSource::Cli, /*persistent_thread_state_available*/ false, ) .await; - assert_eq!(Vec::::new(), tool_names(&tools)); + let parent_tools = installed_tools(runtime, ThreadId::new()).await; + assert_eq!( + tools.iter().map(|tool| tool.spec()).collect::>(), + parent_tools + .iter() + .map(|tool| tool.spec()) + .collect::>(), + ); + for tool in tools { + let Err(error) = tool + .handle(tool_call(&tool.tool_name().name, "ephemeral", json!({}))) + .await + else { + panic!("ephemeral goal execution should fail"); + }; + assert_eq!( + error, + FunctionCallError::RespondToModel( + "Goal tools require a persistent thread.".to_string() + ) + ); + } Ok(()) }