From 9be95745fb84f540582e6ef4144e7edf5bcb7959 Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Tue, 11 Aug 2026 00:46:47 +0000 Subject: [PATCH] Make gRPC code-mode notifications fire-and-forget (#37906) ## What changed - Send notification events without waiting for client acknowledgment, so an unacknowledged notification no longer delays cell completion. - Keep the notification acknowledgment RPC as a compatibility no-op after validating the session and notification ID. - Remove pending-notification tracking and cancellation events from the gRPC session routing path. ## Testing - Verify that a cell completes after emitting an unacknowledged notification and that a later acknowledgment is still accepted. GitOrigin-RevId: 19cf3ed5ce82cf45d84fa75f238c2afa4e83acb5 --- codex-rs/code-mode-host/src/grpc/delegate.rs | 75 ++++--------------- codex-rs/code-mode-host/src/grpc/mod.rs | 5 +- codex-rs/code-mode-host/src/grpc/routing.rs | 67 ----------------- .../code-mode-host/src/grpc/service_tests.rs | 50 +++++-------- codex-rs/code-mode-host/src/grpc/session.rs | 3 - 5 files changed, 32 insertions(+), 168 deletions(-) diff --git a/codex-rs/code-mode-host/src/grpc/delegate.rs b/codex-rs/code-mode-host/src/grpc/delegate.rs index 4b540f8c4b..2a42d78c57 100644 --- a/codex-rs/code-mode-host/src/grpc/delegate.rs +++ b/codex-rs/code-mode-host/src/grpc/delegate.rs @@ -56,7 +56,10 @@ impl CodeModeSessionDelegate for GrpcDelegate { &cancellation, ) .await?; - let mut pending = PendingCallback::tool(Arc::clone(&session), invocation_id); + let mut pending = PendingToolCall { + session: Arc::clone(&session), + id: Some(invocation_id), + }; let result = tokio::select! { biased; result = receiver => result @@ -69,7 +72,7 @@ impl CodeModeSessionDelegate for GrpcDelegate { } }; if result.is_ok() { - pending.disarm(); + pending.id = None; } result }) @@ -92,38 +95,18 @@ impl CodeModeSessionDelegate for GrpcDelegate { .execution_id(cell_id.as_str(), &cancellation) .await?; let notification_id = Uuid::new_v4(); - let (response, receiver) = oneshot::channel(); session - .begin_notification( - notification_id, - proto::Notification { + .send_event( + proto::session_event::Event::Notification(proto::Notification { notification_id: notification_id.to_string(), execution_id, cell_id: cell_id.to_string(), call_id, text, - }, - response, + }), &cancellation, ) - .await?; - let mut pending = PendingCallback::notification(Arc::clone(&session), notification_id); - tokio::select! { - biased; - result = receiver => { - result.map_err(|_| { - "code-mode client closed before acknowledging notification".to_string() - })?; - pending.disarm(); - Ok(()) - } - _ = cancellation.cancelled() => { - Err("code mode notification cancelled".to_string()) - } - _ = session.closed.cancelled() => { - Err("code-mode session closed before acknowledging notification".to_string()) - } - } + .await }) } @@ -134,47 +117,15 @@ impl CodeModeSessionDelegate for GrpcDelegate { } } -enum CallbackKind { - Tool, - Notification, -} - -struct PendingCallback { +struct PendingToolCall { session: Arc, id: Option, - kind: CallbackKind, } -impl PendingCallback { - fn tool(session: Arc, id: Uuid) -> Self { - Self { - session, - id: Some(id), - kind: CallbackKind::Tool, - } - } - - fn notification(session: Arc, id: Uuid) -> Self { - Self { - session, - id: Some(id), - kind: CallbackKind::Notification, - } - } - - fn disarm(&mut self) { - self.id = None; - } -} - -impl Drop for PendingCallback { +impl Drop for PendingToolCall { fn drop(&mut self) { - let Some(id) = self.id.take() else { - return; - }; - match self.kind { - CallbackKind::Tool => self.session.cancel_invocation(id), - CallbackKind::Notification => self.session.cancel_notification(id), + if let Some(id) = self.id.take() { + self.session.cancel_invocation(id); } } } diff --git a/codex-rs/code-mode-host/src/grpc/mod.rs b/codex-rs/code-mode-host/src/grpc/mod.rs index 84145a4bfd..10fe9df4cb 100644 --- a/codex-rs/code-mode-host/src/grpc/mod.rs +++ b/codex-rs/code-mode-host/src/grpc/mod.rs @@ -106,9 +106,8 @@ impl GrpcCodeModeHost { request: proto::AcknowledgeNotificationRequest, ) -> Result, Status> { let _permit = self.state.control_permit()?; - let session = self.state.session(&request.session_id)?; - let notification_id = validation::uuid(&request.notification_id, "notification ID")?; - session.acknowledge_notification(notification_id)?; + self.state.session(&request.session_id)?; + validation::uuid(&request.notification_id, "notification ID")?; Ok(Response::new(proto::AcknowledgeNotificationResponse {})) } diff --git a/codex-rs/code-mode-host/src/grpc/routing.rs b/codex-rs/code-mode-host/src/grpc/routing.rs index e6a29efdce..61fefdccbd 100644 --- a/codex-rs/code-mode-host/src/grpc/routing.rs +++ b/codex-rs/code-mode-host/src/grpc/routing.rs @@ -237,71 +237,4 @@ impl GrpcSession { ); } } - - pub(super) async fn begin_notification( - &self, - notification_id: Uuid, - notification: proto::Notification, - response: oneshot::Sender<()>, - cancellation: &CancellationToken, - ) -> Result<(), String> { - { - let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner); - state - .pending_notifications - .insert(notification_id, response); - state.seen_notifications.remember(notification_id); - } - if let Err(error) = self - .send_event( - proto::session_event::Event::Notification(notification), - cancellation, - ) - .await - { - self.state - .lock() - .unwrap_or_else(PoisonError::into_inner) - .pending_notifications - .remove(¬ification_id); - return Err(error); - } - Ok(()) - } - - pub(super) fn acknowledge_notification(&self, notification_id: Uuid) -> Result<(), Status> { - let response = { - let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner); - match state.pending_notifications.remove(¬ification_id) { - Some(response) => Some(response), - None if state.seen_notifications.contains(¬ification_id) => None, - None => { - return Err(Status::not_found(format!( - "unknown code-mode notification {notification_id}" - ))); - } - } - }; - if let Some(response) = response { - let _ = response.send(()); - } - Ok(()) - } - - pub(super) fn cancel_notification(&self, notification_id: Uuid) { - let pending = self - .state - .lock() - .unwrap_or_else(PoisonError::into_inner) - .pending_notifications - .remove(¬ification_id); - if pending.is_some() { - let _ = self.send_event_now( - proto::session_event::Event::NotificationCancelled(proto::NotificationCancelled { - notification_id: notification_id.to_string(), - }), - /*cell_permit*/ None, - ); - } - } } 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 2ee9e7bcc9..e441796277 100644 --- a/codex-rs/code-mode-host/src/grpc/service_tests.rs +++ b/codex-rs/code-mode-host/src/grpc/service_tests.rs @@ -286,18 +286,16 @@ async fn cancellation_before_wait_admission_is_preserved() { } #[tokio::test] -async fn terminating_a_cell_cancels_unacknowledged_notifications() { +async fn notifications_do_not_delay_cell_completion() { let host = GrpcCodeModeHost::new(); let (session_id, mut session_events) = open_session(&host).await; - let (cell_id, mut execution) = execute_events( - &host, - execute_request( - &session_id, - "execution-notify", - r#"notify("pending"); await new Promise(() => {});"#, - ), - ) - .await; + let mut request = execute_request( + &session_id, + "execution-notify", + r#"notify("pending"); text("done");"#, + ); + request.yield_time_ms = Some(60_000); + let (cell_id, mut execution) = execute_events(&host, request).await; let event = session_events.next().await.unwrap().unwrap(); let Some(proto::session_event::Event::Notification(notification)) = event.event else { panic!("expected pending notification"); @@ -306,35 +304,15 @@ async fn terminating_a_cell_cancels_unacknowledged_notifications() { assert_eq!(notification.cell_id, cell_id); assert_eq!(notification.call_id, "outer-call"); assert_eq!(notification.text, "pending"); - execution.next().await.unwrap().unwrap(); - - let terminated = host - .terminate(Request::new(proto::TerminateRequest { - session_id, - cell_id: cell_id.clone(), - })) - .await - .unwrap() - .into_inner(); assert!(matches!( - terminated.state, - Some(proto::wait_response::State::LiveCell( + execution.next().await.unwrap().unwrap().event, + Some(proto::execute_event::Event::Outcome( proto::ExecutionOutcome { - outcome: Some(proto::execution_outcome::Outcome::Terminated(_)), + outcome: Some(proto::execution_outcome::Outcome::Completed(_)), .. } )) )); - assert_eq!( - session_events.next().await.unwrap().unwrap(), - proto::SessionEvent { - event: Some(proto::session_event::Event::NotificationCancelled( - proto::NotificationCancelled { - notification_id: notification.notification_id, - }, - )), - } - ); assert_eq!( session_events.next().await.unwrap().unwrap(), proto::SessionEvent { @@ -345,4 +323,10 @@ async fn terminating_a_cell_cancels_unacknowledged_notifications() { })), } ); + host.acknowledge_notification(Request::new(proto::AcknowledgeNotificationRequest { + session_id, + notification_id: notification.notification_id, + })) + .await + .expect("legacy notification acknowledgments remain accepted"); } diff --git a/codex-rs/code-mode-host/src/grpc/session.rs b/codex-rs/code-mode-host/src/grpc/session.rs index c88c6add86..4b7754665e 100644 --- a/codex-rs/code-mode-host/src/grpc/session.rs +++ b/codex-rs/code-mode-host/src/grpc/session.rs @@ -64,8 +64,6 @@ pub(super) struct SessionState { pub(super) next_subscription: usize, pub(super) pending_invocations: HashMap, pub(super) seen_invocations: BoundedIds, - pub(super) pending_notifications: HashMap>, - pub(super) seen_notifications: BoundedIds, pub(super) waits: HashMap, pub(super) seen_waits: BoundedIds, pub(super) cancelled_waits: BoundedIds, @@ -253,7 +251,6 @@ impl GrpcSession { wait.cancellation.cancel(); } state.pending_invocations.clear(); - state.pending_notifications.clear(); state.subscriptions.clear(); } let result = self.runtime.shutdown().await.map_err(Status::internal);