diff --git a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-experimental.json.zst b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-experimental.json.zst index 9295fb23f4..4bd6d16f44 100644 Binary files a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-experimental.json.zst and b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-experimental.json.zst differ diff --git a/codex-rs/app-server-protocol/src/protocol/v2/thread_sample.rs b/codex-rs/app-server-protocol/src/protocol/v2/thread_sample.rs index adb65b8a2c..a1d325a82b 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/thread_sample.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/thread_sample.rs @@ -14,9 +14,11 @@ use std::collections::HashMap; #[ts(export_to = "v2/")] pub struct ThreadSampleParams { pub thread_id: String, - /// The loaded source's latest completed, failed, or interrupted turn. - /// Stale or active sources are rejected. + /// The loaded source's latest turn, which must be finished (completed, failed, or interrupted) unless `allowInProgressTurn` is true. pub expected_turn_id: String, + /// Allow sampling the latest turn before it finishes. + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub allow_in_progress_turn: bool, pub prompt: String, /// Applied only when the provider and model support cache-preserving reasoning-effort updates. #[ts(optional = nullable)] diff --git a/codex-rs/app-server/src/thread_sample.rs b/codex-rs/app-server/src/thread_sample.rs index fe96b92cf3..048d9e6c6e 100644 --- a/codex-rs/app-server/src/thread_sample.rs +++ b/codex-rs/app-server/src/thread_sample.rs @@ -23,6 +23,8 @@ use codex_otel::auth_storage::AuthStorageOriginator; use codex_protocol::models::ContentItem; use codex_protocol::models::MessagePhase; use codex_protocol::models::ResponseItem; +use codex_protocol::protocol::AdditionalContextEntry; +use codex_protocol::protocol::AdditionalContextKind; use codex_protocol::protocol::AgentStatus; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::SessionSource; @@ -56,6 +58,11 @@ const SAMPLE_TIMEOUT: Duration = Duration::from_secs(30); // Bytes conservatively bound tokens, including for byte-fallback tokenizers. const MAX_PROMPT_BYTES: usize = 4_000; const MAX_OUTPUT_BYTES: usize = 64 * 1024; +const THREAD_SNAPSHOT_NOTE: &str = concat!( + "\nThis snapshot is a copy of the source thread’s history taken while its turn was still running. ", + "Making the copy did not stop the source turn. ", + "A tool call that was pending at capture may appear as \"aborted\" in this copy.\n", +); type SampleResult = Result; type SampleSender = Arc>>>; @@ -114,7 +121,7 @@ impl SampleTask { cancellation, } = self; let id = parent.startup_metadata().thread_id; - if parent.agent_status().await == AgentStatus::Running { + if !params.allow_in_progress_turn && parent.agent_status().await == AgentStatus::Running { return Err(invalid_request("sample source must be idle")); } let snapshot = parent.config_snapshot().await; @@ -136,16 +143,12 @@ impl SampleTask { }) .await .map_err(|error| invalid_request(error.to_string()))?; - let matches_source = match snapshot.history_mode { + let latest_turn_in_progress = match snapshot.history_mode { ThreadHistoryMode::Legacy => build_legacy_api_turns_from_rollout_items(&history.items) .last() - .is_some_and(|turn| { - turn.id == params.expected_turn_id - && matches!( - turn.status, - TurnStatus::Completed | TurnStatus::Failed | TurnStatus::Interrupted - ) - }), + .filter(|turn| turn.id == params.expected_turn_id) + .map(|turn| turn.status == TurnStatus::InProgress), + // Paginated context can omit the turn start. Use the status observed at admission. ThreadHistoryMode::Paginated => store .list_turns(ListTurnsParams { thread_id: id, @@ -159,21 +162,18 @@ impl SampleTask { .map_err(|error| invalid_request(error.to_string()))? .turns .first() - .is_some_and(|turn| { - turn.turn_id == params.expected_turn_id - && matches!( - turn.status, - StoredTurnStatus::Completed - | StoredTurnStatus::Failed - | StoredTurnStatus::Interrupted - ) - }), + .filter(|turn| turn.turn_id == params.expected_turn_id) + .map(|turn| turn.status == StoredTurnStatus::InProgress), + }; + let Some(source_in_progress) = + latest_turn_in_progress.filter(|running| !running || params.allow_in_progress_turn) + else { + return Err(invalid_request(if params.allow_in_progress_turn { + "sample source turn must be the latest turn" + } else { + "sample source turn must be the latest finished turn" + })); }; - if !matches_source { - return Err(invalid_request( - "sample source turn must be the latest finished turn", - )); - } let (mut options, settings) = Box::pin(async { manager .fork_options_from_parent(id) @@ -190,6 +190,11 @@ impl SampleTask { ) .await; options.config.ephemeral = true; + if params.allow_in_progress_turn { + // Mid-turn copy receives a custom interruption marker THREAD_SNAPSHOT_NOTE + // instead of default interrupt message + options.config.agent_interrupt_message_enabled = false; + } options .config .features @@ -259,7 +264,7 @@ impl SampleTask { Box::pin(fork.thread.restore_thread_settings(*settings)) .await .map_err(|error| internal_error(error.to_string()))?; - let input = TurnInputRequest::user_input(vec![UserInput::Text { + let mut input = TurnInputRequest::user_input(vec![UserInput::Text { text: params.prompt, text_elements: Vec::new(), }]) @@ -273,6 +278,15 @@ impl SampleTask { }) .with_responses_metadata(params.responsesapi_client_metadata) .with_trace(trace); + if source_in_progress { + input.additional_context.insert( + "thread_snapshot".to_string(), + AdditionalContextEntry::new( + THREAD_SNAPSHOT_NOTE.to_string(), + AdditionalContextKind::Application, + ), + ); + } if !matches!( fork.thread .start_turn_if_idle(input) @@ -331,11 +345,12 @@ impl SampleTask { .remove(&fork.thread_id.to_string()); let mut response = result?; cleanup.map_err(|error| internal_error(error.to_string()))?; - if parent - .conversation_history_snapshot() - .await - .history_version() - != revision + if !params.allow_in_progress_turn + && parent + .conversation_history_snapshot() + .await + .history_version() + != revision { return Err(invalid_request("sample source changed")); } diff --git a/codex-rs/app-server/tests/suite/v2/snapshots/all__suite__v2__thread_sample__sample_compacted_parent_with_pending_tool_call.snap b/codex-rs/app-server/tests/suite/v2/snapshots/all__suite__v2__thread_sample__sample_compacted_parent_with_pending_tool_call.snap new file mode 100644 index 0000000000..65b4e210e8 --- /dev/null +++ b/codex-rs/app-server/tests/suite/v2/snapshots/all__suite__v2__thread_sample__sample_compacted_parent_with_pending_tool_call.snap @@ -0,0 +1,74 @@ +--- +source: app-server/tests/suite/v2/thread_sample.rs +expression: "context_snapshot::format_context_snapshot(\"A child samples a compacted parent while a tool call is pending. The parent then receives the real tool result.\",\n&entries, &ContextSnapshotOptions::default().rewrite_known_segments(),)" +--- +Scenario: A child samples a compacted parent while a tool call is pending. The parent then receives the real tool result. + +## Window 1 +-- request 1 (request; Parent starts) -- +00:message/developer: + Help the user with their task. +01:message/developer[2]: + [01] + [02] +02:message/user: + +03:message/user: + Initial turn + +## Window 2 (after request 1: settings changed (parallel_tool_calls, tools)) +-- request 2 (request; Parent compacts) -- +00:message/developer: + Help the user with their task. +01:message/developer[2]: + [01] + [02] +02:message/user: + +03:message/user: + Initial turn +04:function_call/unknown_tool:{} +05:function_call_output:unsupported call: unknown_tool +06:message/user: + + +## Window 3 (after request 2: input diverged at item 04, settings changed (parallel_tool_calls, tools)) +-- request 3 (request; Parent requests a tool) -- +00:message/developer: + Help the user with their task. +01:message/developer[2]: + [01] + [02] +02:message/user: + +03:message/user: + Initial turn +04:message/user: + + Mid-turn compaction summary +-- request 4 (request; Child samples the open turn) -- +05:function_call/parent_tool:{} +06:function_call_output:aborted +07:message/developer: + + This snapshot is a copy of the source thread’s history taken while its turn was still running. Making th...ding at capture may appear as "aborted" in this copy. [hash=2c3f1d1da57b6493] + +08:message/user: + Predict + +## Window 4 (after request 4: input diverged at item 06) +-- request 5 (request; Parent receives the tool result) -- +00:message/developer: + Help the user with their task. +01:message/developer[2]: + [01] + [02] +02:message/user: + +03:message/user: + Initial turn +04:message/user: + + Mid-turn compaction summary +05:function_call/parent_tool:{} +06:function_call_output:parent tool result diff --git a/codex-rs/app-server/tests/suite/v2/thread_sample.rs b/codex-rs/app-server/tests/suite/v2/thread_sample.rs index 0d72d343cf..bfa40673c5 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_sample.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_sample.rs @@ -35,9 +35,12 @@ use codex_protocol::openai_models::ModelServiceTier; use codex_protocol::openai_models::ReasoningEffort; use core_test_support::context_snapshot; use core_test_support::context_snapshot::ContextSnapshotOptions; +use core_test_support::context_snapshot::SnapshotEntry; use core_test_support::responses; use core_test_support::skip_if_host_windows; use core_test_support::skip_if_remote; +use core_test_support::streaming_sse::StreamingSseChunk; +use core_test_support::streaming_sse::start_streaming_sse_server; use pretty_assertions::assert_eq; use serde_json::Value; use serde_json::json; @@ -172,12 +175,14 @@ async fn prediction_fork(client: &mut TestAppServer, parent: &str) -> Result Result<()> { let server = responses::start_mock_server().await; let compact_parent = @@ -303,6 +308,7 @@ async fn sample_preserves_context_and_keeps_source_unchanged( ThreadSampleParams { thread_id: parent.thread.id.clone(), expected_turn_id: completed.turn.id.clone(), + allow_in_progress_turn, prompt: "Continue identically".to_string(), effort: Some(ReasoningEffort::Medium), interrupt_on_cancel: Some(true), @@ -372,6 +378,13 @@ async fn sample_preserves_context_and_keeps_source_unchanged( ); } let prediction = &requests[sample_index]; + let sample_input = prediction.input(); + assert_eq!(marker_count(&sample_input, ""), 0); + assert_eq!( + marker_count(&sample_input, ""), + usize::from(parent_status == TurnStatus::Interrupted), + "only the parent's real interruption marker should be inherited" + ); assert_eq!( comparable_request(prediction.body_json()), comparable_request(requests[sample_index + 1].body_json()) @@ -411,6 +424,329 @@ async fn sample_preserves_context_and_keeps_source_unchanged( Ok(()) } +fn marker_count(input: &[Value], marker: &str) -> usize { + input + .iter() + .filter_map(|item| item["content"].as_array()) + .flatten() + .filter_map(|part| part["text"].as_str()) + .map(|text| text.matches(marker).count()) + .sum() +} + +fn assert_open_sample_input(input: &[Value]) { + assert_eq!( + ( + marker_count(input, ""), + marker_count(input, ""), + ), + (1, 0), + "expected one snapshot note and no interruption marker" + ); + let [.., note, prompt] = input else { + panic!("expected the snapshot note immediately before the sample prompt"); + }; + assert_eq!(note["role"], "developer"); + let content = note["content"].as_array().expect("developer content"); + assert_eq!( + content.last().expect("snapshot note")["text"], + concat!( + "\n", + "This snapshot is a copy of the source thread’s history taken while its turn was still running. ", + "Making the copy did not stop the source turn. ", + "A tool call that was pending at capture may appear as \"aborted\" in this copy.\n", + "" + ) + ); + assert_eq!(prompt["role"], "user"); + assert_eq!(prompt["content"][0]["text"], "Predict"); +} + +#[test_case(ThreadHistoryMode::Legacy, false; "legacy")] +#[test_case(ThreadHistoryMode::Paginated, false; "paginated")] +#[test_case(ThreadHistoryMode::Paginated, true; "paginated_after_mid_turn_compaction")] +#[tokio::test] +async fn sample_in_progress_turn_uses_captured_history( + history_mode: ThreadHistoryMode, + compact_mid_turn: bool, +) -> Result<()> { + let (release_sample, sample_gate) = tokio::sync::oneshot::channel(); + let mut streams = Vec::new(); + if compact_mid_turn { + // A tool call keeps the turn open while its token usage triggers compaction. + streams.push(vec![StreamingSseChunk { + gate: None, + body: responses::sse(vec![ + responses::ev_function_call("before-compaction", "unknown_tool", "{}"), + responses::ev_completed_with_tokens( + "before-compaction", + /*total_tokens*/ 330_000, + ), + ]), + }]); + streams.push(vec![StreamingSseChunk { + gate: None, + body: responses::sse(vec![ + responses::ev_assistant_message("summary", "Mid-turn compaction summary"), + responses::ev_completed_with_tokens("compacted", /*total_tokens*/ 200), + ]), + }]); + } + let child_request_index = streams.len() + 1; + streams.extend([ + vec![StreamingSseChunk { + gate: None, + body: responses::sse(vec![ + responses::ev_function_call("pending-call", "parent_tool", "{}"), + responses::ev_completed("parent"), + ]), + }], + vec![StreamingSseChunk { + gate: Some(sample_gate), + body: responses::sse(vec![ + responses::ev_assistant_message("prediction", "Continue"), + responses::ev_completed("sample"), + ]), + }], + vec![StreamingSseChunk { + gate: None, + body: responses::sse(vec![responses::ev_completed("parent-done")]), + }], + ]); + let (server, _) = start_streaming_sse_server(streams).await; + let home = TempDir::new()?; + MockResponsesConfig::new(server.uri()) + .with_root_config("model_auto_compact_token_limit = 200000") + .write(home.path())?; + let mut client = TestAppServer::builder() + .with_codex_home(home.path()) + .build_initialized_with_timeout(WAIT) + .await?; + let parent = client + .start_thread(ThreadStartParams { + history_mode: Some(history_mode), + base_instructions: compact_mid_turn + .then(|| "Help the user with their task.".to_string()), + dynamic_tools: Some(vec![DynamicToolSpec::Function(DynamicToolFunctionSpec { + name: "parent_tool".to_string(), + description: "Wait for the client.".to_string(), + input_schema: json!({"type":"object","properties":{}}), + defer_loading: false, + })]), + ..Default::default() + }) + .await? + .thread + .id; + let start = client + .send_turn_start_request(TurnStartParams { + thread_id: parent.clone(), + input: input("Initial turn"), + ..Default::default() + }) + .await?; + let started: TurnStartResponse = timeout(WAIT, client.read_response(start)).await??; + let ServerRequest::DynamicToolCall { request_id, .. } = + timeout(WAIT, client.read_stream_until_request_message()).await?? + else { + anyhow::bail!("expected the parent's pending tool call"); + }; + let sample_id = client + .send_raw_request( + "thread/sample", + Some(json!({ + "threadId": parent, "expectedTurnId": started.turn.id, + "allowInProgressTurn": true, "prompt": "Predict" + })), + ) + .await?; + timeout(WAIT, server.wait_for_request_count(child_request_index + 1)).await?; + // Finish the parent after the child has captured the pending call. + client + .send_response( + request_id, + json!({ + "contentItems": [{"type":"inputText","text":"parent tool result"}], + "success": true + }), + ) + .await?; + let completed: TurnCompletedNotification = + timeout(WAIT, client.read_notification("turn/completed")).await??; + assert_eq!(completed.turn.id, started.turn.id); + assert_eq!(completed.turn.status, TurnStatus::Completed); + + let requests = server + .requests() + .await + .iter() + .map(|request| serde_json::from_slice::(request)) + .collect::>>()?; + if compact_mid_turn { + let [initial, compaction, continuation, sample, follow_up] = requests.as_slice() else { + anyhow::bail!("expected five model requests, got {}", requests.len()); + }; + let entries = [ + SnapshotEntry::body(initial).labeled("Parent starts"), + SnapshotEntry::body(compaction).labeled("Parent compacts"), + SnapshotEntry::body(continuation).labeled("Parent requests a tool"), + SnapshotEntry::body(sample).labeled("Child samples the open turn"), + SnapshotEntry::body(follow_up).labeled("Parent receives the tool result"), + ]; + insta::assert_snapshot!( + "sample_compacted_parent_with_pending_tool_call", + context_snapshot::format_context_snapshot( + "A child samples a compacted parent while a tool call is pending. The parent then receives the real tool result.", + &entries, + &ContextSnapshotOptions::default().rewrite_known_segments(), + ) + ); + } + let child = &requests[child_request_index]; + let follow_up = &requests[child_request_index + 1]; + let child_input = child["input"].as_array().expect("sample input"); + let parent_input = follow_up["input"].as_array().expect("parent input"); + assert_eq!(marker_count(parent_input, ""), 0); + assert_eq!(marker_count(parent_input, ""), 0); + for (input, expected_output) in [ + (parent_input, "parent tool result"), + (child_input, "aborted"), + ] { + let call = input + .iter() + .position(|item| item["type"] == "function_call" && item["call_id"] == "pending-call") + .expect("captured call"); + assert_eq!(input[call]["name"], "parent_tool"); + if compact_mid_turn { + assert!(input[..call].iter().any(|item| { + item["content"][0]["text"] + .as_str() + .is_some_and(|text| text.contains("Mid-turn compaction summary")) + })); + } + let output = input + .iter() + .position(|item| { + item["type"] == "function_call_output" && item["call_id"] == "pending-call" + }) + .expect("paired output"); + assert!(call < output); + assert_eq!(input[output]["output"], expected_output); + } + assert_open_sample_input(child_input); + release_sample.send(()).expect("sample stream waiting"); + let response: ThreadSampleResponse = timeout(WAIT, client.read_response(sample_id)).await??; + assert_eq!(response.source_turn_id, started.turn.id); + assert_eq!(response.text.as_deref(), Some("Continue")); + Ok(()) +} + +#[tokio::test] +async fn sample_after_mid_turn_compaction_without_a_pending_tool_call() -> Result<()> { + let (release_parent, parent_gate) = tokio::sync::oneshot::channel(); + let mut first_response = + responses::ev_completed_with_tokens("before-compaction", /*total_tokens*/ 330_000); + first_response["response"]["end_turn"] = json!(false); + let stream = |gate, events| { + vec![StreamingSseChunk { + gate, + body: responses::sse(events), + }] + }; + let (server, _) = start_streaming_sse_server(vec![ + stream(None, vec![first_response]), + stream( + None, + vec![ + responses::ev_assistant_message("summary", "Mid-turn compaction summary"), + responses::ev_completed_with_tokens("compacted", /*total_tokens*/ 200), + ], + ), + stream( + Some(parent_gate), + vec![responses::ev_completed("parent-done")], + ), + stream( + None, + vec![ + responses::ev_assistant_message("prediction", "Continue"), + responses::ev_completed("sample"), + ], + ), + ]) + .await; + let home = TempDir::new()?; + MockResponsesConfig::new(server.uri()) + .with_root_config("model_auto_compact_token_limit = 200000") + .write(home.path())?; + let mut client = TestAppServer::builder() + .with_codex_home(home.path()) + .build_initialized_with_timeout(WAIT) + .await?; + let parent = client + .start_thread(ThreadStartParams { + history_mode: Some(ThreadHistoryMode::Paginated), + ..Default::default() + }) + .await? + .thread; + let start = client + .send_turn_start_request(TurnStartParams { + thread_id: parent.id.clone(), + input: input("Initial turn"), + ..Default::default() + }) + .await?; + let started: TurnStartResponse = timeout(WAIT, client.read_response(start)).await??; + + // The third request is the parent's first model step after compaction. + // Its response is gated, so the turn is still running when the child samples it. + timeout(WAIT, server.wait_for_request_count(/*count*/ 3)).await?; + let path = parent.path.as_ref().expect("persistent source path"); + let before = std::fs::read(path)?; + let response = sample( + &mut client, + ThreadSampleParams { + thread_id: parent.id, + expected_turn_id: started.turn.id.clone(), + allow_in_progress_turn: true, + prompt: "Predict".to_string(), + ..Default::default() + }, + ) + .await?; + assert_eq!( + std::fs::read(path)?, + before, + "sampling changed the parent rollout" + ); + assert_eq!(response.source_turn_id, started.turn.id); + release_parent.send(()).expect("parent response waiting"); + let completed: TurnCompletedNotification = + timeout(WAIT, client.read_notification("turn/completed")).await??; + assert_eq!(completed.turn.id, started.turn.id); + assert_eq!(completed.turn.status, TurnStatus::Completed); + + let requests = server.requests().await; + assert_eq!(requests.len(), 4); + let child: Value = serde_json::from_slice(&requests[3])?; + let input = child["input"].as_array().expect("sample input"); + assert!( + input.iter().any(|item| { + item["content"][0]["text"] + .as_str() + .is_some_and(|text| text.contains("Mid-turn compaction summary")) + }), + "expected the inherited compaction summary" + ); + assert!( + input.iter().all(|item| item.get("call_id").is_none()), + "unexpected tool call or output" + ); + assert_open_sample_input(input); + Ok(()) +} + #[test_case(true, ReasoningEffort::High, ReasoningEffort::Low, ReasoningEffort::Medium; "pinned parent")] #[test_case(false, ReasoningEffort::High, ReasoningEffort::Low, ReasoningEffort::Medium; "parent overrides disabled")] #[test_case(true, ReasoningEffort::Low, ReasoningEffort::Medium, ReasoningEffort::High; "increasing pinned effort")]