diff --git a/codex-rs/core/src/session/inject.rs b/codex-rs/core/src/session/inject.rs index b2d0c05673..50b306a774 100644 --- a/codex-rs/core/src/session/inject.rs +++ b/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>( &self, - input: Vec, - ) -> Result<(), Vec> { + input: Vec, + ) -> Result<(), Vec> { 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(), ) diff --git a/codex-rs/core/src/session/notification_tests.rs b/codex-rs/core/src/session/notification_tests.rs new file mode 100644 index 0000000000..a33bb9f60b --- /dev/null +++ b/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)), + ); +} diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 0152b7b043..56a00c0d61 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/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; diff --git a/codex-rs/core/src/tools/code_mode/delegate.rs b/codex-rs/core/src/tools/code_mode/delegate.rs index 607fd086ff..384aa94894 100644 --- a/codex-rs/core/src/tools/code_mode/delegate.rs +++ b/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>, }, @@ -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(|_| { diff --git a/codex-rs/core/tests/suite/step_settings.rs b/codex-rs/core/tests/suite/step_settings.rs index a0bc4064f0..33ec639d4e 100644 --- a/codex-rs/core/tests/suite/step_settings.rs +++ b/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"; diff --git a/codex-rs/core/tests/suite/step_settings/code_mode_notifications.rs b/codex-rs/core/tests/suite/step_settings/code_mode_notifications.rs new file mode 100644 index 0000000000..3efda75b4f --- /dev/null +++ b/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::>(), + [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::>(), + [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::>(); + assert_eq!(live, expected); + assert_eq!( + live.iter() + .map(|item| item["call_id"].as_str().unwrap()) + .collect::>(), + 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::>(); + 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") + }) +}