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

@@ -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<T> = Pin<Box<dyn Stream<Item = Result<T, Status>> + Send + 'static>>;
 type GrpcFuture<'a, T> = Pin<Box<dyn Future<Output = Result<Response<T>, Status>> + Send + 'a>>;
 
+fn with_heartbeats<T: Send + 'static>(
+    stream: GrpcStream<T>,
+    interval_s: u32,
+    heartbeat: impl Fn() -> T + Send + 'static,
+) -> GrpcStream<T> {
+    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<T>(request: &Request<T>) -> Option<W3cTraceContext> {
     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<Response<GrpcStream<proto::ToolCall>>, 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(

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(),
         }))

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!(

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<i32>,
 ) -> (proto::SessionOpened, GrpcStream<proto::SessionEvent>) {
     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),

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()
                         },
                     )
                 })

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(),
                 })

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,

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,

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 {

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

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<Self>, call: grpc::ToolCall) -> Result<(), String> {
+    fn handle_tool_call(self: &Arc<Self>, 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()

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<SessionInner> {
     let client = Arc::new(client);
     Arc::new(SessionInner {

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(),

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(),
         });

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(),

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<ToolSuggestConfig>,
 
+    /// Code Mode session settings.
+    pub code_mode: Option<CodeModeToml>,
+
     /// Agent-related settings (thread limits, etc.).
     pub agents: Option<AgentsToml>,
 
@@ -712,6 +715,13 @@ pub struct ToolsToml {
     pub update_plan: Option<UpdatePlanToolConfig>,
 }
 
+#[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<u32>,
+}
+
 #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, JsonSchema)]
 #[schemars(deny_unknown_fields)]
 pub struct ExperimentalRequestUserInput {

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"

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()
             })
     }
 

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(())

codex-rs/core/src/config/mod.rs

@@ -1258,6 +1258,7 @@ pub struct CodeModeConfig {
     pub direct_only_tool_namespaces: Vec<String>,
     /// 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(),
     }
 }
 

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<CodeModeDispatchBroker>,
     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;