From 98462cbf15b6fbaa902b66d23e0395176ea46a05 Mon Sep 17 00:00:00 2001 From: Itay Golan Date: Fri, 24 Apr 2026 12:24:15 -0700 Subject: [PATCH] better windowing control --- codex-rs/core/src/client.rs | 12 +- codex-rs/core/src/client_common.rs | 14 ++ codex-rs/core/src/codex.rs | 122 ++++++++++++++--- codex-rs/core/src/codex_tests.rs | 88 ++++++++++++ codex-rs/core/src/stream_events_utils.rs | 15 ++ codex-rs/core/tests/suite/compact.rs | 166 +++++++++++++++++++++++ codex-rs/tools/src/reflections_tool.rs | 2 +- 7 files changed, 399 insertions(+), 20 deletions(-) diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 16f743943a..989ebc85b8 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -1580,7 +1580,7 @@ where let (tx_event, rx_event) = mpsc::channel::>(1600); let (tx_last_response, rx_last_response) = oneshot::channel::(); - tokio::spawn(async move { + let task = tokio::spawn(async move { let mut logged_error = false; let mut tx_last_response = Some(tx_last_response); let mut items_added: Vec = Vec::new(); @@ -1645,8 +1645,16 @@ where } } }); + let abort_handle = task.abort_handle(); + std::mem::drop(task); - (ResponseStream { rx_event }, rx_last_response) + ( + ResponseStream { + rx_event, + abort_handle, + }, + rx_last_response, + ) } /// Handles a 401 response by optionally refreshing ChatGPT tokens once. diff --git a/codex-rs/core/src/client_common.rs b/codex-rs/core/src/client_common.rs index c1dfa29edc..a956968f45 100644 --- a/codex-rs/core/src/client_common.rs +++ b/codex-rs/core/src/client_common.rs @@ -158,6 +158,20 @@ fn strip_total_output_header(output: &str) -> Option<(&str, u32)> { pub struct ResponseStream { pub(crate) rx_event: mpsc::Receiver>, + pub(crate) abort_handle: tokio::task::AbortHandle, +} + +impl ResponseStream { + pub(crate) fn close(&mut self) { + self.abort_handle.abort(); + self.rx_event.close(); + } +} + +impl Drop for ResponseStream { + fn drop(&mut self) { + self.abort_handle.abort(); + } } impl Stream for ResponseStream { diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index c6c5529a77..ba5592b5ff 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -39,6 +39,7 @@ use crate::rollout::find_thread_name_by_id; use crate::session_prefix::format_subagent_notification_message; use crate::skills_load_input_from_config; use crate::stream_events_utils::HandleOutputCtx; +use crate::stream_events_utils::ResponseTerminalControl; use crate::stream_events_utils::handle_non_tool_response_item; use crate::stream_events_utils::handle_output_item_done; use crate::stream_events_utils::last_assistant_message_from_item; @@ -344,6 +345,8 @@ use codex_protocol::config_types::ServiceTier; use codex_protocol::config_types::WindowsSandboxLevel; use codex_protocol::models::ContentItem; use codex_protocol::models::DeveloperInstructions; +use codex_protocol::models::FunctionCallOutputBody; +use codex_protocol::models::FunctionCallOutputPayload; use codex_protocol::models::ResponseInputItem; use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig; @@ -7995,15 +7998,23 @@ async fn drain_in_flight( in_flight: &mut FuturesOrdered>>, sess: Arc, turn_context: Arc, -) -> CodexResult<()> { - let mut reset_reflections_context_window = false; +) -> CodexResult { while let Some(res) = in_flight.next().await { match res { Ok(response_input) => { sess.record_conversation_items(&turn_context, &[response_input.into()]) .await; if sess.take_reflections_context_window_reset_request() { - reset_reflections_context_window = true; + crate::reflections::run_reflections_compact_task_inner( + Arc::clone(&sess), + Arc::clone(&turn_context), + InitialContextInjection::DoNotInject, + CompactionTrigger::Manual, + CompactionReason::UserRequested, + CompactionPhase::MidTurn, + ) + .await?; + return Ok(true); } } Err(err) => { @@ -8011,18 +8022,63 @@ async fn drain_in_flight( } } } - if reset_reflections_context_window { - crate::reflections::run_reflections_compact_task_inner( - Arc::clone(&sess), - Arc::clone(&turn_context), - InitialContextInjection::DoNotInject, - CompactionTrigger::Manual, - CompactionReason::UserRequested, - CompactionPhase::MidTurn, - ) - .await?; + Ok(false) +} + +async fn run_reflections_terminal_reset( + tool_future: BoxFuture<'static, CodexResult>, + sess: Arc, + turn_context: Arc, +) -> CodexResult { + let response_input = tool_future.await?; + + if !sess.take_reflections_context_window_reset_request() { + sess.record_conversation_items(&turn_context, &[response_input.into()]) + .await; + return Ok(false); + } + + match crate::reflections::run_reflections_compact_task_inner( + Arc::clone(&sess), + Arc::clone(&turn_context), + InitialContextInjection::DoNotInject, + CompactionTrigger::Manual, + CompactionReason::UserRequested, + CompactionPhase::MidTurn, + ) + .await + { + Ok(()) => Ok(true), + Err(err) => { + error!("reflections_new_context_window compaction failed: {err}"); + let failure_output = reflections_terminal_reset_failure_output(response_input, &err); + sess.record_conversation_items(&turn_context, &[failure_output.into()]) + .await; + Ok(false) + } + } +} + +fn reflections_terminal_reset_failure_output( + response_input: ResponseInputItem, + err: &CodexErr, +) -> ResponseInputItem { + let call_id = match response_input { + ResponseInputItem::FunctionCallOutput { call_id, .. } + | ResponseInputItem::CustomToolCallOutput { call_id, .. } + | ResponseInputItem::McpToolCallOutput { call_id, .. } + | ResponseInputItem::ToolSearchOutput { call_id, .. } => call_id, + ResponseInputItem::Message { .. } => String::new(), + }; + ResponseInputItem::FunctionCallOutput { + call_id, + output: FunctionCallOutputPayload { + body: FunctionCallOutputBody::Text(format!( + "Failed to start a fresh Reflections context window: {err}" + )), + success: Some(false), + }, } - Ok(()) } #[allow(clippy::too_many_arguments)] @@ -8156,8 +8212,37 @@ async fn try_run_sampling_request( Ok(output_result) => output_result, Err(err) => break Err(err), }; - if let Some(tool_future) = output_result.tool_future { - in_flight.push_back(tool_future); + match output_result.terminal_control { + Some(ResponseTerminalControl::ReflectionsContextWindowReset) => { + in_flight = FuturesOrdered::new(); + stream.close(); + client_session.reset_websocket_session(); + let reset_succeeded = if let Some(tool_future) = output_result.tool_future { + run_reflections_terminal_reset( + tool_future, + Arc::clone(&sess), + Arc::clone(&turn_context), + ) + .await? + } else { + false + }; + if reset_succeeded { + break Ok(SamplingRequestResult { + needs_follow_up: true, + last_agent_message, + }); + } + break Ok(SamplingRequestResult { + needs_follow_up: true, + last_agent_message, + }); + } + None => { + if let Some(tool_future) = output_result.tool_future { + in_flight.push_back(tool_future); + } + } } if let Some(agent_message) = output_result.last_agent_message { last_agent_message = Some(agent_message); @@ -8351,7 +8436,10 @@ async fn try_run_sampling_request( ) .await; - drain_in_flight(&mut in_flight, sess.clone(), turn_context.clone()).await?; + let compacted = drain_in_flight(&mut in_flight, sess.clone(), turn_context.clone()).await?; + if compacted { + client_session.reset_websocket_session(); + } if cancellation_token.is_cancelled() { return Err(CodexErr::TurnAborted); diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index 239791b99c..bf4e42a662 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -5305,6 +5305,94 @@ async fn tool_calls_reopen_mailbox_delivery_for_current_turn() { ); } +#[tokio::test] +async fn reflections_new_context_window_marks_terminal_control() { + let (session, mut turn_context) = make_session_and_context().await; + turn_context.tools_config.reflections = true; + let sess = Arc::new(session); + let tc = Arc::new(turn_context); + + let item = ResponseItem::FunctionCall { + id: None, + name: codex_tools::REFLECTIONS_NEW_CONTEXT_WINDOW_TOOL_NAME.to_string(), + namespace: None, + arguments: "{}".to_string(), + call_id: "call-reflections-reset".to_string(), + }; + let mut ctx = HandleOutputCtx { + sess: Arc::clone(&sess), + turn_context: Arc::clone(&tc), + tool_runtime: test_tool_runtime(Arc::clone(&sess), Arc::clone(&tc)), + cancellation_token: CancellationToken::new(), + }; + + let output = handle_output_item_done(&mut ctx, item, /*previously_active_item*/ None) + .await + .expect("reflections reset tool call should be handled"); + + assert!(output.needs_follow_up); + assert!(output.tool_future.is_some()); + assert_eq!( + output.terminal_control, + Some(crate::stream_events_utils::ResponseTerminalControl::ReflectionsContextWindowReset) + ); +} + +#[tokio::test] +async fn reflections_terminal_reset_compaction_failure_allows_same_context_continuation() { + let (sess, tc, rx) = make_session_and_context_with_rx().await; + sess.request_reflections_context_window_reset(); + + let tool_future: BoxFuture<'static, CodexResult> = Box::pin(async { + Ok(ResponseInputItem::FunctionCallOutput { + call_id: "call-reflections-reset".to_string(), + output: FunctionCallOutputPayload { + body: FunctionCallOutputBody::Text( + "A fresh Reflections context window will start after this tool result is recorded." + .to_string(), + ), + success: Some(true), + }, + }) + }); + + let compacted = run_reflections_terminal_reset(tool_future, Arc::clone(&sess), Arc::clone(&tc)) + .await + .expect("compaction failure should be surfaced without failing the sampling loop"); + + assert!(!compacted); + assert!(!sess.take_reflections_context_window_reset_request()); + let history = sess.clone_history().await; + assert!( + history.raw_items().iter().any(|item| { + let ResponseItem::FunctionCallOutput { call_id, output } = item else { + return false; + }; + call_id == "call-reflections-reset" + && output.success == Some(false) + && matches!( + &output.body, + FunctionCallOutputBody::Text(text) + if text.contains("Failed to start a fresh Reflections context window") + ) + }), + "compaction failure should be recorded as a failed tool output" + ); + tokio::time::timeout(StdDuration::from_secs(2), async { + loop { + let event = rx + .recv() + .await + .expect("expected reflections compaction error event"); + if matches!(event.msg, EventMsg::Error(_)) { + break; + } + } + }) + .await + .expect("timed out waiting for reflections compaction error event"); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn abort_review_task_emits_exited_then_aborted_and_records_history() { let (sess, tc, rx) = make_session_and_context_with_rx().await; diff --git a/codex-rs/core/src/stream_events_utils.rs b/codex-rs/core/src/stream_events_utils.rs index d491bb6adb..e11d1d52bb 100644 --- a/codex-rs/core/src/stream_events_utils.rs +++ b/codex-rs/core/src/stream_events_utils.rs @@ -5,6 +5,7 @@ use base64::Engine; use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use codex_protocol::config_types::ModeKind; use codex_protocol::items::TurnItem; +use codex_tools::REFLECTIONS_NEW_CONTEXT_WINDOW_TOOL_NAME; use codex_utils_stream_parser::strip_citations; use tokio_util::sync::CancellationToken; @@ -191,6 +192,12 @@ pub(crate) struct OutputItemResult { pub last_agent_message: Option, pub needs_follow_up: bool, pub tool_future: Option>, + pub terminal_control: Option, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum ResponseTerminalControl { + ReflectionsContextWindowReset, } pub(crate) struct HandleOutputCtx { @@ -227,6 +234,13 @@ pub(crate) async fn handle_output_item_done( record_completed_response_item(ctx.sess.as_ref(), ctx.turn_context.as_ref(), &item) .await; + let terminal_control = if call.tool_name.namespace.is_none() + && call.tool_name.name.as_str() == REFLECTIONS_NEW_CONTEXT_WINDOW_TOOL_NAME + { + Some(ResponseTerminalControl::ReflectionsContextWindowReset) + } else { + None + }; let cancellation_token = ctx.cancellation_token.child_token(); let tool_future: InFlightFuture<'static> = Box::pin( ctx.tool_runtime @@ -235,6 +249,7 @@ pub(crate) async fn handle_output_item_done( ); output.needs_follow_up = true; + output.terminal_control = terminal_control; output.tool_future = Some(tool_future); } // No tool call: convert messages/reasoning into turn items and mark them as complete. diff --git a/codex-rs/core/tests/suite/compact.rs b/codex-rs/core/tests/suite/compact.rs index 8bf82d54d3..b354f6e84f 100644 --- a/codex-rs/core/tests/suite/compact.rs +++ b/codex-rs/core/tests/suite/compact.rs @@ -21,13 +21,19 @@ use codex_protocol::protocol::RolloutLine; use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::WarningEvent; use codex_protocol::user_input::UserInput; +use codex_tools::REFLECTIONS_NEW_CONTEXT_WINDOW_TOOL_NAME; use core_test_support::context_snapshot; use core_test_support::context_snapshot::ContextSnapshotOptions; use core_test_support::context_snapshot::ContextSnapshotRenderMode; use core_test_support::responses::ev_local_shell_call; +use core_test_support::responses::ev_message_item_added; +use core_test_support::responses::ev_output_text_delta; use core_test_support::responses::ev_reasoning_item; +use core_test_support::responses::ev_response_created; use core_test_support::responses::mount_models_once; use core_test_support::skip_if_no_network; +use core_test_support::streaming_sse::StreamingSseChunk; +use core_test_support::streaming_sse::start_streaming_sse_server; use core_test_support::test_codex::test_codex; use core_test_support::wait_for_event; use core_test_support::wait_for_event_match; @@ -494,6 +500,166 @@ async fn reflections_manual_compact_writes_sidecar_without_summarizer_request() assert!(!rollout.contains("logs/cw00000/transcript.md")); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn reflections_new_context_window_skips_sibling_tools_and_resumes_fresh() { + skip_if_no_network!(); + + const AFTER_RESET_TEXT: &str = "SHOULD_NOT_RECORD_AFTER_RESET"; + const AFTER_RESET_NOTE: &str = "call-note-after-reset"; + const LIST_LOGS_CALL: &str = "call-list-logs"; + + let (after_reset_gate_tx, after_reset_gate_rx) = tokio::sync::oneshot::channel(); + let first_turn = sse(vec![ + ev_response_created("r1"), + ev_function_call( + "call-reflections-reset", + REFLECTIONS_NEW_CONTEXT_WINDOW_TOOL_NAME, + "{}", + ), + ]); + let after_reset_chunk = sse(vec![ + ev_message_item_added("msg-empty-after-reset", ""), + ev_output_text_delta(AFTER_RESET_TEXT), + ev_assistant_message("msg-after-reset", AFTER_RESET_TEXT), + ev_function_call( + AFTER_RESET_NOTE, + "reflections_write_note", + r#"{"note_id":"dangling","operation":"append","content":"DANGLING_NOTE"}"#, + ), + ev_function_call(DUMMY_CALL_ID, DUMMY_FUNCTION_NAME, "{}"), + ev_completed_with_tokens("r1", /*total_tokens*/ 100), + ]); + let list_logs_turn = sse(vec![ + ev_response_created("r2"), + ev_function_call( + LIST_LOGS_CALL, + "reflections_list", + r#"{"collection":"logs","start":1,"stop":200}"#, + ), + ev_completed_with_tokens("r2", /*total_tokens*/ 100), + ]); + let final_turn = sse(vec![ + ev_response_created("r3"), + ev_assistant_message("m2", "AFTER_REFLECTIONS_TOOL_RESET"), + ev_completed_with_tokens("r3", /*total_tokens*/ 100), + ]); + let (streaming_server, _completion_receivers) = start_streaming_sse_server(vec![ + vec![ + StreamingSseChunk { + gate: None, + body: first_turn, + }, + StreamingSseChunk { + gate: Some(after_reset_gate_rx), + body: after_reset_chunk, + }, + ], + vec![StreamingSseChunk { + gate: None, + body: list_logs_turn, + }], + vec![StreamingSseChunk { + gate: None, + body: final_turn, + }], + ]) + .await; + + let mut builder = test_codex().with_config(|config| { + config.compaction_strategy = CompactionStrategyConfig::Reflections; + }); + let test = builder + .build_with_streaming_server(&streaming_server) + .await + .unwrap(); + let codex = test.codex.clone(); + let rollout_path = test.session_configured.rollout_path.expect("rollout path"); + let sidecar_path = rollout_path.with_extension("reflections"); + + codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "start work and reset with reflections".into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + }) + .await + .unwrap(); + wait_for_event(&codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await; + + let _ = after_reset_gate_tx.send(()); + codex.submit(Op::Shutdown).await.unwrap(); + wait_for_event(&codex, |ev| matches!(ev, EventMsg::ShutdownComplete)).await; + + let request_bodies: Vec = streaming_server + .requests() + .await + .into_iter() + .map(|request| String::from_utf8_lossy(&request).to_string()) + .collect(); + streaming_server.shutdown().await; + assert_eq!(request_bodies.len(), 3); + assert!( + body_contains_text(&request_bodies[1], "Your context window was reset"), + "follow-up request should resume from the fresh Reflections handoff" + ); + assert!( + !request_bodies[1].contains(&format!("unsupported call: {DUMMY_FUNCTION_NAME}")), + "sibling tool calls after the reset control tool should not be executed" + ); + assert!( + request_bodies[2].contains("cw00000"), + "reflections_list(logs) should return the completed reset window" + ); + assert!( + !request_bodies[2].contains("cw00001"), + "reflections_list(logs) should not expose a partial second window" + ); + assert!( + !request_bodies[2].contains(AFTER_RESET_TEXT) + && !request_bodies[2].contains(AFTER_RESET_NOTE) + && !request_bodies[2].contains("DANGLING_NOTE"), + "reflections_list(logs) should not include events emitted after the reset tool" + ); + + let rollout = std::fs::read_to_string(&rollout_path).unwrap_or_else(|err| { + panic!( + "failed to read rollout file {}: {err}", + rollout_path.display() + ) + }); + assert!( + rollout.contains(REFLECTIONS_NEW_CONTEXT_WINDOW_TOOL_NAME), + "reset tool call should be recorded before compaction" + ); + assert!( + !rollout.contains(DUMMY_FUNCTION_NAME) + && !rollout.contains(AFTER_RESET_TEXT) + && !rollout.contains("msg-empty-after-reset") + && !rollout.contains(AFTER_RESET_NOTE) + && !rollout.contains("DANGLING_NOTE"), + "output after the reset control tool should not be recorded" + ); + + let state_path = sidecar_path.join("state.json"); + let state = std::fs::read_to_string(&state_path) + .unwrap_or_else(|err| panic!("failed to read {}: {err}", state_path.display())); + assert!(state.contains("\"latest_window\": \"cw00000\"")); + assert!(!state.contains("cw00001")); + + let transcript_path = sidecar_path.join("logs/cw00000/transcript.md"); + let transcript = std::fs::read_to_string(&transcript_path) + .unwrap_or_else(|err| panic!("failed to read {}: {err}", transcript_path.display())); + assert!( + !transcript.contains(AFTER_RESET_TEXT) + && !transcript.contains(AFTER_RESET_NOTE) + && !transcript.contains("DANGLING_NOTE"), + "transcript should stop at the reset control tool" + ); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn reflections_near_limit_reminder_is_recorded_once_and_dropped_after_compaction() { skip_if_no_network!(); diff --git a/codex-rs/tools/src/reflections_tool.rs b/codex-rs/tools/src/reflections_tool.rs index 5346516aa3..b6ef30cba2 100644 --- a/codex-rs/tools/src/reflections_tool.rs +++ b/codex-rs/tools/src/reflections_tool.rs @@ -25,7 +25,7 @@ pub fn create_reflections_new_context_window_tool( "under the Reflections notes directory" }; let mut description = format!( - "Starts a fresh context window for the same task. Use this after you have saved concise recovery notes {recovery_notes_location} when the current context is large or the next steps should continue from durable logs." + "Starts a fresh context window for the same task. Use this after you have saved concise recovery notes {recovery_notes_location} when the current context is large or the next steps should continue from durable logs. This is a control-flow tool: after the tool result is recorded, the current response stops and the next model request resumes from the Reflections handoff in a fresh context window." ); if let Some(usage_hint) = usage_hint { description.push_str("\n\n");