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;