diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index e06dba9b68..176ea9413f 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -7009,7 +7009,9 @@ async fn handle_output_item_done_records_image_save_history_message() { let mut ctx = HandleOutputCtx { sess: Arc::clone(&session), turn_context: Arc::clone(&turn_context), - turn_store: Arc::new(codex_extension_api::ExtensionData::new()), + turn_store: Arc::new(codex_extension_api::ExtensionData::new( + turn_context.sub_id.clone(), + )), tool_runtime: test_tool_runtime(Arc::clone(&session), Arc::clone(&turn_context)), cancellation_token: CancellationToken::new(), }; @@ -7062,7 +7064,9 @@ async fn handle_output_item_done_skips_image_save_message_when_save_fails() { let mut ctx = HandleOutputCtx { sess: Arc::clone(&session), turn_context: Arc::clone(&turn_context), - turn_store: Arc::new(codex_extension_api::ExtensionData::new()), + turn_store: Arc::new(codex_extension_api::ExtensionData::new( + turn_context.sub_id.clone(), + )), tool_runtime: test_tool_runtime(Arc::clone(&session), Arc::clone(&turn_context)), cancellation_token: CancellationToken::new(), }; @@ -8950,7 +8954,7 @@ async fn tool_calls_reopen_mailbox_delivery_for_current_turn() { let mut ctx = HandleOutputCtx { sess: Arc::clone(&sess), turn_context: Arc::clone(&tc), - turn_store: Arc::new(codex_extension_api::ExtensionData::new()), + turn_store: Arc::new(codex_extension_api::ExtensionData::new(tc.sub_id.clone())), tool_runtime: test_tool_runtime(Arc::clone(&sess), Arc::clone(&tc)), cancellation_token: CancellationToken::new(), }; diff --git a/codex-rs/core/src/session/turn_tests.rs b/codex-rs/core/src/session/turn_tests.rs index 4f19474d38..f1cf86b46c 100644 --- a/codex-rs/core/src/session/turn_tests.rs +++ b/codex-rs/core/src/session/turn_tests.rs @@ -43,7 +43,7 @@ async fn plan_mode_uses_contributed_turn_item_for_last_agent_message() { let mut builder = codex_extension_api::ExtensionRegistryBuilder::new(); builder.turn_item_contributor(Arc::new(RewriteAgentMessageContributor)); session.services.extensions = Arc::new(builder.build()); - let turn_store = ExtensionData::new(); + let turn_store = ExtensionData::new(turn_context.sub_id.clone()); let mut state = PlanModeStreamState::new(&turn_context.sub_id); let mut last_agent_message = None; let item = assistant_output_text("original assistant text"); diff --git a/codex-rs/core/src/state/turn.rs b/codex-rs/core/src/state/turn.rs index 0dd0c09379..70aae74e87 100644 --- a/codex-rs/core/src/state/turn.rs +++ b/codex-rs/core/src/state/turn.rs @@ -9,6 +9,7 @@ use tokio::sync::Notify; use tokio_util::sync::CancellationToken; use tokio_util::task::AbortOnDropHandle; +use codex_extension_api::ExtensionData; use codex_protocol::dynamic_tools::DynamicToolResponse; use codex_protocol::models::ResponseInputItem; use codex_protocol::request_permissions::RequestPermissionProfile; @@ -75,6 +76,7 @@ pub(crate) struct RunningTask { pub(crate) cancellation_token: CancellationToken, pub(crate) handle: AbortOnDropHandle<()>, pub(crate) turn_context: Arc, + pub(crate) turn_extension_data: Arc, // Timer recorded when the task drops to capture the full turn duration. pub(crate) _timer: Option, } diff --git a/codex-rs/core/src/stream_events_utils_tests.rs b/codex-rs/core/src/stream_events_utils_tests.rs index 46c7ffed19..2d095906ed 100644 --- a/codex-rs/core/src/stream_events_utils_tests.rs +++ b/codex-rs/core/src/stream_events_utils_tests.rs @@ -214,7 +214,7 @@ async fn handle_non_tool_response_item_runs_turn_item_contributors_only_when_req let mut builder = codex_extension_api::ExtensionRegistryBuilder::new(); builder.turn_item_contributor(Arc::new(TestTurnItemContributor)); session.services.extensions = Arc::new(builder.build()); - let turn_store = ExtensionData::new(); + let turn_store = ExtensionData::new(turn_context.sub_id.clone()); let item = assistant_output_text( "helloignored by memory parser world", ); @@ -288,8 +288,8 @@ async fn handle_output_item_done_returns_contributed_last_agent_message() { let item = assistant_output_text("original assistant text"); let mut ctx = HandleOutputCtx { sess: session, - turn_context, - turn_store: Arc::new(ExtensionData::new()), + turn_context: Arc::clone(&turn_context), + turn_store: Arc::new(ExtensionData::new(turn_context.sub_id.clone())), tool_runtime, cancellation_token: CancellationToken::new(), }; @@ -310,7 +310,7 @@ async fn finalized_turn_item_defers_mailbox_for_contributed_visible_text() { let mut builder = codex_extension_api::ExtensionRegistryBuilder::new(); builder.turn_item_contributor(Arc::new(RewriteAgentMessageContributor)); session.services.extensions = Arc::new(builder.build()); - let turn_store = ExtensionData::new(); + let turn_store = ExtensionData::new(turn_context.sub_id.clone()); let item = assistant_output_text("hidden only"); let finalized = finalize_non_tool_response_item( @@ -336,7 +336,7 @@ async fn finalized_turn_item_keeps_mailbox_open_for_commentary_text() { let mut builder = codex_extension_api::ExtensionRegistryBuilder::new(); builder.turn_item_contributor(Arc::new(RewriteAgentMessageContributor)); session.services.extensions = Arc::new(builder.build()); - let turn_store = ExtensionData::new(); + let turn_store = ExtensionData::new(turn_context.sub_id.clone()); let item = assistant_output_text_with_phase("still working", Some(MessagePhase::Commentary)); let finalized = finalize_non_tool_response_item( diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 5e6b931500..58ead0fd25 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -367,6 +367,7 @@ impl Session { } self.emit_turn_start_lifecycle(turn_context.extension_data.as_ref()); + let turn_extension_data = Arc::clone(&turn_context.extension_data); let mut active = self.active_turn.lock().await; let turn = active.get_or_insert_with(ActiveTurn::default); debug_assert!(turn.tasks.is_empty()); @@ -439,6 +440,7 @@ impl Session { task, cancellation_token, turn_context: Arc::clone(&turn_context), + turn_extension_data, _timer: timer, }; turn.add_task(running_task);