Fix turn extension data task plumbing (#22646)

## Summary
- carry the per-turn extension data through RunningTask so abort
handling can rebuild SessionTaskContext
- update stale test ExtensionData::new() callsites to pass the turn id

## Testing
- Not run after PR branch creation; CI will cover.
This commit is contained in:
jif-oai
2026-05-14 16:00:06 +02:00
committed by GitHub
parent 9ea38136b0
commit 12bfb57139
5 changed files with 17 additions and 9 deletions

View File

@@ -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(),
};

View File

@@ -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");