mirror of
https://github.com/openai/codex.git
synced 2026-11-14 21:40:39 +00:00
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
This commit is contained in:
@@ -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()
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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(),
|
||||
});
|
||||
|
||||
@@ -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(),
|
||||
|
||||
Reference in New Issue
Block a user