openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-19 17:28:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

Make gRPC code-mode notifications fire-and-forget (#37906)

Browse Source 
## 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
This commit is contained in:
Channing Conger
2026-08-11 00:52:35 +00:00
committed by copyberry
parent 260261ed8f
commit 9be95745fb
5 changed files with 32 additions and 168 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/code-mode-host/src/grpc/delegate.rs
+13 -62
View File
@@ -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<GrpcSession>,
id: Option<Uuid>,
kind: CallbackKind,
}
impl PendingCallback {
fn tool(session: Arc<GrpcSession>, id: Uuid) -> Self {
Self {
session,
id: Some(id),
kind: CallbackKind::Tool,
}
}
fn notification(session: Arc<GrpcSession>, 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);
}
}
}
codex-rs/code-mode-host/src/grpc/mod.rs
+2 -3
View File
@@ -106,9 +106,8 @@ impl GrpcCodeModeHost {
request: proto::AcknowledgeNotificationRequest,
) -> Result<Response<proto::AcknowledgeNotificationResponse>, 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 {}))
}
codex-rs/code-mode-host/src/grpc/routing.rs
-67
View File
@@ -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(&notification_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(&notification_id) {
Some(response) => Some(response),
None if state.seen_notifications.contains(&notification_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(&notification_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,
);
}
}
}
codex-rs/code-mode-host/src/grpc/service_tests.rs
+17 -33
View File
@@ -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");
}
codex-rs/code-mode-host/src/grpc/session.rs
-3
View File
@@ -64,8 +64,6 @@ pub(super) struct SessionState {
pub(super) next_subscription: usize,
pub(super) pending_invocations: HashMap<Uuid, PendingInvocation>,
pub(super) seen_invocations: BoundedIds<Uuid>,
pub(super) pending_notifications: HashMap<Uuid, oneshot::Sender<()>>,
pub(super) seen_notifications: BoundedIds<Uuid>,
pub(super) waits: HashMap<String, ActiveWait>,
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);