From 14813d36167e5148e51830624d0e76241049d5e1 Mon Sep 17 00:00:00 2001 From: surajs-oai Date: Fri, 13 Nov 2026 20:24:20 +0000 Subject: [PATCH] Add configurable idle heartbeats to Code Mode gRPC streams (#61385) ## What changed - Add `code_mode.heartbeat_interval_s` to configure heartbeats after inactivity on session-event and tool-call streams. The default value of `0` disables heartbeats. - Pass the interval through session options and gRPC requests, preserving provider defaults when overriding it. - Ignore heartbeat messages in the client without admitting tool callbacks, and reject tool-call heartbeats that carry invocation data. ## Testing Add coverage for disabled heartbeats, repeated idle heartbeats and timer resets after traffic, callback handling, and configuration loading. Extend TCP integration tests to verify heartbeats on both streams and continued session operation after idle periods. GitOrigin-RevId: bf0b6d3081b5541837d07f35931e0dd3e5472c3b --- codex-rs/code-mode-host/src/grpc/mod.rs | 43 +++++++++++++- .../src/grpc/robustness_tests.rs | 12 ++++ codex-rs/code-mode-host/src/grpc/routing.rs | 1 + .../code-mode-host/src/grpc/service_tests.rs | 58 +++++++++++++++++++ codex-rs/code-mode-host/src/lib.rs | 1 + codex-rs/code-mode-host/tests/grpc.rs | 29 +++++++++- codex-rs/code-mode-host/tests/grpc_tcp.rs | 1 + codex-rs/code-mode-host/tests/grpc_timing.rs | 1 + .../src/grpc/codex.code_mode.v1.proto | 12 ++++ codex-rs/code-mode-protocol/src/session.rs | 7 +++ .../code-mode/src/grpc_session/callbacks.rs | 13 ++++- .../src/grpc_session/callbacks_tests.rs | 22 +++++++ .../src/grpc_session/conversion_tests.rs | 1 + codex-rs/code-mode/src/grpc_session/mod.rs | 2 + .../code-mode/src/grpc_session/state_tests.rs | 1 + codex-rs/config/src/config_toml.rs | 10 ++++ codex-rs/core/config.schema.json | 20 +++++++ codex-rs/core/src/code_mode_host.rs | 9 ++- codex-rs/core/src/config/config_tests.rs | 4 ++ codex-rs/core/src/config/mod.rs | 7 +++ codex-rs/core/src/tools/code_mode/mod.rs | 16 ++++- 21 files changed, 260 insertions(+), 10 deletions(-) diff --git a/codex-rs/code-mode-host/src/grpc/mod.rs b/codex-rs/code-mode-host/src/grpc/mod.rs index ea7b9a6d5f..600016b766 100644 --- a/codex-rs/code-mode-host/src/grpc/mod.rs +++ b/codex-rs/code-mode-host/src/grpc/mod.rs @@ -9,6 +9,7 @@ mod waits; use std::future::Future; use std::pin::Pin; use std::sync::Arc; +use std::time::Duration; use std::time::Instant; use codex_code_mode_protocol::CellId; @@ -25,6 +26,7 @@ use codex_protocol::protocol::W3cTraceContext; use futures::Stream; use futures::StreamExt; use tokio::sync::mpsc; +use tokio::time::MissedTickBehavior; use tokio_stream::wrappers::ReceiverStream; use tonic::Request; use tonic::Response; @@ -38,6 +40,21 @@ use self::waits::WaitRegistration; type GrpcStream = Pin> + Send + 'static>>; type GrpcFuture<'a, T> = Pin, Status>> + Send + 'a>>; +fn with_heartbeats( + stream: GrpcStream, + interval_s: u32, + heartbeat: impl Fn() -> T + Send + 'static, +) -> GrpcStream { + if interval_s == 0 { + return stream; + } + let duration = Duration::from_secs(u64::from(interval_s)); + let mut interval = tokio::time::interval_at(tokio::time::Instant::now() + duration, duration); + interval.set_missed_tick_behavior(MissedTickBehavior::Delay); + let stream = tokio_stream::StreamExt::timeout_repeating(stream, interval); + Box::pin(stream.map(move |item| item.unwrap_or_else(|_| Ok(heartbeat())))) +} + fn trace_context_from_request(request: &Request) -> Option { request .metadata() @@ -102,13 +119,23 @@ impl GrpcCodeModeHost { } Some(proto::OutputTruncation::Server) => CodeModeOutputTruncation::Server, }; - Ok(Response::new(self.state.open_session( + let events = self.state.open_session( CodeModeSessionOptions { cell_execution_limits: limits, output_truncation, + ..Default::default() }, max_pending_delegate_calls, - )?)) + )?; + Ok(Response::new(with_heartbeats( + events, + request.heartbeat_interval_s, + || proto::SessionEvent { + event: Some(proto::session_event::Event::Heartbeat( + proto::StreamHeartbeat {}, + )), + }, + ))) } #[tracing::instrument( @@ -135,7 +162,17 @@ impl GrpcCodeModeHost { ) -> Result>, Status> { let _permit = self.state.request_permit()?; let session = self.state.session(&request.session_id)?; - Ok(Response::new(session.subscribe(request.tool_names)?)) + let calls = session.subscribe(request.tool_names)?; + let session_id = session.id; + Ok(Response::new(with_heartbeats( + calls, + request.heartbeat_interval_s, + move || proto::ToolCall { + session_id: session_id.to_string(), + heartbeat: Some(proto::StreamHeartbeat {}), + ..Default::default() + }, + ))) } async fn complete_tool_request( diff --git a/codex-rs/code-mode-host/src/grpc/robustness_tests.rs b/codex-rs/code-mode-host/src/grpc/robustness_tests.rs index fe6cdf56b0..b8ad5f19f9 100644 --- a/codex-rs/code-mode-host/src/grpc/robustness_tests.rs +++ b/codex-rs/code-mode-host/src/grpc/robustness_tests.rs @@ -94,6 +94,7 @@ async fn rejects_oversized_identifiers_and_subscription_filters() { ); assert_invalid( host.subscribe_to_tool_calls(Request::new(proto::SubscribeToToolCallsRequest { + heartbeat_interval_s: 0, session_id: session_id.clone(), tool_names: vec![proto::ToolName { name: oversized_id, @@ -104,6 +105,7 @@ async fn rejects_oversized_identifiers_and_subscription_filters() { ); assert_invalid( host.subscribe_to_tool_calls(Request::new(proto::SubscribeToToolCallsRequest { + heartbeat_interval_s: 0, session_id: session_id.clone(), tool_names: vec![ proto::ToolName { @@ -274,6 +276,7 @@ async fn oversized_encoded_tool_invocation_fails_without_closing_its_session() { let (session_id, _events) = open_session(&host).await; let mut subscription = host .subscribe_to_tool_calls(Request::new(proto::SubscribeToToolCallsRequest { + heartbeat_interval_s: 0, session_id: session_id.clone(), tool_names: Vec::new(), })) @@ -333,6 +336,7 @@ async fn unmatched_tool_subscription_fails_without_closing_its_session() { let (session_id, _events) = open_session(&host).await; let _unrelated = host .subscribe_to_tool_calls(Request::new(proto::SubscribeToToolCallsRequest { + heartbeat_interval_s: 0, session_id: session_id.clone(), tool_names: vec![proto::ToolName { name: "other".to_string(), @@ -374,6 +378,7 @@ async fn unmatched_tool_subscription_fails_without_closing_its_session() { let mut matching = host .subscribe_to_tool_calls(Request::new(proto::SubscribeToToolCallsRequest { + heartbeat_interval_s: 0, session_id, tool_names: vec![proto::ToolName { name: "echo".to_string(), @@ -407,6 +412,7 @@ async fn missing_selected_subscription_retries_another_matching_subscription() { let (session_id, _events) = open_session(&host).await; let mut first = host .subscribe_to_tool_calls(Request::new(proto::SubscribeToToolCallsRequest { + heartbeat_interval_s: 0, session_id: session_id.clone(), tool_names: Vec::new(), })) @@ -415,6 +421,7 @@ async fn missing_selected_subscription_retries_another_matching_subscription() { .into_inner(); let mut second = host .subscribe_to_tool_calls(Request::new(proto::SubscribeToToolCallsRequest { + heartbeat_interval_s: 0, session_id: session_id.clone(), tool_names: Vec::new(), })) @@ -496,6 +503,7 @@ async fn saturated_subscription_does_not_block_independently_filtered_tools() { let (session_id, _events) = open_session(&host).await; let mut slow = host .subscribe_to_tool_calls(Request::new(proto::SubscribeToToolCallsRequest { + heartbeat_interval_s: 0, session_id: session_id.clone(), tool_names: vec![proto::ToolName { name: "slow".to_string(), @@ -507,6 +515,7 @@ async fn saturated_subscription_does_not_block_independently_filtered_tools() { .into_inner(); let mut fast = host .subscribe_to_tool_calls(Request::new(proto::SubscribeToToolCallsRequest { + heartbeat_interval_s: 0, session_id: session_id.clone(), tool_names: vec![proto::ToolName { name: "fast".to_string(), @@ -602,6 +611,7 @@ async fn dropping_subscriptions_only_retires_sessions_with_unread_calls() { let (session_id, _events) = open_session(&host).await; let idle = host .subscribe_to_tool_calls(Request::new(proto::SubscribeToToolCallsRequest { + heartbeat_interval_s: 0, session_id: session_id.clone(), tool_names: Vec::new(), })) @@ -622,6 +632,7 @@ async fn dropping_subscriptions_only_retires_sessions_with_unread_calls() { let first = host .subscribe_to_tool_calls(Request::new(proto::SubscribeToToolCallsRequest { + heartbeat_interval_s: 0, session_id: session_id.clone(), tool_names: Vec::new(), })) @@ -630,6 +641,7 @@ async fn dropping_subscriptions_only_retires_sessions_with_unread_calls() { .into_inner(); let mut second = host .subscribe_to_tool_calls(Request::new(proto::SubscribeToToolCallsRequest { + heartbeat_interval_s: 0, session_id: session_id.clone(), tool_names: Vec::new(), })) diff --git a/codex-rs/code-mode-host/src/grpc/routing.rs b/codex-rs/code-mode-host/src/grpc/routing.rs index 6fe2f76fdf..f39f7d7597 100644 --- a/codex-rs/code-mode-host/src/grpc/routing.rs +++ b/codex-rs/code-mode-host/src/grpc/routing.rs @@ -146,6 +146,7 @@ impl GrpcSession { input_json, sequence, traceparent, + heartbeat: None, }; if message.encoded_len() > MAX_FRAME_BYTES { return Err(format!( diff --git a/codex-rs/code-mode-host/src/grpc/service_tests.rs b/codex-rs/code-mode-host/src/grpc/service_tests.rs index 48888e738e..6407a4bf3f 100644 --- a/codex-rs/code-mode-host/src/grpc/service_tests.rs +++ b/codex-rs/code-mode-host/src/grpc/service_tests.rs @@ -5,6 +5,9 @@ use codex_code_mode_protocol::grpc::code_mode_host_server::CodeModeHost; use futures::FutureExt; use futures::StreamExt; use pretty_assertions::assert_eq; +use std::time::Duration; +use tokio::time::sleep; +use tokio_stream::wrappers::UnboundedReceiverStream; use tonic::Code; use tonic::Request; @@ -15,12 +18,63 @@ pub(super) async fn open_session( (opened.session_id, stream) } +#[tokio::test(start_paused = true)] +async fn zero_interval_keeps_both_rpc_streams_quiet() { + let host = GrpcCodeModeHost::new(); + let (session_id, mut lease) = open_session(&host).await; + let mut calls = host + .subscribe_to_tool_calls(Request::new(proto::SubscribeToToolCallsRequest { + session_id, + tool_names: Vec::new(), + heartbeat_interval_s: 0, + })) + .await + .expect("subscribe to tool calls") + .into_inner(); + + assert!(lease.next().now_or_never().is_none()); + assert!(calls.next().now_or_never().is_none()); + sleep(Duration::from_secs(121)).await; + assert!(lease.next().now_or_never().is_none()); + assert!(calls.next().now_or_never().is_none()); +} + +#[tokio::test(start_paused = true)] +async fn idle_heartbeats_repeat_and_reset_after_traffic() { + let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); + let mut stream = super::with_heartbeats( + Box::pin(UnboundedReceiverStream::new(rx)), + /*interval_s*/ 30, + || "heartbeat", + ); + sleep(Duration::from_secs(29)).await; + tx.send(Ok("event")).unwrap(); + assert_eq!(stream.next().await.unwrap().unwrap(), "event"); + + for _ in 0..2 { + sleep(Duration::from_secs(29)).await; + assert!(stream.next().now_or_never().is_none()); + sleep(Duration::from_secs(2)).await; + let next = stream.next().now_or_never().expect("idle heartbeat due"); + assert_eq!(next.transpose().unwrap(), Some("heartbeat")); + } + sleep(Duration::from_secs(31)).await; + tx.send(Err(tonic::Status::cancelled("closed"))).unwrap(); + drop(tx); + assert_eq!( + stream.next().await.unwrap().unwrap_err().code(), + Code::Cancelled + ); + assert!(stream.next().await.is_none()); +} + async fn open_session_with_truncation( host: &GrpcCodeModeHost, output_truncation: Option, ) -> (proto::SessionOpened, GrpcStream) { let mut stream = host .open_session(Request::new(proto::OpenSessionRequest { + heartbeat_interval_s: 0, max_pending_delegate_calls: None, cell_execution_limits: None, output_truncation, @@ -211,6 +265,7 @@ async fn filtered_subscriptions_match_default_namespace_aliases() { let (session_id, _events) = open_session(&host).await; let mut subscription = host .subscribe_to_tool_calls(Request::new(proto::SubscribeToToolCallsRequest { + heartbeat_interval_s: 0, session_id: session_id.clone(), tool_names: vec![proto::ToolName { name: "echo".to_string(), @@ -277,6 +332,7 @@ async fn filtered_subscriptions_receive_ordered_calls_and_unary_completions() { let (session_id, mut session_events) = open_session(&host).await; let mut matching = host .subscribe_to_tool_calls(Request::new(proto::SubscribeToToolCallsRequest { + heartbeat_interval_s: 0, session_id: session_id.clone(), tool_names: vec![proto::ToolName { name: "echo".to_string(), @@ -288,6 +344,7 @@ async fn filtered_subscriptions_receive_ordered_calls_and_unary_completions() { .into_inner(); let mut unrelated = host .subscribe_to_tool_calls(Request::new(proto::SubscribeToToolCallsRequest { + heartbeat_interval_s: 0, session_id: session_id.clone(), tool_names: vec![proto::ToolName { name: "other".to_string(), @@ -606,6 +663,7 @@ async fn unknown_output_truncation_is_rejected_before_opening_a_session() { let host = GrpcCodeModeHost::new(); let result = host .open_session(Request::new(proto::OpenSessionRequest { + heartbeat_interval_s: 0, max_pending_delegate_calls: None, cell_execution_limits: None, output_truncation: Some(42), diff --git a/codex-rs/code-mode-host/src/lib.rs b/codex-rs/code-mode-host/src/lib.rs index d1a801a6e6..28120949c8 100644 --- a/codex-rs/code-mode-host/src/lib.rs +++ b/codex-rs/code-mode-host/src/lib.rs @@ -436,6 +436,7 @@ impl HostState { output_truncation: output_truncation .map(CodeModeOutputTruncation::from) .unwrap_or_default(), + ..Default::default() }, ) }) diff --git a/codex-rs/code-mode-host/tests/grpc.rs b/codex-rs/code-mode-host/tests/grpc.rs index d3994545af..91822fa579 100644 --- a/codex-rs/code-mode-host/tests/grpc.rs +++ b/codex-rs/code-mode-host/tests/grpc.rs @@ -9,6 +9,7 @@ use codex_code_mode::CodeModeNestedToolCall; use codex_code_mode::CodeModeSession; use codex_code_mode::CodeModeSessionCellExecutionLimits; use codex_code_mode::CodeModeSessionDelegate; +use codex_code_mode::CodeModeSessionOptions; use codex_code_mode::CodeModeSessionProvider; use codex_code_mode::CodeModeToolKind; use codex_code_mode::ExecuteRequest; @@ -222,7 +223,10 @@ async fn tcp_session_persists_values_and_forwards_tools_notifications_and_closur let delegate = Arc::new(RecordingDelegate::default()); let first_delegate = Arc::new(RecordingDelegate::default()); let session = provider - .create_session() + .create_session_with_options(CodeModeSessionOptions { + heartbeat_interval_s: 1, + ..CodeModeSessionOptions::default() + }) .await .map_err(anyhow::Error::msg)?; @@ -243,6 +247,8 @@ async fn tcp_session_persists_values_and_forwards_tools_notifications_and_closur } ); + // Leave both streams idle long enough for the host to send a heartbeat. + tokio::time::sleep(Duration::from_secs(2)).await; let mut callback = request( r#"const result = await tools.echo({ value: String(load("key")) }); notify("notice"); text(result.value);"#, ); @@ -1171,6 +1177,7 @@ async fn dropping_a_grpc_lease_retires_its_server_session() -> Result<()> { .context("connect raw gRPC client")?; let mut lease = client .open_session(grpc::OpenSessionRequest { + heartbeat_interval_s: 1, max_pending_delegate_calls: None, cell_execution_limits: None, output_truncation: None, @@ -1186,12 +1193,32 @@ async fn dropping_a_grpc_lease_retires_its_server_session() -> Result<()> { let Some(grpc::session_event::Event::Opened(opened)) = first.event else { anyhow::bail!("raw gRPC session did not start with an opening event"); }; + let mut subscription = client + .subscribe_to_tool_calls(grpc::SubscribeToToolCallsRequest { + heartbeat_interval_s: 1, + session_id: opened.session_id.clone(), + tool_names: Vec::new(), + }) + .await? + .into_inner(); + let (event, call) = timeout(TEST_TIMEOUT, async { + tokio::try_join!(lease.message(), subscription.message()) + }) + .await + .context("missing an idle Code Mode stream heartbeat")??; + assert!(matches!( + event.and_then(|event| event.event), + Some(grpc::session_event::Event::Heartbeat(_)) + )); + assert!(call.is_some_and(|call| call.heartbeat.is_some())); + drop(subscription); drop(lease); timeout(TEST_TIMEOUT, async { loop { match client .subscribe_to_tool_calls(grpc::SubscribeToToolCallsRequest { + heartbeat_interval_s: 0, session_id: opened.session_id.clone(), tool_names: Vec::new(), }) diff --git a/codex-rs/code-mode-host/tests/grpc_tcp.rs b/codex-rs/code-mode-host/tests/grpc_tcp.rs index eb81054e46..f3e67d8890 100644 --- a/codex-rs/code-mode-host/tests/grpc_tcp.rs +++ b/codex-rs/code-mode-host/tests/grpc_tcp.rs @@ -149,6 +149,7 @@ async fn tcp_listener_opens_a_grpc_session() -> Result<()> { } let response = client .open_session(grpc::OpenSessionRequest { + heartbeat_interval_s: 0, max_pending_delegate_calls: None, cell_execution_limits: None, output_truncation: None, diff --git a/codex-rs/code-mode-host/tests/grpc_timing.rs b/codex-rs/code-mode-host/tests/grpc_timing.rs index e97e3c26aa..057e4aab89 100644 --- a/codex-rs/code-mode-host/tests/grpc_timing.rs +++ b/codex-rs/code-mode-host/tests/grpc_timing.rs @@ -39,6 +39,7 @@ async fn execution_timing_includes_javascript_but_excludes_delayed_reads() -> Re let mut client = CodeModeHostClient::connect(host.endpoint).await?; let mut events = client .open_session(grpc::OpenSessionRequest { + heartbeat_interval_s: 0, max_pending_delegate_calls: None, cell_execution_limits: None, output_truncation: None, diff --git a/codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto b/codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto index 80c212421b..0cf339a0d9 100644 --- a/codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto +++ b/codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto @@ -64,6 +64,9 @@ message OpenSessionRequest { // Absent preserves the legacy 1,024 callback budget. Hosts may impose a lower // shared limit, and older hosts ignore this field and retain their own limit. optional uint32 max_pending_delegate_calls = 3; + reserved 4; // Boolean heartbeat opt-in used by earlier preview builds. + // Seconds of stream inactivity before a heartbeat. Zero disables heartbeats. + uint32 heartbeat_interval_s = 5; } enum OutputTruncation { @@ -84,6 +87,7 @@ message SessionEvent { Notification notification = 3; NotificationCancelled notification_cancelled = 4; CellClosed cell_closed = 5; + StreamHeartbeat heartbeat = 6; } } @@ -102,8 +106,13 @@ message CloseSessionResponse {} message SubscribeToToolCallsRequest { string session_id = 1; repeated ToolName tool_names = 2; + reserved 3; // Boolean heartbeat opt-in used by earlier preview builds. + // Seconds of stream inactivity before a heartbeat. Zero disables heartbeats. + uint32 heartbeat_interval_s = 4; } +message StreamHeartbeat {} + message ToolCall { string session_id = 1; @@ -123,6 +132,9 @@ message ToolCall { // gRPC metadata is fixed when the subscription stream opens, so each call // carries its own context in the message. optional string traceparent = 10; + + // A heartbeat carries only the session ID and must not be handled as a tool call. + optional StreamHeartbeat heartbeat = 11; } message CompleteToolCallRequest { diff --git a/codex-rs/code-mode-protocol/src/session.rs b/codex-rs/code-mode-protocol/src/session.rs index d70fcd4e18..8250117ddc 100644 --- a/codex-rs/code-mode-protocol/src/session.rs +++ b/codex-rs/code-mode-protocol/src/session.rs @@ -57,6 +57,8 @@ pub enum CodeModeOutputTruncation { pub struct CodeModeSessionOptions { pub cell_execution_limits: CodeModeSessionCellExecutionLimits, pub output_truncation: CodeModeOutputTruncation, + /// Interval for idle gRPC heartbeats. Zero disables them. + pub heartbeat_interval_s: u32, } #[derive(Clone, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)] @@ -236,6 +238,11 @@ pub trait CodeModeSessionProvider: Send + Sync { fn create_session(&self) -> CodeModeSessionProviderFuture<'_>; + /// Returns the default options for callers that need to override one setting. + fn default_session_options(&self) -> CodeModeSessionOptions { + CodeModeSessionOptions::default() + } + /// Creates a session whose cells share the supplied execution limits. /// /// Existing providers remain compatible with unlimited sessions, but must diff --git a/codex-rs/code-mode/src/grpc_session/callbacks.rs b/codex-rs/code-mode/src/grpc_session/callbacks.rs index d423e4cdfc..010ec35d1d 100644 --- a/codex-rs/code-mode/src/grpc_session/callbacks.rs +++ b/codex-rs/code-mode/src/grpc_session/callbacks.rs @@ -84,7 +84,8 @@ impl SessionInner { grpc::session_event::Event::Notification(notification) => { self.handle_notification(notification) } - grpc::session_event::Event::NotificationCancelled(_) => Ok(()), + grpc::session_event::Event::NotificationCancelled(_) + | grpc::session_event::Event::Heartbeat(_) => Ok(()), grpc::session_event::Event::CellClosed(closed) => { let cell = self .state @@ -97,13 +98,21 @@ impl SessionInner { } } - fn handle_tool_call(self: &Arc, call: grpc::ToolCall) -> Result<(), String> { + fn handle_tool_call(self: &Arc, mut call: grpc::ToolCall) -> Result<(), String> { if call.session_id != self.id { return Err(format!( "gRPC code-mode tool invocation belongs to session {} instead of {}", call.session_id, self.id )); } + if call.heartbeat.take().is_some() { + // A heartbeat carries only its marker and the verified session ID. + call.session_id.clear(); + if call != grpc::ToolCall::default() { + return Err("gRPC code-mode heartbeat carried a tool invocation".to_string()); + } + return Ok(()); + } let admission = self .state .lock() diff --git a/codex-rs/code-mode/src/grpc_session/callbacks_tests.rs b/codex-rs/code-mode/src/grpc_session/callbacks_tests.rs index 1ebfc416fc..5d4475b0df 100644 --- a/codex-rs/code-mode/src/grpc_session/callbacks_tests.rs +++ b/codex-rs/code-mode/src/grpc_session/callbacks_tests.rs @@ -7,6 +7,7 @@ use std::sync::atomic::AtomicBool; use std::time::Duration; use codex_code_mode_protocol::CodeModeOutputTruncation; +use codex_code_mode_protocol::grpc; use http_body_util::Full; use pretty_assertions::assert_eq; use tokio::sync::Notify; @@ -48,6 +49,27 @@ async fn missing_invocation_completion_racing_cancellation_preserves_session() { assert_eq!(session.require_open(), Ok(())); } +#[tokio::test] +async fn tool_call_heartbeats_do_not_enter_callback_admission() { + let client = GrpcClient::new(BoxCloneSyncService::new(service_fn(|_request| async { + panic!("heartbeat must not send an RPC"); + }))); + let session = session(client); + let heartbeat = grpc::ToolCall { + session_id: session.id.clone(), + heartbeat: Some(grpc::StreamHeartbeat {}), + ..Default::default() + }; + assert_eq!(session.handle_tool_call(heartbeat.clone()), Ok(())); + assert_eq!( + session.handle_tool_call(grpc::ToolCall { + invocation_id: "not-a-heartbeat".to_string(), + ..heartbeat + }), + Err("gRPC code-mode heartbeat carried a tool invocation".to_string()) + ); +} + fn session(client: GrpcClient) -> Arc { let client = Arc::new(client); Arc::new(SessionInner { diff --git a/codex-rs/code-mode/src/grpc_session/conversion_tests.rs b/codex-rs/code-mode/src/grpc_session/conversion_tests.rs index d3250f46f3..c9fbb9f4c1 100644 --- a/codex-rs/code-mode/src/grpc_session/conversion_tests.rs +++ b/codex-rs/code-mode/src/grpc_session/conversion_tests.rs @@ -63,6 +63,7 @@ fn execute_request_preserves_tool_schemas_namespaces_and_limits() { #[test] fn tool_call_decodes_structured_input_and_namespace() { let call = grpc::ToolCall { + heartbeat: None, session_id: "session".to_string(), execution_id: "execution".to_string(), cell_id: "cell".to_string(), diff --git a/codex-rs/code-mode/src/grpc_session/mod.rs b/codex-rs/code-mode/src/grpc_session/mod.rs index 66d10ac3c8..c76a7594f3 100644 --- a/codex-rs/code-mode/src/grpc_session/mod.rs +++ b/codex-rs/code-mode/src/grpc_session/mod.rs @@ -176,6 +176,7 @@ impl GrpcCodeModeSessionProvider { let mut retries = 0; loop { let mut request = tonic::Request::new(grpc::OpenSessionRequest { + heartbeat_interval_s: options.heartbeat_interval_s, max_pending_delegate_calls: Some( codex_code_mode_protocol::host::MAX_PENDING_DELEGATE_CALLS as u32, ), @@ -289,6 +290,7 @@ impl GrpcCodeModeSessionProvider { session.id = %inner.id, ); let mut request = inner.request(grpc::SubscribeToToolCallsRequest { + heartbeat_interval_s: options.heartbeat_interval_s, session_id: inner.id.clone(), tool_names: Vec::new(), }); diff --git a/codex-rs/code-mode/src/grpc_session/state_tests.rs b/codex-rs/code-mode/src/grpc_session/state_tests.rs index 1b7ab62dc8..1646bd6c65 100644 --- a/codex-rs/code-mode/src/grpc_session/state_tests.rs +++ b/codex-rs/code-mode/src/grpc_session/state_tests.rs @@ -34,6 +34,7 @@ fn request(execution_id: &str) -> grpc::ExecuteRequest { fn tool_call(execution_id: &str, invocation_id: u128) -> grpc::ToolCall { let invocation_id = Uuid::from_u128(invocation_id).to_string(); grpc::ToolCall { + heartbeat: None, session_id: "session".to_string(), execution_id: execution_id.to_string(), cell_id: "cell".to_string(), diff --git a/codex-rs/config/src/config_toml.rs b/codex-rs/config/src/config_toml.rs index 340861ba53..f33ea8007f 100644 --- a/codex-rs/config/src/config_toml.rs +++ b/codex-rs/config/src/config_toml.rs @@ -484,6 +484,9 @@ pub struct ConfigToml { /// Additional discoverable tools that can be suggested for installation. pub tool_suggest: Option, + /// Code Mode session settings. + pub code_mode: Option, + /// Agent-related settings (thread limits, etc.). pub agents: Option, @@ -712,6 +715,13 @@ pub struct ToolsToml { pub update_plan: Option, } +#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq, JsonSchema)] +#[schemars(deny_unknown_fields)] +pub struct CodeModeToml { + /// Idle seconds between gRPC stream heartbeats. Zero (the default) disables them. + pub heartbeat_interval_s: Option, +} + #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, JsonSchema)] #[schemars(deny_unknown_fields)] pub struct ExperimentalRequestUserInput { diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 4940a64483..9069235c69 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -710,6 +710,18 @@ }, "type": "object" }, + "CodeModeToml": { + "additionalProperties": false, + "properties": { + "heartbeat_interval_s": { + "description": "Idle seconds between gRPC stream heartbeats. Zero (the default) disables them.", + "format": "uint32", + "minimum": 0.0, + "type": "integer" + } + }, + "type": "object" + }, "ComputerUseConfigToml": { "additionalProperties": false, "properties": { @@ -7880,6 +7892,14 @@ ], "description": "Cloud-owned feature settings." }, + "code_mode": { + "allOf": [ + { + "$ref": "#/definitions/CodeModeToml" + } + ], + "description": "Code Mode session settings." + }, "compact_prompt": { "description": "Compact prompt used for history compaction.", "type": "string" diff --git a/codex-rs/core/src/code_mode_host.rs b/codex-rs/core/src/code_mode_host.rs index 976b637178..deca0e6b9c 100644 --- a/codex-rs/core/src/code_mode_host.rs +++ b/codex-rs/core/src/code_mode_host.rs @@ -65,6 +65,13 @@ impl CodeModeSessionProvider for LocalServerOutputProvider { self.create_session_with_limits(CodeModeSessionCellExecutionLimits::default()) } + fn default_session_options(&self) -> CodeModeSessionOptions { + CodeModeSessionOptions { + output_truncation: CodeModeOutputTruncation::Server, + ..CodeModeSessionOptions::default() + } + } + fn create_session_with_limits<'a>( &'a self, limits: CodeModeSessionCellExecutionLimits, @@ -72,7 +79,7 @@ impl CodeModeSessionProvider for LocalServerOutputProvider { self.inner .create_session_with_options(CodeModeSessionOptions { cell_execution_limits: limits, - output_truncation: CodeModeOutputTruncation::Server, + ..self.default_session_options() }) } diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index c110754403..0e3b64e849 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -667,6 +667,9 @@ async fn load_config_resolves_code_mode_config() -> std::io::Result<()> { let codex_home = tempdir()?; let config_toml: ConfigToml = toml::from_str( r#" +[code_mode] +heartbeat_interval_s = 30 + [features.code_mode] enabled = true default_exec_yield_time_ms = 10000 @@ -700,6 +703,7 @@ disable_in_process_fallback = true vec!["mcp__history".to_string(), "mcp__notes".to_string()] ); assert!(config.code_mode.disable_in_process_fallback); + assert_eq!(config.code_mode.heartbeat_interval_s, 30); assert!(config.features.enabled(Feature::CodeMode)); assert!(config.features.enabled(Feature::CodeModeHost)); Ok(()) diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index e41f69cf0f..84c60748cc 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -1258,6 +1258,7 @@ pub struct CodeModeConfig { pub direct_only_tool_namespaces: Vec, /// Keep code mode fail-closed when the standalone host is unavailable. pub disable_in_process_fallback: bool, + pub heartbeat_interval_s: u32, } impl Default for CodeModeConfig { @@ -1269,6 +1270,7 @@ impl Default for CodeModeConfig { excluded_tool_namespaces: Vec::new(), direct_only_tool_namespaces: Vec::new(), disable_in_process_fallback: false, + heartbeat_interval_s: 0, } } } @@ -2873,6 +2875,11 @@ fn resolve_code_mode_config(config_toml: &ConfigToml) -> CodeModeConfig { disable_in_process_fallback: host .and_then(|config| config.disable_in_process_fallback) .unwrap_or_default(), + heartbeat_interval_s: config_toml + .code_mode + .as_ref() + .and_then(|config| config.heartbeat_interval_s) + .unwrap_or_default(), } } diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index fe59c484e6..3b92fe0d41 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -17,6 +17,7 @@ use codex_code_mode::CellId; use codex_code_mode::CodeModeNestedToolCall; use codex_code_mode::CodeModeOutputTruncation; use codex_code_mode::CodeModeSession; +use codex_code_mode::CodeModeSessionOptions; use codex_code_mode::CodeModeSessionProvider; use codex_code_mode::CodeModeToolKind; use codex_code_mode::RuntimeResponse; @@ -79,6 +80,7 @@ pub(crate) struct CodeModeService { availability: Result<(), String>, dispatch_broker: Arc, default_exec_yield_time_ms: u64, + heartbeat_interval_s: u32, shutdown_token: CancellationToken, unavailable_warning_emitted: AtomicBool, } @@ -98,6 +100,7 @@ impl CodeModeService { availability, dispatch_broker, default_exec_yield_time_ms: config.default_exec_yield_time_ms, + heartbeat_interval_s: config.heartbeat_interval_s, shutdown_token: CancellationToken::new(), unavailable_warning_emitted: AtomicBool::new(false), } @@ -252,9 +255,16 @@ impl CodeModeService { _ = self.shutdown_token.cancelled() => { return Err("code mode session is shutting down".to_string()); } - session = self - .session_provider - .create_session() => session?, + session = async { + if self.heartbeat_interval_s == 0 { + self.session_provider.create_session().await + } else { + self.session_provider.create_session_with_options(CodeModeSessionOptions { + heartbeat_interval_s: self.heartbeat_interval_s, + ..self.session_provider.default_session_options() + }).await + } + } => session?, }; if self.shutdown_token.is_cancelled() { let _ = session.shutdown().await;