Preserve originating budgets for code mode notifications (#44867)

## Why

A yielded code mode cell can emit a notification after the session switches models. Its output should retain the originating model's truncation budget instead of using the receiving model's budget.

## What changed

Carry the cell's output token limit, including the serialization allowance, through notification dispatch and attach it as `history_truncation_token_limit` metadata. Allow `inject_if_running` to accept response item envelopes so queued notifications preserve this metadata in history and rollouts.

## Testing

Add regression coverage for delayed notifications across a model switch, verifying per-cell truncation budgets, full notification text in raw events, and persisted rollout metadata. Add a task-completion test that checks notification budgets are preserved alongside queued user input.

GitOrigin-RevId: 28380a039a6d41ea324401fcf1f611d81a279477

codex-rs/core/src/session/inject.rs

@@ -13,10 +13,10 @@ impl Session {
         clippy::await_holding_invalid_type,
         reason = "active turn checks and turn state updates must remain atomic"
     )]
-    pub async fn inject_if_running(
+    pub(crate) async fn inject_if_running<T: Into<ResponseItemEnvelope>>(
         &self,
-        input: Vec<ResponseItem>,
-    ) -> Result<(), Vec<ResponseItem>> {
+        input: Vec<T>,
+    ) -> Result<(), Vec<T>> {
         let mut active = self.active_turn.lock().await;
         match active.as_mut() {
             Some(active_turn) => {
@@ -25,7 +25,7 @@ impl Session {
                         active_turn.turn_state.as_ref(),
                         input
                             .into_iter()
-                            .map(ResponseItemEnvelope::new)
+                            .map(Into::into)
                             .map(PendingTurnInput::ResponseItem)
                             .collect(),
                     )

codex-rs/core/src/session/notification_tests.rs

@@ -0,0 +1,66 @@
+//! Task completion preserves a queued notification's budget alongside real user input.
+
+use super::*;
+use pretty_assertions::assert_eq;
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn task_finish_preserves_notification_budget_and_queued_user_input() {
+    let (session, turn, _rx) = make_session_and_context_with_rx().await;
+    session
+        .spawn_task(
+            Arc::clone(&turn),
+            Vec::new(),
+            NeverEndingTask {
+                kind: TaskKind::Regular,
+                listen_to_cancellation_token: false,
+            },
+        )
+        .await;
+    let text = "notification diagnostic line\n".repeat(500);
+    let notification = |text| ResponseItem::CustomToolCallOutput {
+        id: None,
+        call_id: "call-a".to_string(),
+        name: Some("exec".to_string()),
+        output: FunctionCallOutputPayload::from_text(text),
+        internal_chat_message_metadata_passthrough: None,
+    };
+    session
+        .inject_if_running(vec![ResponseItemEnvelope {
+            item: notification(text.clone()),
+            metadata: Some(CodexHarnessMetadata {
+                history_truncation_token_limit: Some(120),
+                ..Default::default()
+            }),
+        }])
+        .await
+        .unwrap();
+    let user_input = vec![UserInput::Text {
+        text: "keep this queued user input".to_string(),
+        text_elements: Vec::new(),
+    }];
+    assert!(matches!(
+        submit_steer_only(&session, user_input.clone(), &turn.sub_id).await,
+        TurnInputSubmission::Steered { .. }
+    ));
+    let mut current = turn.initial_settings.as_ref().clone();
+    Arc::make_mut(&mut current.model_info).truncation_policy =
+        codex_protocol::openai_models::TruncationPolicyConfig::tokens(/*limit*/ 400);
+    turn.current_settings.store(Arc::new(current));
+
+    session
+        .on_task_finished(Arc::clone(&turn), /*task_result*/ Ok(None))
+        .await;
+
+    let expected = vec![
+        notification(codex_utils_output_truncation::truncate_text(
+            &text,
+            codex_utils_output_truncation::TruncationPolicy::Tokens(120),
+        )),
+        session.response_item_from_user_input(user_input),
+    ];
+    let history = session.clone_history().await;
+    assert_eq!(
+        strip_response_item_ids(&strip_metadata_from_items(&raw_history_items(&history))),
+        strip_response_item_ids(&strip_metadata_from_items(&expected)),
+    );
+}

codex-rs/core/src/session/tests.rs

@@ -1,3 +1,6 @@
+#[path = "notification_tests.rs"]
+mod notification_tests;
+
 use super::mcp_refresh::McpRefresh;
 use super::step_settings::ResolvedStepSettings;
 use super::step_settings::StepSettings;

codex-rs/core/src/tools/code_mode/delegate.rs

@@ -8,10 +8,13 @@ use codex_code_mode::CodeModeNestedToolCall;
 use codex_code_mode::CodeModeSessionDelegate;
 use codex_code_mode::NotificationFuture;
 use codex_code_mode::ToolInvocationFuture;
+use codex_history::CodexHarnessMetadata;
+use codex_history::ResponseItemEnvelope;
 use codex_protocol::ResponseItemId;
 use codex_protocol::ThreadId;
 use codex_protocol::models::FunctionCallOutputPayload;
 use codex_protocol::models::ResponseItem;
+use codex_utils_output_truncation::with_serialization_allowance;
 use serde_json::Value as JsonValue;
 use tokio::sync::oneshot;
 use tokio::sync::watch;
@@ -138,6 +141,7 @@ impl CodeModeDispatchBroker {
                         call_id,
                         cell_id,
                         text,
+                        output_token_limit,
                         cancellation_token,
                         response_tx,
                     } => {
@@ -148,7 +152,8 @@ impl CodeModeDispatchBroker {
                         )
                         .await
                         {
-                            host.notify(call_id, cell_id, text).await
+                            host.notify(call_id, cell_id, text, output_token_limit)
+                                .await
                         } else {
                             remove_dispatch_gate(&dispatch_gates, &cell_id);
                             Err("code mode notification cancelled".to_string())
@@ -370,6 +375,14 @@ impl CodeModeSessionDelegate for CodeModeCellDelegate {
                     call_id,
                     cell_id,
                     text,
+                    output_token_limit: with_serialization_allowance(
+                        self.step_context
+                            .settings
+                            .model_info
+                            .truncation_policy
+                            .into(),
+                    )
+                    .token_budget(),
                     cancellation_token: cancellation_token.clone(),
                     response_tx,
                 })
@@ -403,6 +416,7 @@ enum DispatchMessage {
         call_id: String,
         cell_id: CellId,
         text: String,
+        output_token_limit: usize,
         cancellation_token: CancellationToken,
         response_tx: oneshot::Sender<Result<(), String>>,
     },
@@ -445,17 +459,29 @@ impl CoreTurnHost {
         async move { invocation?.await.map_err(|error| error.to_string()) }
     }
 
-    async fn notify(&self, call_id: String, cell_id: CellId, text: String) -> Result<(), String> {
+    async fn notify(
+        &self,
+        call_id: String,
+        cell_id: CellId,
+        text: String,
+        output_token_limit: usize,
+    ) -> Result<(), String> {
         if text.trim().is_empty() {
             return Ok(());
         }
         self.session
-            .inject_if_running(vec![ResponseItem::CustomToolCallOutput {
-                id: None,
-                call_id,
-                name: Some(PUBLIC_TOOL_NAME.to_string()),
-                output: FunctionCallOutputPayload::from_text(text),
-                internal_chat_message_metadata_passthrough: None,
+            .inject_if_running(vec![ResponseItemEnvelope {
+                item: ResponseItem::CustomToolCallOutput {
+                    id: None,
+                    call_id,
+                    name: Some(PUBLIC_TOOL_NAME.to_string()),
+                    output: FunctionCallOutputPayload::from_text(text),
+                    internal_chat_message_metadata_passthrough: None,
+                },
+                metadata: Some(CodexHarnessMetadata {
+                    history_truncation_token_limit: Some(output_token_limit),
+                    ..Default::default()
+                }),
             }])
             .await
             .map_err(|_| {

codex-rs/core/tests/suite/step_settings.rs

@@ -111,6 +111,8 @@ use test_case::test_case;
 use super::rmcp_client::remote_aware_environment_id;
 use super::rmcp_client::remote_aware_stdio_server_bin;
 
+mod code_mode_notifications;
+
 const MODEL_A: &str = "step-settings-a";
 const MODEL_B: &str = "step-settings-b";
 const MODEL_C: &str = "step-settings-c";

codex-rs/core/tests/suite/step_settings/code_mode_notifications.rs

@@ -0,0 +1,192 @@
+//! Delayed notifications retain their producing model's budget in history and rollouts.
+
+use super::*;
+use pretty_assertions::assert_eq;
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn yielded_notifications_keep_originating_budgets_after_model_switch() -> Result<()> {
+    skip_if_no_network!(Ok(()));
+    let server = start_mock_server().await;
+    let test = step_settings_test()
+        .with_config(|config| {
+            config.features.enable(Feature::CodeMode).unwrap();
+            for model in &mut config.model_catalog.as_mut().expect("models").models {
+                model.tool_mode = Some(ToolMode::CodeModeOnly);
+                model.use_responses_lite = false;
+                model.experimental_supported_tools = vec!["test_sync_tool".to_string()];
+                model.truncation_policy =
+                    TruncationPolicyConfig::tokens(if model.slug == MODEL_B { 400 } else { 100 });
+            }
+        })
+        .build_with_auto_env(&server)
+        .await?;
+    let barrier = r#"await tools.test_sync_tool({barrier: {
+        id: "notification-origin", participants: 2, timeout_ms: 60000
+    }});"#;
+    let a_text = "A_NOTIFICATION diagnostic line\n".repeat(500);
+    let b_text = "B_NOTIFICATION diagnostic line\n".repeat(500);
+    let initial = mount_sse_sequence(
+        &server,
+        vec![
+            sse(vec![
+                ev_response_created("resp-a"),
+                ev_custom_tool_call(
+                    "call-a",
+                    "exec",
+                    &format!(
+                        "// @exec: {{\"yield_time_ms\": 1}}\n{barrier}\nnotify({});",
+                        serde_json::to_string(&a_text)?,
+                    ),
+                ),
+                ev_completed("resp-a"),
+            ]),
+            sse_completed("resp-a-done"),
+        ],
+    )
+    .await;
+    test.submit_text_turn("start A and leave its cell running")
+        .await?;
+    let initial_requests = initial.requests();
+    assert_eq!(
+        initial_requests
+            .iter()
+            .map(|request| request.body_json()["model"].clone())
+            .collect::<Vec<_>>(),
+        [json!(MODEL_A), json!(MODEL_A)],
+    );
+    let (a_output, _) = initial_requests[1]
+        .custom_tool_call_output_content_and_success("call-a")
+        .expect("A's output");
+    let a_output = a_output.expect("yielded output");
+    let cell_id = a_output
+        .strip_prefix("Script running with cell ID ")
+        .and_then(|rest| rest.lines().next())
+        .expect("A's running cell id");
+    let responses = mount_sse_sequence(
+        &server,
+        vec![
+            sse(vec![
+                ev_response_created("resp-b"),
+                ev_custom_tool_call(
+                    "call-b",
+                    "exec",
+                    &format!(
+                        "// @exec: {{\"yield_time_ms\": 60000}}\n{barrier}\nnotify({});",
+                        serde_json::to_string(&b_text)?,
+                    ),
+                ),
+                ev_completed("resp-b"),
+            ]),
+            sse(vec![
+                ev_response_created("resp-wait-a"),
+                ev_function_call(
+                    "call-wait-a",
+                    "wait",
+                    &json!({
+                        "cell_id": cell_id, "yield_time_ms": 60000,
+                    })
+                    .to_string(),
+                ),
+                ev_completed("resp-wait-a"),
+            ]),
+            sse_completed("resp-b-done"),
+        ],
+    )
+    .await;
+    // B's cell releases the barrier, so A cannot notify while A is still active.
+    test.codex
+        .start_or_steer_turn(
+            TurnInputRequest::user_input(vec![UserInput::Text {
+                text: "run B and release A".to_string(),
+                text_elements: Vec::new(),
+            }])
+            .with_thread_settings(ThreadSettingsOverrides {
+                model: Some(MODEL_B.to_string()),
+                ..Default::default()
+            }),
+        )
+        .await?;
+    let mut raw_notifications = Vec::new();
+    wait_for_event(&test.codex, |event| {
+        if let EventMsg::RawResponseItem(event) = event {
+            let item = serde_json::to_value(&event.item).expect("raw item");
+            if is_notification(&item) {
+                raw_notifications.push(item);
+            }
+        }
+        matches!(event, EventMsg::TurnComplete(_))
+    })
+    .await;
+    let requests = responses.requests();
+    assert_eq!(
+        requests
+            .iter()
+            .map(|request| request.body_json()["model"].clone())
+            .collect::<Vec<_>>(),
+        [json!(MODEL_B), json!(MODEL_B), json!(MODEL_B)],
+    );
+    assert_eq!(raw_notifications.len(), 2, "both cells must notify");
+    let mut expected = Vec::new();
+    let mut expected_saved = Vec::new();
+    for raw in &raw_notifications {
+        let (text, budget) = match raw["call_id"].as_str().expect("call id") {
+            "call-a" => (&a_text, 120),
+            "call-b" => (&b_text, 480),
+            other => panic!("unexpected notification call: {other}"),
+        };
+        assert_eq!(raw["name"], "exec");
+        assert_eq!(
+            raw["output"], *text,
+            "raw events retain full notification text"
+        );
+        let mut bounded = raw.clone();
+        bounded["output"] = json!(truncate_text(text, TruncationPolicy::Tokens(budget)));
+        expected.push(bounded);
+        expected_saved.push((raw.clone(), Some(budget)));
+    }
+    let live = requests[2]
+        .inputs_of_type("custom_tool_call_output")
+        .into_iter()
+        .filter(is_notification)
+        .collect::<Vec<_>>();
+    assert_eq!(live, expected);
+    assert_eq!(
+        live.iter()
+            .map(|item| item["call_id"].as_str().unwrap())
+            .collect::<std::collections::BTreeSet<_>>(),
+        std::collections::BTreeSet::from(["call-a", "call-b"])
+    );
+
+    test.codex.shutdown_and_wait().await?;
+    let rollout_path = test.codex.rollout_path().expect("rollout path");
+    let history = codex_rollout::RolloutRecorder::get_rollout_history(&rollout_path).await?;
+    let saved = history
+        .get_rollout_items()
+        .iter()
+        .filter_map(|item| {
+            let RolloutItem::ResponseItem(envelope) = item else {
+                return None;
+            };
+            let item = serde_json::to_value(&envelope.item).expect("saved item");
+            is_notification(&item).then(|| {
+                (
+                    item,
+                    envelope
+                        .metadata
+                        .as_ref()
+                        .and_then(|metadata| metadata.history_truncation_token_limit),
+                )
+            })
+        })
+        .collect::<Vec<_>>();
+    assert_eq!(saved, expected_saved);
+
+    Ok(())
+}
+
+fn is_notification(item: &Value) -> bool {
+    item["type"] == "custom_tool_call_output"
+        && item["output"].as_str().is_some_and(|text| {
+            text.starts_with("A_NOTIFICATION") || text.starts_with("B_NOTIFICATION")
+        })
+}