diff --git a/codex-rs/code-mode-host/tests/grpc.rs b/codex-rs/code-mode-host/tests/grpc.rs index c2af0e3e85..8f584a2203 100644 --- a/codex-rs/code-mode-host/tests/grpc.rs +++ b/codex-rs/code-mode-host/tests/grpc.rs @@ -52,6 +52,8 @@ mod large_tool_delegate; mod network_policy_tests; #[path = "support/recording_delegate.rs"] mod recording_delegate; +#[path = "grpc/session_recovery_tests.rs"] +mod session_recovery_tests; use host::HostHarness; use large_tool_delegate::LargeToolResultDelegate; diff --git a/codex-rs/code-mode-host/tests/grpc/session_recovery_tests.rs b/codex-rs/code-mode-host/tests/grpc/session_recovery_tests.rs new file mode 100644 index 0000000000..fa75218808 --- /dev/null +++ b/codex-rs/code-mode-host/tests/grpc/session_recovery_tests.rs @@ -0,0 +1,101 @@ +//! Exercises missing-session recovery while the host's lease streams remain open. + +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +use anyhow::Result; +use codex_code_mode::CodeModeSessionProvider; +use codex_code_mode::GrpcCodeModeSessionProvider; +use codex_code_mode::NoopCodeModeSessionDelegate; +use codex_code_mode_host::GrpcCodeModeHost; +use codex_code_mode_protocol::grpc::code_mode_host_server::CodeModeHostServer; +use pretty_assertions::assert_eq; +use tokio::net::TcpListener; +use tokio_stream::wrappers::TcpListenerStream; +use tokio_util::task::AbortOnDropHandle; +use tonic::Code; +use tonic::Status; +use tonic::transport::Server; + +use super::execute; +use super::request; +use super::text_response; + +#[tokio::test] +async fn only_missing_sessions_reopen_without_replaying_execution() -> Result<()> { + for (code, expected_cell, expected_value) in [ + (Code::NotFound, "g2:1", "undefined"), + (Code::InvalidArgument, "2", "1"), + ] { + let rejection = Arc::new(Mutex::new(/*t*/ None)); + let next_rejection = Arc::clone(&rejection); + let request_count = Arc::new(AtomicUsize::new(/*v*/ 0)); + let count = Arc::clone(&request_count); + let service = CodeModeHostServer::with_interceptor( + GrpcCodeModeHost::new(), + move |request: tonic::Request<()>| { + count.fetch_add(/*val*/ 1, Ordering::SeqCst); + match next_rejection.lock().unwrap().take() { + Some(rejection) => Err(rejection), + None => Ok(request), + } + }, + ); + let listener = TcpListener::bind("127.0.0.1:0").await?; + let endpoint = format!("http://{}", listener.local_addr()?); + let _server = AbortOnDropHandle::new(tokio::spawn( + Server::builder() + .add_service(service) + .serve_with_incoming(TcpListenerStream::new(listener)), + )); + let session = GrpcCodeModeSessionProvider::new(endpoint) + .create_session() + .await + .map_err(anyhow::Error::msg)?; + execute( + &session, + request(r#"store("before", 1);"#), + Arc::new(NoopCodeModeSessionDelegate), + ) + .await?; + + // Reject the next Execute without closing the lease or subscription streams. + let failure = Status::new(code, "injected RPC rejection"); + *rejection.lock().unwrap() = Some(failure.clone()); + let requests_before_rejection = request_count.load(Ordering::SeqCst); + let failed = execute( + &session, + request(r#"store("before", "replayed");"#), + Arc::new(NoopCodeModeSessionDelegate), + ) + .await + .unwrap_err(); + assert_eq!( + failed.to_string(), + format!("gRPC code-mode execution failed: {failure}") + ); + assert_eq!( + request_count.load(Ordering::SeqCst), + requests_before_rejection + 1, + "the failed call must return without opening a session or replaying Execute" + ); + let actual = execute( + &session, + request(r#"text(String(load("before")))"#), + Arc::new(NoopCodeModeSessionDelegate), + ) + .await?; + assert_eq!( + actual, + text_response( + expected_cell, + expected_value, + actual.code_mode_host_duration() + ) + ); + session.shutdown().await.map_err(anyhow::Error::msg)?; + } + Ok(()) +} 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 84fb22a373..4e4134c21c 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 @@ -5,6 +5,10 @@ package codex.code_mode.v1; // Hosts stateful JavaScript execution and delegates nested tool calls to the // session owner. Large tool inputs and outputs stay off the session event // stream so independent HTTP/2 streams can make progress concurrently. +// +// Execute (including its stream), Wait, and Terminate reserve NOT_FOUND for an +// unknown session_id. Clients may retire the session binding on this status. +// Missing cells use WaitResponse.missing_cell instead. service CodeModeHost { // Opens a session lease. The first event is always SessionOpened; dropping // this stream closes the session and terminates its active cells. @@ -19,6 +23,7 @@ service CodeModeHost { // Each result receives its own HTTP/2 stream, preventing a large response // from blocking unrelated tool completions or session control events. + // NOT_FOUND may identify a missing invocation rather than a missing session. rpc CompleteToolCall(CompleteToolCallRequest) returns (CompleteToolCallResponse); rpc AcknowledgeNotification(AcknowledgeNotificationRequest) diff --git a/codex-rs/code-mode/src/grpc_session/callbacks.rs b/codex-rs/code-mode/src/grpc_session/callbacks.rs index 21f74cec0c..94a1631bc1 100644 --- a/codex-rs/code-mode/src/grpc_session/callbacks.rs +++ b/codex-rs/code-mode/src/grpc_session/callbacks.rs @@ -193,6 +193,8 @@ impl SessionInner { tokio::select! { biased; _ = cancellation.cancelled() => {} + // NotFound can mean an invocation in a live session was canceled. + // Check cancellation before treating a completion error as fatal. result = deadline::request( self, "tool invocation completion", @@ -266,3 +268,7 @@ impl SessionInner { Ok(()) } } + +#[cfg(test)] +#[path = "callbacks_tests.rs"] +mod tests; diff --git a/codex-rs/code-mode/src/grpc_session/callbacks_tests.rs b/codex-rs/code-mode/src/grpc_session/callbacks_tests.rs new file mode 100644 index 0000000000..6351d3ce50 --- /dev/null +++ b/codex-rs/code-mode/src/grpc_session/callbacks_tests.rs @@ -0,0 +1,50 @@ +//! Keeps late tool-completion errors scoped to the canceled invocation. + +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::atomic::AtomicBool; + +use pretty_assertions::assert_eq; +use tokio_util::sync::CancellationToken; +use tokio_util::task::TaskTracker; +use tower::service_fn; +use tower::util::BoxCloneSyncService; + +use super::super::GrpcClient; +use super::super::state::SessionState; +use super::super::transport::SharedTransport; +use super::SessionInner; + +#[tokio::test] +async fn missing_invocation_completion_racing_cancellation_preserves_session() { + let cancellation = CancellationToken::new(); + let completion_cancellation = cancellation.clone(); + let client = GrpcClient::new(BoxCloneSyncService::new(service_fn(move |_request| { + // Cancel after the biased select has polled cancellation, before the RPC fails. + completion_cancellation.cancel(); + async { Ok(tonic::Status::not_found("unknown code-mode tool invocation").into_http()) } + }))); + let session = SessionInner { + id: "session".to_string(), + client: client.clone(), + runtime: tokio::runtime::Handle::current(), + state: Mutex::new(SessionState::default()), + wait_slots: Mutex::new(HashMap::new()), + shutdown_requested: AtomicBool::new(/*v*/ false), + shutdown_result: Mutex::new(/*t*/ None), + stopped: CancellationToken::new(), + stream_tasks: TaskTracker::new(), + _transport: Arc::new(SharedTransport::Connected(client)), + }; + + session + .complete_tool_call( + "invocation".to_string(), + cancellation, + Ok(serde_json::Value::Null), + ) + .await; + + assert_eq!(session.require_open(), Ok(())); +} diff --git a/codex-rs/code-mode/src/grpc_session/deadline.rs b/codex-rs/code-mode/src/grpc_session/deadline.rs index 266b1ed5e5..f172343c84 100644 --- a/codex-rs/code-mode/src/grpc_session/deadline.rs +++ b/codex-rs/code-mode/src/grpc_session/deadline.rs @@ -42,6 +42,25 @@ pub(super) async fn request( } } +/// Invalidates the session on RPCs where `NotFound` means the session is missing. +pub(super) async fn session_request( + session: &SessionInner, + operation: &str, + runtime_timeout: Duration, + rpc: impl Future>, +) -> Result { + request(session, operation, runtime_timeout, async { + let result = rpc.await; + if let Err(error) = &result + && error.code() == tonic::Code::NotFound + { + session.fail(failure(operation, error)); + } + result + }) + .await +} + pub(super) fn failure(operation: &str, error: impl Display) -> String { let mut message = format!("gRPC code-mode {operation} failed: {error}"); if message.len() > MAX_ERROR_BYTES { diff --git a/codex-rs/code-mode/src/grpc_session/operations.rs b/codex-rs/code-mode/src/grpc_session/operations.rs index 922a036a3b..bd266bedf8 100644 --- a/codex-rs/code-mode/src/grpc_session/operations.rs +++ b/codex-rs/code-mode/src/grpc_session/operations.rs @@ -130,11 +130,15 @@ impl SessionInner { { request.metadata_mut().insert("traceparent", traceparent); } - let mut stream = - deadline::request(&self, "execution", Duration::ZERO, client.execute(request)) - .await? - .into_inner(); - let first = deadline::request( + let mut stream = deadline::session_request( + &self, + "execution", + Duration::ZERO, + client.execute(request), + ) + .await? + .into_inner(); + let first = deadline::session_request( &self, "execution starting event", Duration::ZERO, @@ -217,7 +221,7 @@ impl SessionInner { let outcome = tokio::select! { biased; _ = response_tx.closed() => return, - outcome = deadline::request( + outcome = deadline::session_request( &self, "execution outcome", runtime_timeout, @@ -317,7 +321,8 @@ impl SessionInner { grpc::yield_observation_request::Observation::WaitId(wait_id), ) }); - let response = deadline::request(self, "wait", runtime_timeout, client.wait(request)).await; + let response = + deadline::session_request(self, "wait", runtime_timeout, client.wait(request)).await; drop(watcher); cancellation.disarm(); self.prune_wait_slots(); @@ -351,7 +356,7 @@ impl SessionInner { pub(super) async fn terminate(&self, cell_id: CellId) -> Result { self.require_open()?; let mut client = self.client(); - let response = deadline::request( + let response = deadline::session_request( self, "termination", Duration::ZERO,