openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-05 15:18:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

Simplify goal extension plumbing

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

34
codex-rs/app-server/src/extensions.rs
View File

@@ -11,7 +11,6 @@ use codex_extension_api::AgentSpawnFuture;
use codex_extension_api::AgentSpawner;
use codex_extension_api::ExtensionEvent;
use codex_extension_api::ExtensionEventFuture;
use codex_extension_api::ExtensionEventMsg;
use codex_extension_api::ExtensionEventSink;
use codex_extension_api::ExtensionRegistry;
use codex_extension_api::ExtensionRegistryBuilder;
@@ -60,8 +59,8 @@ struct AppServerExtensionEventSink {
impl ExtensionEventSink for AppServerExtensionEventSink {
fn emit<'a>(&'a self, event: ExtensionEvent) -> ExtensionEventFuture<'a> {
Box::pin(async move {
match event.msg {
ExtensionEventMsg::ThreadGoalUpdated(thread_goal_event) => {
match event {
ExtensionEvent::ThreadGoalUpdated(thread_goal_event) => {
let notification =
ServerNotification::ThreadGoalUpdated(ThreadGoalUpdatedNotification {
thread_id: thread_goal_event.thread_id.to_string(),
@@ -182,23 +181,20 @@ mod tests {
objective: &str,
turn_id: &str,
) -> ExtensionEvent {
ExtensionEvent {
id: "call-1".to_string(),
msg: ExtensionEventMsg::ThreadGoalUpdated(ThreadGoalUpdatedEvent {
ExtensionEvent::ThreadGoalUpdated(ThreadGoalUpdatedEvent {
thread_id,
turn_id: Some(turn_id.to_string()),
goal: ThreadGoal {
thread_id,
turn_id: Some(turn_id.to_string()),
goal: ThreadGoal {
thread_id,
objective: objective.to_string(),
status: ThreadGoalStatus::Active,
token_budget: Some(123),
tokens_used: 45,
time_used_seconds: 6,
created_at: 7,
updated_at: 8,
},
}),
}
objective: objective.to_string(),
status: ThreadGoalStatus::Active,
token_budget: Some(123),
tokens_used: 45,
time_used_seconds: 6,
created_at: 7,
updated_at: 8,
},
})
}
fn app_server_goal_update(

97
codex-rs/app-server/src/request_processors/thread_goal_processor.rs
View File

@@ -143,67 +143,42 @@ impl ThreadGoalRequestProcessor {
.get_thread_goal(thread_id)
.await
.map_err(|err| invalid_request(err.to_string()))?;
let (goal, previous_goal) = match (objective, existing_goal) {
(Some(objective), Some(existing_goal)) => {
let goal = state_db
.thread_goals()
.update_thread_goal(
thread_id,
codex_state::GoalUpdate {
objective: Some(objective.to_string()),
status,
token_budget: params.token_budget,
expected_goal_id: Some(existing_goal.goal_id.clone()),
},
)
.await
.map_err(|err| invalid_request(err.to_string()))?
.ok_or_else(|| {
invalid_request(format!(
"cannot update goal for thread {thread_id}: no goal exists"
))
})?;
(goal, Some(existing_goal))
}
(Some(objective), None) => {
let goal = state_db
.thread_goals()
.replace_thread_goal(
thread_id,
objective,
status.unwrap_or(codex_state::ThreadGoalStatus::Active),
params.token_budget.flatten(),
)
.await
.map_err(|err| invalid_request(err.to_string()))?;
(goal, None)
}
(None, Some(existing_goal)) => {
let goal = state_db
.thread_goals()
.update_thread_goal(
thread_id,
codex_state::GoalUpdate {
objective: None,
status,
token_budget: params.token_budget,
expected_goal_id: Some(existing_goal.goal_id.clone()),
},
)
.await
.map_err(|err| invalid_request(err.to_string()))?
.ok_or_else(|| {
invalid_request(format!(
"cannot update goal for thread {thread_id}: no goal exists"
))
})?;
(goal, Some(existing_goal))
}
(None, None) => {
return Err(invalid_request(format!(
"cannot update goal for thread {thread_id}: no goal exists"
)));
}
let no_goal_error = || {
invalid_request(format!(
"cannot update goal for thread {thread_id}: no goal exists"
))
};
let (goal, previous_goal) = if let Some(existing_goal) = existing_goal {
let goal = state_db
.thread_goals()
.update_thread_goal(
thread_id,
codex_state::GoalUpdate {
objective: objective.map(str::to_string),
status,
token_budget: params.token_budget,
expected_goal_id: Some(existing_goal.goal_id.clone()),
},
)
.await
.map_err(|err| invalid_request(err.to_string()))?
.ok_or_else(no_goal_error)?;
(goal, Some(existing_goal))
} else {
let Some(objective) = objective else {
return Err(no_goal_error());
};
let goal = state_db
.thread_goals()
.replace_thread_goal(
thread_id,
objective,
status.unwrap_or(codex_state::ThreadGoalStatus::Active),
params.token_budget.flatten(),
)
.await
.map_err(|err| invalid_request(err.to_string()))?;
(goal, None)
};
if should_set_thread_preview
&& let Err(err) = state_db

29
codex-rs/core/src/tasks/idle_extension.rs
View File

@@ -4,6 +4,7 @@ use std::sync::Arc;
use codex_extension_api::ResponseInjectionItem;
use codex_extension_api::ThreadIdleRequest;
use codex_extension_api::ThreadIdleTurnContributor;
use codex_protocol::config_types::CollaborationMode;
use codex_protocol::config_types::ModeKind;
@@ -124,19 +125,13 @@ async fn has_pending_work(session: &Session) -> bool {
}
struct IdleTurnCandidate {
contributor_index: usize,
contributor: Arc<dyn ThreadIdleTurnContributor>,
request: ThreadIdleRequest,
}
async fn next_idle_turn_candidate(session: &Session) -> Option<IdleTurnCandidate> {
let collaboration_mode = session.collaboration_mode().await;
for (contributor_index, contributor) in session
.services
.extensions
.thread_idle_turn_contributors()
.iter()
.enumerate()
{
for contributor in session.services.extensions.thread_idle_turn_contributors() {
if !idle_turn_policy_allows_mode(contributor.idle_turn_policy(), &collaboration_mode) {
continue;
}
@@ -151,7 +146,7 @@ async fn next_idle_turn_candidate(session: &Session) -> Option<IdleTurnCandidate
};
if is_non_empty_idle_input(&request.item) {
return Some(IdleTurnCandidate {
contributor_index,
contributor: Arc::clone(contributor),
request,
});
}
@@ -175,19 +170,15 @@ fn idle_turn_policy_allows_mode(
}
async fn should_start_idle_turn(session: &Session, candidate: &IdleTurnCandidate) -> bool {
let Some(contributor) = session
.services
.extensions
.thread_idle_turn_contributors()
.get(candidate.contributor_index)
else {
return false;
};
let collaboration_mode = session.collaboration_mode().await;
if !idle_turn_policy_allows_mode(contributor.idle_turn_policy(), &collaboration_mode) {
if !idle_turn_policy_allows_mode(
candidate.contributor.idle_turn_policy(),
&collaboration_mode,
) {
return false;
}
contributor
candidate
.contributor
.should_start_thread_idle_turn(codex_extension_api::ThreadIdleTurnStartInput {
request: &candidate.request,
session_store: &session.services.session_extension_data,

14
codex-rs/ext/extension-api/src/capabilities/events.rs
View File

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

1
codex-rs/ext/extension-api/src/capabilities/mod.rs
View File

@@ -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
codex-rs/ext/extension-api/src/lib.rs
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;

7
codex-rs/ext/extension-api/src/registry.rs
View File

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

21
codex-rs/ext/goal/src/events.rs
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
codex-rs/ext/goal/src/extension.rs
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,
)

20
codex-rs/ext/goal/src/runtime.rs
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 })
}

21
codex-rs/ext/goal/src/tool.rs
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)
}

23
codex-rs/ext/goal/tests/goal_extension_backend.rs
View File

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