Simplify goal extension plumbing

This commit is contained in:
Eric Traut
2026-05-27 21:36:25 -07:00
parent 933771d891
commit 0471e80ece
12 changed files with 82 additions and 189 deletions
@@ -5,24 +5,16 @@ use codex_protocol::protocol::ThreadGoalUpdatedEvent;
pub type ExtensionEventFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
/// Extension-generated event with a host-owned delivery correlation id.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ExtensionEvent {
pub id: String,
pub msg: ExtensionEventMsg,
}
/// Events that extensions can ask the host to deliver.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ExtensionEventMsg {
pub enum ExtensionEvent {
ThreadGoalUpdated(ThreadGoalUpdatedEvent),
}
/// Host-provided sink for extension-generated events.
///
/// Extensions construct extension events with the correlation id appropriate
/// for the callback they are handling, then leave persistence, ordering,
/// transport fanout, and logging decisions to the host.
/// Extensions construct typed extension events, then leave persistence,
/// ordering, transport fanout, and logging decisions to the host.
pub trait ExtensionEventSink: Send + Sync {
/// Queue one extension event for host-owned delivery.
fn emit<'a>(&'a self, event: ExtensionEvent) -> ExtensionEventFuture<'a>;
@@ -6,7 +6,6 @@ pub use agent::AgentSpawnFuture;
pub use agent::AgentSpawner;
pub use events::ExtensionEvent;
pub use events::ExtensionEventFuture;
pub use events::ExtensionEventMsg;
pub use events::ExtensionEventSink;
pub use events::NoopExtensionEventSink;
pub use response_items::NoopResponseItemInjector;
-1
View File
@@ -8,7 +8,6 @@ pub use capabilities::AgentSpawnFuture;
pub use capabilities::AgentSpawner;
pub use capabilities::ExtensionEvent;
pub use capabilities::ExtensionEventFuture;
pub use capabilities::ExtensionEventMsg;
pub use capabilities::ExtensionEventSink;
pub use capabilities::NoopExtensionEventSink;
pub use capabilities::NoopResponseItemInjector;
@@ -127,7 +127,6 @@ impl<C: Sync> ExtensionRegistryBuilder<C> {
/// Finishes construction and returns the immutable registry.
pub fn build(self) -> ExtensionRegistry<C> {
ExtensionRegistry {
event_sink: self.event_sink,
thread_lifecycle_contributors: self.thread_lifecycle_contributors,
thread_idle_turn_contributors: self.thread_idle_turn_contributors,
turn_lifecycle_contributors: self.turn_lifecycle_contributors,
@@ -144,7 +143,6 @@ impl<C: Sync> ExtensionRegistryBuilder<C> {
/// Immutable typed registry produced after extensions are installed.
pub struct ExtensionRegistry<C: Sync> {
event_sink: Arc<dyn ExtensionEventSink>,
thread_lifecycle_contributors: Vec<Arc<dyn ThreadLifecycleContributor<C>>>,
thread_idle_turn_contributors: Vec<Arc<dyn ThreadIdleTurnContributor>>,
turn_lifecycle_contributors: Vec<Arc<dyn TurnLifecycleContributor>>,
@@ -158,11 +156,6 @@ pub struct ExtensionRegistry<C: Sync> {
}
impl<C: Sync> ExtensionRegistry<C> {
/// Returns the host event sink retained by this registry.
pub fn event_sink(&self) -> Arc<dyn ExtensionEventSink> {
Arc::clone(&self.event_sink)
}
/// Returns the registered thread-lifecycle contributors.
pub fn thread_lifecycle_contributors(&self) -> &[Arc<dyn ThreadLifecycleContributor<C>>] {
&self.thread_lifecycle_contributors
+6 -15
View File
@@ -1,7 +1,6 @@
use std::sync::Arc;
use codex_extension_api::ExtensionEvent;
use codex_extension_api::ExtensionEventMsg;
use codex_extension_api::ExtensionEventSink;
use codex_protocol::protocol::ThreadGoal;
use codex_protocol::protocol::ThreadGoalUpdatedEvent;
@@ -16,21 +15,13 @@ impl GoalEventEmitter {
Self { sink }
}
pub(crate) async fn thread_goal_updated(
&self,
event_id: impl Into<String>,
turn_id: Option<String>,
goal: ThreadGoal,
) {
pub(crate) async fn thread_goal_updated(&self, turn_id: Option<String>, goal: ThreadGoal) {
self.sink
.emit(ExtensionEvent {
id: event_id.into(),
msg: ExtensionEventMsg::ThreadGoalUpdated(ThreadGoalUpdatedEvent {
thread_id: goal.thread_id,
turn_id,
goal,
}),
})
.emit(ExtensionEvent::ThreadGoalUpdated(ThreadGoalUpdatedEvent {
thread_id: goal.thread_id,
turn_id,
goal,
}))
.await;
}
}
-3
View File
@@ -216,7 +216,6 @@ where
if let Err(err) = runtime
.account_active_goal_progress(
turn_id,
&format!("{turn_id}:turn-stop"),
codex_state::GoalAccountingMode::ActiveOnly,
BudgetLimitedGoalDisposition::ClearActive,
)
@@ -242,7 +241,6 @@ where
if let Err(err) = runtime
.account_active_goal_progress(
turn_id,
&format!("{turn_id}:turn-abort"),
codex_state::GoalAccountingMode::ActiveOnly,
BudgetLimitedGoalDisposition::ClearActive,
)
@@ -327,7 +325,6 @@ where
let progress = match runtime
.account_active_goal_progress(
turn_id,
input.call_id,
codex_state::GoalAccountingMode::ActiveOnly,
BudgetLimitedGoalDisposition::KeepActive,
)
+3 -17
View File
@@ -88,7 +88,6 @@ impl GoalRuntimeHandle {
if let Some(turn_id) = self.inner.accounting_state.current_turn_id() {
self.account_active_goal_progress(
turn_id.as_str(),
&format!("{turn_id}:external-goal-mutation"),
codex_state::GoalAccountingMode::ActiveOnly,
BudgetLimitedGoalDisposition::ClearActive,
)
@@ -97,7 +96,6 @@ impl GoalRuntimeHandle {
}
self.account_idle_goal_progress(
&format!("{}:external-goal-mutation", self.inner.thread_id),
codex_state::GoalAccountingMode::ActiveOnly,
BudgetLimitedGoalDisposition::ClearActive,
)
@@ -187,10 +185,8 @@ impl GoalRuntimeHandle {
return Ok(());
}
let progress_event_id = format!("{turn_id}:usage-limit-progress");
self.account_active_goal_progress(
turn_id,
progress_event_id.as_str(),
codex_state::GoalAccountingMode::ActiveOnly,
BudgetLimitedGoalDisposition::ClearActive,
)
@@ -216,11 +212,7 @@ impl GoalRuntimeHandle {
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,
)
.thread_goal_updated(Some(turn_id.to_string()), goal)
.await;
Ok(())
}
@@ -308,7 +300,6 @@ impl GoalRuntimeHandle {
pub(crate) async fn account_active_goal_progress(
&self,
turn_id: &str,
event_id: &str,
mode: codex_state::GoalAccountingMode,
budget_limited_goal_disposition: BudgetLimitedGoalDisposition,
) -> Result<Option<AccountedGoalProgress>, String> {
@@ -348,11 +339,7 @@ impl GoalRuntimeHandle {
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(),
)
.thread_goal_updated(Some(turn_id.to_string()), goal.clone())
.await;
Some(AccountedGoalProgress { goal, goal_id })
}
@@ -362,7 +349,6 @@ impl GoalRuntimeHandle {
async fn account_idle_goal_progress(
&self,
event_id: &str,
mode: codex_state::GoalAccountingMode,
budget_limited_goal_disposition: BudgetLimitedGoalDisposition,
) -> Result<Option<AccountedGoalProgress>, String> {
@@ -401,7 +387,7 @@ impl GoalRuntimeHandle {
let goal = protocol_goal_from_state(goal);
self.inner
.event_emitter
.thread_goal_updated(event_id.to_string(), /*turn_id*/ None, goal.clone())
.thread_goal_updated(/*turn_id*/ None, goal.clone())
.await;
Some(AccountedGoalProgress { goal, goal_id })
}
+5 -16
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.event_emitter
.thread_goal_updated(turn_id, goal.clone())
.await;
goal_response(Some(goal), CompletionBudgetReport::Omit)
}
@@ -230,7 +231,6 @@ impl GoalToolExecutor {
| ThreadGoalStatus::UsageLimited
| ThreadGoalStatus::BudgetLimited => unreachable!("status validated above"),
},
invocation.call_id.as_str(),
BudgetLimitedGoalDisposition::ClearActive,
)
.await?;
@@ -262,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.event_emitter
.thread_goal_updated(turn_id, goal.clone())
.await;
goal_response(
Some(goal),
@@ -274,21 +275,9 @@ impl GoalToolExecutor {
)
}
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)
.await;
}
async fn account_active_goal_progress(
&self,
mode: codex_state::GoalAccountingMode,
event_id: &str,
budget_limited_goal_disposition: BudgetLimitedGoalDisposition,
) -> Result<Option<ThreadGoal>, FunctionCallError> {
let Some(turn_id) = self.accounting_state.current_turn_id() else {
@@ -331,7 +320,7 @@ impl GoalToolExecutor {
);
let goal = protocol_goal_from_state(goal);
self.event_emitter
.thread_goal_updated(event_id.to_string(), Some(turn_id), goal.clone())
.thread_goal_updated(Some(turn_id), goal.clone())
.await;
Some(goal)
}
@@ -5,7 +5,6 @@ use std::time::Duration;
use codex_extension_api::ExtensionData;
use codex_extension_api::ExtensionEvent;
use codex_extension_api::ExtensionEventMsg;
use codex_extension_api::ExtensionEventSink;
use codex_extension_api::ExtensionRegistryBuilder;
use codex_extension_api::FunctionCallError;
@@ -263,7 +262,6 @@ async fn tool_finish_accounts_active_goal_progress_and_emits_event() -> anyhow::
assert_eq!(
vec![CapturedGoalEvent {
event_id: "call-shell".to_string(),
turn_id: Some("turn-1".to_string()),
status: ThreadGoalStatus::Active,
tokens_used: 23,
@@ -331,13 +329,11 @@ async fn budget_limited_goal_keeps_accruing_until_turn_stop() -> anyhow::Result<
assert_eq!(
vec![
CapturedGoalEvent {
event_id: "call-shell".to_string(),
turn_id: Some("turn-1".to_string()),
status: ThreadGoalStatus::BudgetLimited,
tokens_used: 25,
},
CapturedGoalEvent {
event_id: "turn-1:turn-stop".to_string(),
turn_id: Some("turn-1".to_string()),
status: ThreadGoalStatus::BudgetLimited,
tokens_used: 35,
@@ -449,13 +445,11 @@ async fn usage_limit_turn_error_accounts_and_marks_goal_terminal() -> anyhow::Re
assert_eq!(
vec![
CapturedGoalEvent {
event_id: "turn-1:usage-limit-progress".to_string(),
turn_id: Some("turn-1".to_string()),
status: ThreadGoalStatus::Active,
tokens_used: 23,
},
CapturedGoalEvent {
event_id: "turn-1:usage-limit".to_string(),
turn_id: Some("turn-1".to_string()),
status: ThreadGoalStatus::UsageLimited,
tokens_used: 23,
@@ -634,13 +628,11 @@ async fn usage_limit_budget_limited_goal_accounts_remaining_progress() -> anyhow
assert_eq!(
vec![
CapturedGoalEvent {
event_id: "turn-1:usage-limit-progress".to_string(),
turn_id: Some("turn-1".to_string()),
status: ThreadGoalStatus::BudgetLimited,
tokens_used: 35,
},
CapturedGoalEvent {
event_id: "turn-1:usage-limit".to_string(),
turn_id: Some("turn-1".to_string()),
status: ThreadGoalStatus::UsageLimited,
tokens_used: 35,
@@ -790,13 +782,11 @@ async fn update_goal_can_block_and_accounts_final_progress() -> anyhow::Result<(
assert_eq!(
vec![
CapturedGoalEvent {
event_id: "call-update-goal".to_string(),
turn_id: Some("turn-1".to_string()),
status: ThreadGoalStatus::Active,
tokens_used: 23,
},
CapturedGoalEvent {
event_id: "call-update-goal".to_string(),
turn_id: Some("turn-1".to_string()),
status: ThreadGoalStatus::Blocked,
tokens_used: 23,
@@ -876,13 +866,11 @@ async fn update_goal_can_complete_and_reports_final_budget() -> anyhow::Result<(
assert_eq!(
vec![
CapturedGoalEvent {
event_id: "call-update-goal".to_string(),
turn_id: Some("turn-1".to_string()),
status: ThreadGoalStatus::Active,
tokens_used: 23,
},
CapturedGoalEvent {
event_id: "call-update-goal".to_string(),
turn_id: Some("turn-1".to_string()),
status: ThreadGoalStatus::Complete,
tokens_used: 23,
@@ -935,7 +923,6 @@ async fn external_goal_mutation_start_accounts_active_goal_progress() -> anyhow:
assert_eq!(23, goal.tokens_used);
assert_eq!(
vec![CapturedGoalEvent {
event_id: "turn-1:external-goal-mutation".to_string(),
turn_id: Some("turn-1".to_string()),
status: ThreadGoalStatus::Active,
tokens_used: 23,
@@ -1093,7 +1080,6 @@ async fn idle_continuation_request_rehydrates_active_goal_idle_accounting() -> a
);
assert_eq!(
vec![CapturedGoalEvent {
event_id: format!("{thread_id}:external-goal-mutation"),
turn_id: None,
status: ThreadGoalStatus::Active,
tokens_used: 0,
@@ -1430,13 +1416,13 @@ impl RecordingEventSink {
fn goal_events(&self) -> Vec<CapturedGoalEvent> {
self.events()
.iter()
.filter_map(|event| match &event.msg {
ExtensionEventMsg::ThreadGoalUpdated(updated) => Some(CapturedGoalEvent {
event_id: event.id.clone(),
.map(|event| {
let ExtensionEvent::ThreadGoalUpdated(updated) = event;
CapturedGoalEvent {
turn_id: updated.turn_id.clone(),
status: updated.goal.status,
tokens_used: updated.goal.tokens_used,
}),
}
})
.collect()
}
@@ -1459,7 +1445,6 @@ impl ExtensionEventSink for RecordingEventSink {
#[derive(Debug, PartialEq, Eq)]
struct CapturedGoalEvent {
event_id: String,
turn_id: Option<String>,
status: ThreadGoalStatus,
tokens_used: i64,