From 83121bf58de06acf3f1ff6617a73ad12181553ef Mon Sep 17 00:00:00 2001 From: Eric Traut Date: Wed, 27 May 2026 12:30:41 -0700 Subject: [PATCH] Preserve extension event ordering --- codex-rs/app-server/src/extensions.rs | 73 +++++++++++++------ .../extension-api/src/capabilities/events.rs | 12 ++- .../ext/extension-api/src/capabilities/mod.rs | 1 + codex-rs/ext/extension-api/src/lib.rs | 1 + codex-rs/ext/goal/src/events.rs | 20 ++--- codex-rs/ext/goal/src/runtime.rs | 35 +++++---- codex-rs/ext/goal/src/tool.rs | 19 ++--- .../ext/goal/tests/goal_extension_backend.rs | 3 +- 8 files changed, 105 insertions(+), 59 deletions(-) diff --git a/codex-rs/app-server/src/extensions.rs b/codex-rs/app-server/src/extensions.rs index 216fb11198..6cdd4ac206 100644 --- a/codex-rs/app-server/src/extensions.rs +++ b/codex-rs/app-server/src/extensions.rs @@ -9,6 +9,7 @@ use codex_core::ThreadManager; use codex_core::config::Config; use codex_extension_api::AgentSpawnFuture; use codex_extension_api::AgentSpawner; +use codex_extension_api::ExtensionEventFuture; use codex_extension_api::ExtensionEventSink; use codex_extension_api::ExtensionRegistry; use codex_extension_api::ExtensionRegistryBuilder; @@ -59,24 +60,23 @@ struct AppServerExtensionEventSink { } impl ExtensionEventSink for AppServerExtensionEventSink { - fn emit(&self, event: Event) { - match event.msg { - EventMsg::ThreadGoalUpdated(thread_goal_event) => { - let outgoing = Arc::clone(&self.outgoing); - let notification = - ServerNotification::ThreadGoalUpdated(ThreadGoalUpdatedNotification { - thread_id: thread_goal_event.thread_id.to_string(), - turn_id: thread_goal_event.turn_id, - goal: thread_goal_event.goal.into(), - }); - let _handle = tokio::spawn(async move { - outgoing.send_server_notification(notification).await; - }); + fn emit<'a>(&'a self, event: Event) -> ExtensionEventFuture<'a> { + Box::pin(async move { + match event.msg { + EventMsg::ThreadGoalUpdated(thread_goal_event) => { + let notification = + ServerNotification::ThreadGoalUpdated(ThreadGoalUpdatedNotification { + thread_id: thread_goal_event.thread_id.to_string(), + turn_id: thread_goal_event.turn_id, + goal: thread_goal_event.goal.into(), + }); + self.outgoing.send_server_notification(notification).await; + } + msg => { + tracing::debug!(event_id = %event.id, ?msg, "dropping unsupported extension event"); + } } - msg => { - tracing::debug!(event_id = %event.id, ?msg, "dropping unsupported extension event"); - } - } + }) } } @@ -138,13 +138,17 @@ mod tests { .expect("prefill should fit in one-slot channel"); let sink = app_server_extension_event_sink(outgoing); - sink.emit(thread_goal_update_event( - thread_id, - "wait for capacity", - "turn-1", - )); + let emit = tokio::spawn(async move { + sink.emit(thread_goal_update_event( + thread_id, + "wait for capacity", + "turn-1", + )) + .await; + }); let _prefill = recv_goal_update(&mut outgoing_rx).await; + emit.await.expect("event emission should complete"); let notification = recv_goal_update(&mut outgoing_rx).await; assert_eq!( @@ -153,6 +157,31 @@ mod tests { ); } + #[tokio::test] + async fn app_server_event_sink_preserves_goal_update_order() { + let (outgoing_tx, mut outgoing_rx) = mpsc::channel(2); + let outgoing = Arc::new(OutgoingMessageSender::new( + outgoing_tx, + AnalyticsEventsClient::disabled(), + )); + let thread_id = ThreadId::default(); + let sink = app_server_extension_event_sink(outgoing); + + sink.emit(thread_goal_update_event(thread_id, "first", "turn-1")) + .await; + sink.emit(thread_goal_update_event(thread_id, "second", "turn-2")) + .await; + + assert_eq!( + app_server_goal_update(thread_id, "first", "turn-1"), + recv_goal_update(&mut outgoing_rx).await + ); + assert_eq!( + app_server_goal_update(thread_id, "second", "turn-2"), + recv_goal_update(&mut outgoing_rx).await + ); + } + fn thread_goal_update_event(thread_id: ThreadId, objective: &str, turn_id: &str) -> Event { Event { id: "call-1".to_string(), diff --git a/codex-rs/ext/extension-api/src/capabilities/events.rs b/codex-rs/ext/extension-api/src/capabilities/events.rs index d8fda1aea7..f51b6e1d47 100644 --- a/codex-rs/ext/extension-api/src/capabilities/events.rs +++ b/codex-rs/ext/extension-api/src/capabilities/events.rs @@ -1,13 +1,17 @@ use codex_protocol::protocol::Event; +use std::future::Future; +use std::pin::Pin; -/// Host-provided fire-and-forget sink for extension-generated events. +pub type ExtensionEventFuture<'a> = Pin + Send + 'a>>; + +/// Host-provided sink for extension-generated events. /// /// Extensions construct protocol events with the correlation id appropriate for /// the callback they are handling, then leave persistence, ordering, transport /// fanout, and logging decisions to the host. pub trait ExtensionEventSink: Send + Sync { /// Queue one protocol event for host-owned delivery. - fn emit(&self, event: Event); + fn emit<'a>(&'a self, event: Event) -> ExtensionEventFuture<'a>; } /// Event sink used when the host does not expose extension event emission. @@ -15,5 +19,7 @@ pub trait ExtensionEventSink: Send + Sync { pub struct NoopExtensionEventSink; impl ExtensionEventSink for NoopExtensionEventSink { - fn emit(&self, _event: Event) {} + fn emit<'a>(&'a self, _event: Event) -> ExtensionEventFuture<'a> { + Box::pin(std::future::ready(())) + } } diff --git a/codex-rs/ext/extension-api/src/capabilities/mod.rs b/codex-rs/ext/extension-api/src/capabilities/mod.rs index 37c36e573b..0fabed3929 100644 --- a/codex-rs/ext/extension-api/src/capabilities/mod.rs +++ b/codex-rs/ext/extension-api/src/capabilities/mod.rs @@ -4,6 +4,7 @@ mod response_items; pub use agent::AgentSpawnFuture; pub use agent::AgentSpawner; +pub use events::ExtensionEventFuture; pub use events::ExtensionEventSink; pub use events::NoopExtensionEventSink; pub use response_items::NoopResponseItemInjector; diff --git a/codex-rs/ext/extension-api/src/lib.rs b/codex-rs/ext/extension-api/src/lib.rs index a49680fff8..3a75597dc4 100644 --- a/codex-rs/ext/extension-api/src/lib.rs +++ b/codex-rs/ext/extension-api/src/lib.rs @@ -5,6 +5,7 @@ mod state; pub use capabilities::AgentSpawnFuture; pub use capabilities::AgentSpawner; +pub use capabilities::ExtensionEventFuture; pub use capabilities::ExtensionEventSink; pub use capabilities::NoopExtensionEventSink; pub use capabilities::NoopResponseItemInjector; diff --git a/codex-rs/ext/goal/src/events.rs b/codex-rs/ext/goal/src/events.rs index ab9eda4056..75f8c592a6 100644 --- a/codex-rs/ext/goal/src/events.rs +++ b/codex-rs/ext/goal/src/events.rs @@ -16,19 +16,21 @@ impl GoalEventEmitter { Self { sink } } - pub(crate) fn thread_goal_updated( + pub(crate) async fn thread_goal_updated( &self, event_id: impl Into, turn_id: Option, goal: ThreadGoal, ) { - self.sink.emit(Event { - id: event_id.into(), - msg: EventMsg::ThreadGoalUpdated(ThreadGoalUpdatedEvent { - thread_id: goal.thread_id, - turn_id, - goal, - }), - }); + self.sink + .emit(Event { + id: event_id.into(), + msg: EventMsg::ThreadGoalUpdated(ThreadGoalUpdatedEvent { + thread_id: goal.thread_id, + turn_id, + goal, + }), + }) + .await; } } diff --git a/codex-rs/ext/goal/src/runtime.rs b/codex-rs/ext/goal/src/runtime.rs index bdfee04fbc..51716aa17e 100644 --- a/codex-rs/ext/goal/src/runtime.rs +++ b/codex-rs/ext/goal/src/runtime.rs @@ -216,11 +216,14 @@ impl GoalRuntimeHandle { .record_terminal_if_status_changed(previous_status, &goal); self.inner.accounting_state.clear_active_goal(); let goal = protocol_goal_from_state(goal); - self.inner.event_emitter.thread_goal_updated( - format!("{turn_id}:usage-limit"), - Some(turn_id.to_string()), - goal, - ); + self.inner + .event_emitter + .thread_goal_updated( + format!("{turn_id}:usage-limit"), + Some(turn_id.to_string()), + goal, + ) + .await; Ok(()) } @@ -380,11 +383,14 @@ impl GoalRuntimeHandle { budget_limited_goal_disposition, ); let goal = protocol_goal_from_state(goal); - self.inner.event_emitter.thread_goal_updated( - event_id.to_string(), - Some(turn_id.to_string()), - goal.clone(), - ); + self.inner + .event_emitter + .thread_goal_updated( + event_id.to_string(), + Some(turn_id.to_string()), + goal.clone(), + ) + .await; Some(AccountedGoalProgress { goal, goal_id }) } codex_state::GoalAccountingOutcome::Unchanged(_) => None, @@ -430,11 +436,10 @@ impl GoalRuntimeHandle { budget_limited_goal_disposition, ); let goal = protocol_goal_from_state(goal); - self.inner.event_emitter.thread_goal_updated( - event_id.to_string(), - /*turn_id*/ None, - goal.clone(), - ); + self.inner + .event_emitter + .thread_goal_updated(event_id.to_string(), /*turn_id*/ None, goal.clone()) + .await; Some(AccountedGoalProgress { goal, goal_id }) } codex_state::GoalAccountingOutcome::Unchanged(_) => { diff --git a/codex-rs/ext/goal/src/tool.rs b/codex-rs/ext/goal/src/tool.rs index 5dc547e521..6071a55103 100644 --- a/codex-rs/ext/goal/src/tool.rs +++ b/codex-rs/ext/goal/src/tool.rs @@ -201,7 +201,8 @@ impl GoalToolExecutor { .mark_current_turn_goal_active(goal.goal_id.clone()); self.metrics.record_created(); let goal = protocol_goal_from_state(goal); - self.emit_goal_updated_from_tool_call(&invocation, turn_id, goal.clone()); + self.emit_goal_updated_from_tool_call(&invocation, turn_id, goal.clone()) + .await; goal_response(Some(goal), CompletionBudgetReport::Omit) } @@ -261,7 +262,8 @@ impl GoalToolExecutor { .record_terminal_if_status_changed(previous_status, &goal); let goal = protocol_goal_from_state(goal); let turn_id = self.accounting_state.clear_current_turn_goal(); - self.emit_goal_updated_from_tool_call(&invocation, turn_id, goal.clone()); + self.emit_goal_updated_from_tool_call(&invocation, turn_id, goal.clone()) + .await; goal_response( Some(goal), if args.status == ThreadGoalStatus::Complete { @@ -272,14 +274,15 @@ impl GoalToolExecutor { ) } - fn emit_goal_updated_from_tool_call( + async fn emit_goal_updated_from_tool_call( &self, invocation: &ToolCall, turn_id: Option, goal: ThreadGoal, ) { self.event_emitter - .thread_goal_updated(invocation.call_id.clone(), turn_id, goal); + .thread_goal_updated(invocation.call_id.clone(), turn_id, goal) + .await; } async fn account_active_goal_progress( @@ -327,11 +330,9 @@ impl GoalToolExecutor { budget_limited_goal_disposition, ); let goal = protocol_goal_from_state(goal); - self.event_emitter.thread_goal_updated( - event_id.to_string(), - Some(turn_id), - goal.clone(), - ); + self.event_emitter + .thread_goal_updated(event_id.to_string(), Some(turn_id), goal.clone()) + .await; Some(goal) } codex_state::GoalAccountingOutcome::Unchanged(_) => None, diff --git a/codex-rs/ext/goal/tests/goal_extension_backend.rs b/codex-rs/ext/goal/tests/goal_extension_backend.rs index e2e2a3e8ff..9ecb167cbe 100644 --- a/codex-rs/ext/goal/tests/goal_extension_backend.rs +++ b/codex-rs/ext/goal/tests/goal_extension_backend.rs @@ -1435,8 +1435,9 @@ impl RecordingEventSink { } impl ExtensionEventSink for RecordingEventSink { - fn emit(&self, event: Event) { + fn emit<'a>(&'a self, event: Event) -> codex_extension_api::ExtensionEventFuture<'a> { self.events().push(event); + Box::pin(std::future::ready(())) } }