better windowing control

codex-rs/core/src/client.rs

@@ -1580,7 +1580,7 @@ where
     let (tx_event, rx_event) = mpsc::channel::<Result<ResponseEvent>>(1600);
     let (tx_last_response, rx_last_response) = oneshot::channel::<LastResponse>();
 
-    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<ResponseItem> = 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.

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<Result<ResponseEvent>>,
+    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 {

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<BoxFuture<'static, CodexResult<ResponseInputItem>>>,
     sess: Arc<Session>,
     turn_context: Arc<TurnContext>,
-) -> CodexResult<()> {
-    let mut reset_reflections_context_window = false;
+) -> CodexResult<bool> {
     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<ResponseInputItem>>,
+    sess: Arc<Session>,
+    turn_context: Arc<TurnContext>,
+) -> CodexResult<bool> {
+    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);

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<ResponseInputItem>> = 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;

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<String>,
     pub needs_follow_up: bool,
     pub tool_future: Option<InFlightFuture<'static>>,
+    pub terminal_control: Option<ResponseTerminalControl>,
+}
+
+#[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.

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<String> = 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!();

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