Preserve parent cache affinity for ephemeral forks (#44862)

## Why

ChatGPT derives Responses cache affinity from the `session-id` header. Ephemeral forks need to reuse their parent's cache routing while retaining their own session and thread identities.

## What changed

- Inherit the parent session ID as the `prompt_cache_key` for ephemeral root forks.
- Use the prompt cache key for root-agent Responses `session-id` headers, including WebSocket handshakes. Preserve the actual session identity in turn metadata and leave non-root-agent routing unchanged.
- Keep enabled goal tool definitions visible on ephemeral threads, but reject execution with `Goal tools require a persistent thread.` and disable automatic goal continuation without persistent state.

## Testing

Add regression coverage for inherited cache routing with distinct fork identities, matching parent and fork tool definitions, WebSocket session headers, and rejection of ephemeral goal tool execution.

GitOrigin-RevId: d235e97630068b27f7ebd562dd23ac6266ebc1b9

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: {

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]

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(),

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(())
+}

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::<GoalExtensionConfig>()
             .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<dyn for<'call> ToolExecutor<ToolCall<'call>>>
+            })
+            .collect()
     }
 }
 

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<Arc<GoalAccountingState>>,
 }
 
@@ -54,6 +55,7 @@ struct GoalRuntimeInner {
     root_accounting_state: Option<Arc<GoalAccountingState>>,
     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(());
         }

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<codex_state::StateRuntime>,
@@ -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<ToolCall<'call>> 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,

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::<String>::new(), tool_names(&tools));
+    let parent_tools = installed_tools(runtime, ThreadId::new()).await;
+    assert_eq!(
+        tools.iter().map(|tool| tool.spec()).collect::<Vec<_>>(),
+        parent_tools
+            .iter()
+            .map(|tool| tool.spec())
+            .collect::<Vec<_>>(),
+    );
+    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(())
 }