Preserve extension event ordering

This commit is contained in:
Eric Traut
2026-05-27 12:30:41 -07:00
parent c8221f8c3f
commit 83121bf58d
8 changed files with 105 additions and 59 deletions

View File

@@ -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(),

View File

@@ -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<Box<dyn Future<Output = ()> + 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(()))
}
}

View File

@@ -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;

View File

@@ -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;

View File

@@ -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<String>,
turn_id: Option<String>,
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;
}
}

View File

@@ -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(_) => {

View File

@@ -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<String>,
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,

View File

@@ -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(()))
}
}