diff --git a/codex-rs/code-mode-host/tests/grpc.rs b/codex-rs/code-mode-host/tests/grpc.rs index aa44f19510..c81890998d 100644 --- a/codex-rs/code-mode-host/tests/grpc.rs +++ b/codex-rs/code-mode-host/tests/grpc.rs @@ -1201,6 +1201,17 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> { .await .map_err(anyhow::Error::msg)?; + let saved = execute( + &session, + request(r#"store("saved", "before restart"); text(load("saved"));"#), + delegate.clone(), + ) + .await?; + assert_eq!( + saved, + text_response("1", "before restart", saved.code_mode_host_duration()) + ); + let mut pending = request("await tools.echo({generation: 1}); await new Promise(() => {});"); pending.enabled_tools = vec![tool("echo")]; pending.yield_time_ms = Some(/*value*/ 1); @@ -1209,7 +1220,7 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> { .await .map_err(anyhow::Error::msg)?; let old_cell_id = started.cell_id.clone(); - assert_eq!(old_cell_id, cell_id("1")); + assert_eq!(old_cell_id, cell_id("2")); assert!(matches!( started.initial_response().await, Ok(RuntimeResponse::Yielded { .. }) @@ -1264,8 +1275,38 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> { .context("host loss did not retire the original generation's cell")?; let _replacement = HostHarness::start(&listen_url).await?; + assert!(session.terminate(old_cell_id.clone()).await.is_err()); + + let mut rejected = request(r#"await tools.echo({must_not_run: true});"#); + rejected.enabled_tools = vec![tool("echo")]; + let first = session.execute(rejected.clone(), delegate.clone(), /*preempt*/ None); + let concurrent = session.execute(rejected.clone(), delegate.clone(), /*preempt*/ None); + let queued = session.execute(rejected, delegate.clone(), /*preempt*/ None); + let queued_wait = session.wait( + WaitRequest { + cell_id: old_cell_id.clone(), + yield_time_ms: 1, + }, + /*preempt*/ None, + ); + let (first, concurrent) = timeout(TEST_TIMEOUT, async { tokio::join!(first, concurrent) }) + .await + .context("concurrent recovery calls did not finish")?; + let reset = first.err().expect("first call must report the reset"); + assert_eq!( + reset, + concat!( + "The exec runtime was replaced. Stored values and running cells were lost. ", + "Prior external tool operations may have completed; inspect their results before retrying. ", + "You may continue with a fresh exec call.", + ) + ); + assert_eq!(concurrent.err().expect("concurrent call must fail"), reset); + // Calls created before recovery must still fail even if first polled afterward. + assert_eq!(queued.await.err().expect("queued call must fail"), reset); + assert_eq!(queued_wait.await.unwrap_err(), reset); let mut callback = request( - r#"const result = await tools.echo({generation: 2}); notify("reconnected"); text(result.value);"#, + r#"const result = await tools.echo({generation: 2}); notify("reconnected"); text([result.value, load("saved") ?? null]);"#, ); callback.tool_call_id = "reconnected-call".to_string(); callback.enabled_tools = vec![tool("echo")]; @@ -1286,7 +1327,7 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> { callback_response, text_response( callback_cell_id.as_str(), - "output", + r#"["output",null]"#, callback_response.code_mode_host_duration() ) ); diff --git a/codex-rs/code-mode-host/tests/grpc/admission_tests.rs b/codex-rs/code-mode-host/tests/grpc/admission_tests.rs index 492b21d196..6847bf8174 100644 --- a/codex-rs/code-mode-host/tests/grpc/admission_tests.rs +++ b/codex-rs/code-mode-host/tests/grpc/admission_tests.rs @@ -271,6 +271,19 @@ async fn routing_metadata_is_scoped_to_each_session_on_a_shared_provider() -> Re ); if !retired { *rejection.lock().unwrap() = None; + } else { + let reset = execute( + &parent, + request(r#"store("preserved", "must not run");"#), + Arc::new(NoopCodeModeSessionDelegate), + ) + .await + .expect_err("route loss must report the reset before executing new code"); + assert!( + reset + .to_string() + .contains("Stored values and running cells were lost") + ); } // The old lease and callback streams are still open. Only the explicit // route-loss response may replace the parent; the child must keep its state. 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 index 98b5366b38..b11dbca224 100644 --- a/codex-rs/code-mode-host/tests/grpc/session_recovery_tests.rs +++ b/codex-rs/code-mode-host/tests/grpc/session_recovery_tests.rs @@ -89,6 +89,20 @@ async fn only_missing_sessions_reopen_without_replaying_execution() -> Result<() requests_before_rejection + 1, "the failed call must return without opening a session or replaying Execute" ); + if code == Code::NotFound { + let reset = execute( + &session, + request(r#"store("before", "must not run");"#), + Arc::new(NoopCodeModeSessionDelegate), + ) + .await + .unwrap_err(); + assert!( + reset + .to_string() + .contains("Stored values and running cells were lost") + ); + } let actual = execute( &session, request(r#"text(String(load("before")))"#), diff --git a/codex-rs/code-mode/src/grpc_session/reconnect.rs b/codex-rs/code-mode/src/grpc_session/reconnect.rs index 5d40f8902c..5338a4f3c0 100644 --- a/codex-rs/code-mode/src/grpc_session/reconnect.rs +++ b/codex-rs/code-mode/src/grpc_session/reconnect.rs @@ -75,8 +75,10 @@ impl CodeModeSession for ReconnectableSession { delegate: Arc, preempt: Option, ) -> CodeModeSessionResultFuture<'a, StartedCell> { + // Capture before polling so queued calls cannot silently cross a runtime reset. + let generation = self.inner.binding().map(|binding| binding.generation); Box::pin(async move { - let binding = self.inner.get_or_open_binding().await?; + let binding = self.inner.binding_for_request(generation).await?; let delegate = Arc::new(GenerationDelegate { delegate, generation: binding.generation, @@ -91,8 +93,9 @@ impl CodeModeSession for ReconnectableSession { request: WaitRequest, preempt: Option, ) -> CodeModeSessionResultFuture<'a, WaitOutcome> { + let generation = self.inner.binding().map(|binding| binding.generation); Box::pin(async move { - let binding = self.inner.get_or_open_binding().await?; + let binding = self.inner.binding_for_request(generation).await?; let request = WaitRequest { cell_id: generation::remote_cell_id(binding.generation, &request.cell_id)?, yield_time_ms: request.yield_time_ms, @@ -104,7 +107,11 @@ impl CodeModeSession for ReconnectableSession { fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome> { Box::pin(async move { - let binding = self.inner.get_or_open_binding().await?; + // Cleanup must not replace the runtime and hide its loss from the next request. + let binding = self + .inner + .binding() + .ok_or_else(|| SHUTDOWN_ERROR.to_string())?; let cell_id = generation::remote_cell_id(binding.generation, &cell_id)?; let outcome = binding.session.terminate(cell_id).await?; Ok(generation::public_wait_outcome(binding.generation, outcome)) @@ -125,6 +132,18 @@ impl Drop for ReconnectableSession { } impl ReconnectInner { + async fn binding_for_request(&self, generation: Option) -> Result { + let binding = self.get_or_open_binding().await?; + if generation != Some(binding.generation) { + return Err(concat!( + "The exec runtime was replaced. Stored values and running cells were lost. ", + "Prior external tool operations may have completed; inspect their results before retrying. ", + "You may continue with a fresh exec call.", + ).into()); + } + Ok(binding) + } + async fn get_or_open_binding(&self) -> Result { if self.shutdown_requested.is_cancelled() { return Err(SHUTDOWN_ERROR.to_string()); @@ -148,12 +167,7 @@ impl ReconnectInner { return Ok(binding); } - let previous_binding = self - .binding - .lock() - .unwrap_or_else(PoisonError::into_inner) - .clone(); - if let Some(binding) = previous_binding { + if let Some(binding) = self.binding() { wait_for_watch(binding.session.inner.request_shutdown()).await?; } @@ -186,12 +200,15 @@ impl ReconnectInner { } fn live_binding(&self) -> Option { + self.binding() + .filter(|binding| !binding.session.inner.stopped.is_cancelled()) + } + + fn binding(&self) -> Option { self.binding .lock() .unwrap_or_else(PoisonError::into_inner) - .as_ref() - .filter(|binding| !binding.session.inner.stopped.is_cancelled()) - .cloned() + .clone() } fn request_shutdown(self: &Arc) -> ShutdownResultReceiver {