openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-08 20:40:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add goal extension API for thread goals

Browse Source
This commit is contained in:
jif-oai
2026-05-29 15:59:59 +01:00
parent c4f509101b
commit f76b1da0a6
5 changed files with 352 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/ext/goal/src/api.rs
+269
View File
@@ -0,0 +1,269 @@
use std::collections::HashMap;
use std::fmt;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::PoisonError;
use std::sync::Weak;
use codex_protocol::ThreadId;
use codex_protocol::protocol::ThreadGoal;
use codex_protocol::protocol::ThreadGoalStatus;
use codex_protocol::protocol::validate_thread_goal_objective;
use crate::runtime::GoalRuntimeHandle;
use crate::runtime::PreviousGoalSnapshot;
use crate::tool::fill_empty_thread_preview_if_possible;
use crate::tool::protocol_goal_from_state;
use crate::tool::state_status_from_protocol;
use crate::tool::validate_goal_budget;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum GoalApiError {
InvalidRequest(String),
Internal(String),
}
impl fmt::Display for GoalApiError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::InvalidRequest(message) | Self::Internal(message) => f.write_str(message),
}
}
}
impl std::error::Error for GoalApiError {}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum GoalObjectiveUpdate<'a> {
Keep,
Set(&'a str),
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum GoalTokenBudgetUpdate {
Keep,
Set(Option<i64>),
}
#[derive(Clone, Copy, Debug)]
pub struct GoalSetRequest<'a> {
pub thread_id: ThreadId,
pub objective: GoalObjectiveUpdate<'a>,
pub status: Option<ThreadGoalStatus>,
pub token_budget: GoalTokenBudgetUpdate,
}
#[derive(Clone, Debug)]
pub struct GoalSetOutcome {
pub goal: ThreadGoal,
state_goal: codex_state::ThreadGoal,
previous_goal: Option<PreviousGoalSnapshot>,
}
impl GoalSetOutcome {
pub async fn apply_runtime_effects(&self, goal_api: &GoalApi) {
if let Some(runtime) = goal_api.runtime_for_thread(self.goal.thread_id)
&& let Err(err) = runtime
.apply_external_goal_set(self.state_goal.clone(), self.previous_goal.clone())
.await
{
tracing::warn!("failed to apply external goal status runtime effects: {err}");
}
}
}
#[derive(Debug, Default)]
pub struct GoalApi {
runtimes: Mutex<HashMap<String, Weak<GoalRuntimeHandle>>>,
}
impl GoalApi {
pub fn new() -> Self {
Self::default()
}
pub async fn get_thread_goal(
&self,
state_db: &codex_state::StateRuntime,
thread_id: ThreadId,
) -> Result<Option<ThreadGoal>, GoalApiError> {
state_db
.thread_goals()
.get_thread_goal(thread_id)
.await
.map(|goal| goal.map(protocol_goal_from_state))
.map_err(|err| GoalApiError::Internal(format!("failed to read thread goal: {err}")))
}
pub async fn set_thread_goal(
&self,
state_db: &codex_state::StateRuntime,
request: GoalSetRequest<'_>,
) -> Result<GoalSetOutcome, GoalApiError> {
let GoalSetRequest {
thread_id,
objective,
status,
token_budget,
} = request;
let status = status.map(state_status_from_protocol);
let objective = match objective {
GoalObjectiveUpdate::Keep => None,
GoalObjectiveUpdate::Set(objective) => Some(objective.trim()),
};
let token_budget = match token_budget {
GoalTokenBudgetUpdate::Keep => None,
GoalTokenBudgetUpdate::Set(token_budget) => Some(token_budget),
};
if let Some(objective) = objective {
validate_thread_goal_objective(objective).map_err(GoalApiError::InvalidRequest)?;
}
if objective.is_some() || token_budget.is_some() {
validate_goal_budget(token_budget.flatten()).map_err(GoalApiError::InvalidRequest)?;
}
if let Some(runtime) = self.runtime_for_thread(thread_id)
&& let Err(err) = runtime.prepare_external_goal_mutation().await
{
tracing::warn!("failed to prepare external goal mutation: {err}");
}
let (goal, previous_goal) = if let Some(objective) = objective {
let existing_goal = state_db
.thread_goals()
.get_thread_goal(thread_id)
.await
.map_err(|err| {
GoalApiError::Internal(format!("failed to read thread goal: {err}"))
})?;
if let Some(existing_goal) = existing_goal.as_ref() {
let previous_goal = PreviousGoalSnapshot::from(existing_goal);
state_db
.thread_goals()
.update_thread_goal(
thread_id,
codex_state::GoalUpdate {
objective: Some(objective.to_string()),
status,
token_budget,
expected_goal_id: Some(existing_goal.goal_id.clone()),
},
)
.await
.map_err(|err| {
GoalApiError::Internal(format!("failed to update thread goal: {err}"))
})?
.ok_or_else(|| {
GoalApiError::InvalidRequest(format!(
"cannot update goal for thread {thread_id}: no goal exists"
))
})
.map(|goal| (goal, Some(previous_goal)))?
} else {
state_db
.thread_goals()
.replace_thread_goal(
thread_id,
objective,
status.unwrap_or(codex_state::ThreadGoalStatus::Active),
token_budget.flatten(),
)
.await
.map_err(|err| {
GoalApiError::Internal(format!("failed to replace thread goal: {err}"))
})
.map(|goal| (goal, None))?
}
} else {
let existing_goal = state_db
.thread_goals()
.get_thread_goal(thread_id)
.await
.map_err(|err| {
GoalApiError::Internal(format!("failed to read thread goal: {err}"))
})?
.ok_or_else(|| {
GoalApiError::InvalidRequest(format!(
"cannot update goal for thread {thread_id}: no goal exists"
))
})?;
let previous_goal = PreviousGoalSnapshot::from(&existing_goal);
state_db
.thread_goals()
.update_thread_goal(
thread_id,
codex_state::GoalUpdate {
objective: None,
status,
token_budget,
expected_goal_id: None,
},
)
.await
.map_err(|err| {
GoalApiError::Internal(format!("failed to update thread goal: {err}"))
})?
.ok_or_else(|| {
GoalApiError::InvalidRequest(format!(
"cannot update goal for thread {thread_id}: no goal exists"
))
})
.map(|goal| (goal, Some(previous_goal)))?
};
if objective.is_some() {
fill_empty_thread_preview_if_possible(state_db, thread_id, &goal).await;
}
Ok(GoalSetOutcome {
goal: protocol_goal_from_state(goal.clone()),
state_goal: goal,
previous_goal,
})
}
pub async fn clear_thread_goal(
&self,
state_db: &codex_state::StateRuntime,
thread_id: ThreadId,
) -> Result<bool, GoalApiError> {
if let Some(runtime) = self.runtime_for_thread(thread_id)
&& let Err(err) = runtime.prepare_external_goal_mutation().await
{
tracing::warn!("failed to prepare external goal mutation: {err}");
}
let cleared = state_db
.thread_goals()
.delete_thread_goal(thread_id)
.await
.map_err(|err| GoalApiError::Internal(format!("failed to clear thread goal: {err}")))?;
if cleared
&& let Some(runtime) = self.runtime_for_thread(thread_id)
&& let Err(err) = runtime.apply_external_goal_clear().await
{
tracing::warn!("failed to apply external goal clear runtime effects: {err}");
}
Ok(cleared)
}
pub(crate) fn register_runtime(&self, runtime: &Arc<GoalRuntimeHandle>) {
self.runtimes()
.insert(runtime.thread_id().to_string(), Arc::downgrade(runtime));
}
fn runtime_for_thread(&self, thread_id: ThreadId) -> Option<Arc<GoalRuntimeHandle>> {
let key = thread_id.to_string();
let runtime = self.runtimes().get(&key).and_then(Weak::upgrade);
if runtime.is_none() {
self.runtimes().remove(&key);
}
runtime
}
fn runtimes(&self) -> std::sync::MutexGuard<'_, HashMap<String, Weak<GoalRuntimeHandle>>> {
self.runtimes.lock().unwrap_or_else(PoisonError::into_inner)
}
}
codex-rs/ext/goal/src/extension.rs
+7
View File
@@ -30,6 +30,7 @@ use codex_protocol::protocol::TokenUsageInfo;
use crate::accounting::BudgetLimitedGoalDisposition;
use crate::accounting::GoalAccountingState;
use crate::api::GoalApi;
use crate::events::GoalEventEmitter;
use crate::metrics::GoalMetrics;
use crate::runtime::GoalRuntimeConfig;
@@ -55,6 +56,7 @@ pub struct GoalExtension<C> {
event_emitter: GoalEventEmitter,
metrics: GoalMetrics,
thread_manager: Weak<ThreadManager>,
goal_api: Arc<GoalApi>,
goals_enabled: Arc<dyn Fn(&C) -> bool + Send + Sync>,
}
@@ -70,6 +72,7 @@ impl<C> GoalExtension<C> {
event_sink: Arc<dyn ExtensionEventSink>,
metrics_client: Option<MetricsClient>,
thread_manager: Weak<ThreadManager>,
goal_api: Arc<GoalApi>,
goals_enabled: impl Fn(&C) -> bool + Send + Sync + 'static,
) -> Self {
Self {
@@ -77,6 +80,7 @@ impl<C> GoalExtension<C> {
event_emitter: GoalEventEmitter::new(event_sink),
metrics: GoalMetrics::new(metrics_client),
thread_manager,
goal_api,
goals_enabled: Arc::new(goals_enabled),
}
}
@@ -118,6 +122,7 @@ where
)
});
runtime.set_enabled(enabled);
self.goal_api.register_runtime(&runtime);
}
async fn on_thread_resume(&self, input: ThreadResumeInput<'_>) {
@@ -390,6 +395,7 @@ pub fn install_with_backend<C>(
state_dbs: Arc<codex_state::StateRuntime>,
metrics_client: Option<MetricsClient>,
thread_manager: Weak<ThreadManager>,
goal_api: Arc<GoalApi>,
goals_enabled: impl Fn(&C) -> bool + Send + Sync + 'static,
) where
C: Send + Sync + 'static,
@@ -399,6 +405,7 @@ pub fn install_with_backend<C>(
registry.event_sink(),
metrics_client,
thread_manager,
Arc::clone(&goal_api),
goals_enabled,
));
registry.thread_lifecycle_contributor(extension.clone());
codex-rs/ext/goal/src/lib.rs
+7
View File
@@ -5,6 +5,7 @@
//! accounting that can be represented with today's extension API.
mod accounting;
mod api;
mod events;
mod extension;
mod metrics;
@@ -13,6 +14,12 @@ mod spec;
mod steering;
mod tool;
pub use api::GoalApi;
pub use api::GoalApiError;
pub use api::GoalObjectiveUpdate;
pub use api::GoalSetOutcome;
pub use api::GoalSetRequest;
pub use api::GoalTokenBudgetUpdate;
pub use extension::GoalExtension;
pub use extension::GoalExtensionConfig;
pub use extension::install_with_backend;
codex-rs/ext/goal/src/tool.rs
+5 -3
View File
@@ -363,7 +363,7 @@ where
.map_err(|err| FunctionCallError::RespondToModel(err.to_string()))
}
fn validate_goal_budget(value: Option<i64>) -> Result<(), String> {
pub(crate) fn validate_goal_budget(value: Option<i64>) -> Result<(), String> {
if let Some(value) = value
&& value <= 0
{
@@ -402,7 +402,7 @@ impl GoalToolResponse {
}
}
async fn fill_empty_thread_preview_if_possible(
pub(crate) async fn fill_empty_thread_preview_if_possible(
state_db: &codex_state::StateRuntime,
thread_id: ThreadId,
goal: &codex_state::ThreadGoal,
@@ -441,7 +441,9 @@ fn protocol_status_from_state(status: codex_state::ThreadGoalStatus) -> ThreadGo
}
}
fn state_status_from_protocol(status: ThreadGoalStatus) -> codex_state::ThreadGoalStatus {
pub(crate) fn state_status_from_protocol(
status: ThreadGoalStatus,
) -> codex_state::ThreadGoalStatus {
match status {
ThreadGoalStatus::Active => codex_state::ThreadGoalStatus::Active,
ThreadGoalStatus::Paused => codex_state::ThreadGoalStatus::Paused,
codex-rs/ext/goal/tests/goal_extension_backend.rs
+64 -32
View File
@@ -19,8 +19,11 @@ use codex_extension_api::ToolFinishInput;
use codex_extension_api::ToolPayload;
use codex_extension_api::TurnStartInput;
use codex_extension_api::TurnStopInput;
use codex_goal_extension::GoalApi;
use codex_goal_extension::GoalObjectiveUpdate;
use codex_goal_extension::GoalRuntimeHandle;
use codex_goal_extension::PreviousGoalSnapshot;
use codex_goal_extension::GoalSetRequest;
use codex_goal_extension::GoalTokenBudgetUpdate;
use codex_goal_extension::install_with_backend;
use codex_protocol::ThreadId;
use codex_protocol::config_types::CollaborationMode;
@@ -769,7 +772,7 @@ async fn external_goal_mutation_start_accounts_active_goal_progress() -> anyhow:
}
#[tokio::test]
async fn external_goal_set_active_resets_baseline_without_live_thread() -> anyhow::Result<()> {
async fn goal_api_external_set_active_resets_baseline_without_live_thread() -> anyhow::Result<()> {
let runtime = test_runtime().await?;
let thread_id = test_thread_id()?;
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
@@ -806,38 +809,19 @@ async fn external_goal_set_active_resets_baseline_without_live_thread() -> anyho
),
)
.await;
harness
.runtime_handle()
.prepare_external_goal_mutation()
.await
.map_err(anyhow::Error::msg)?;
let previous_goal = runtime
.thread_goals()
.get_thread_goal(thread_id)
.await?
.ok_or_else(|| anyhow::anyhow!("goal should exist"))?;
let updated_goal = runtime
.thread_goals()
.update_thread_goal(
thread_id,
codex_state::GoalUpdate {
objective: Some("new objective".to_string()),
status: Some(codex_state::ThreadGoalStatus::Active),
token_budget: None,
expected_goal_id: Some(previous_goal.goal_id.clone()),
let outcome = harness
.goal_api
.set_thread_goal(
runtime.as_ref(),
GoalSetRequest {
thread_id,
objective: GoalObjectiveUpdate::Set("new objective"),
status: Some(ThreadGoalStatus::Active),
token_budget: GoalTokenBudgetUpdate::Keep,
},
)
.await?
.ok_or_else(|| anyhow::anyhow!("goal update should succeed"))?;
harness
.runtime_handle()
.apply_external_goal_set(
updated_goal,
Some(PreviousGoalSnapshot::from(&previous_goal)),
)
.await
.map_err(anyhow::Error::msg)?;
.await?;
outcome.apply_runtime_effects(&harness.goal_api).await;
harness
.record_token_usage(
@@ -908,6 +892,48 @@ async fn thread_resume_rehydrates_active_goal_idle_accounting() -> anyhow::Resul
Ok(())
}
#[tokio::test]
async fn goal_api_sets_gets_and_clears_thread_goal() -> anyhow::Result<()> {
let runtime = test_runtime().await?;
let thread_id = test_thread_id()?;
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
let api = GoalApi::new();
let set = api
.set_thread_goal(
runtime.as_ref(),
GoalSetRequest {
thread_id,
objective: GoalObjectiveUpdate::Set(" ship goal API ownership "),
status: None,
token_budget: GoalTokenBudgetUpdate::Set(Some(123)),
},
)
.await?;
let get = api
.get_thread_goal(runtime.as_ref(), thread_id)
.await?
.ok_or_else(|| anyhow::anyhow!("goal should exist"))?;
let metadata = runtime
.get_thread(thread_id)
.await?
.ok_or_else(|| anyhow::anyhow!("seeded thread metadata should exist"))?;
assert_eq!(set.goal, get);
assert_eq!("ship goal API ownership", get.objective);
assert_eq!(ThreadGoalStatus::Active, get.status);
assert_eq!(Some(123), get.token_budget);
assert_eq!(Some("ship goal API ownership"), metadata.preview.as_deref());
assert!(api.clear_thread_goal(runtime.as_ref(), thread_id).await?);
assert_eq!(
None,
api.get_thread_goal(runtime.as_ref(), thread_id).await?
);
assert!(!api.clear_thread_goal(runtime.as_ref(), thread_id).await?);
Ok(())
}
async fn installed_tools(
runtime: Arc<codex_state::StateRuntime>,
thread_id: ThreadId,
@@ -928,11 +954,13 @@ async fn installed_tools_with_start(
persistent_thread_state_available: bool,
) -> Vec<Arc<dyn ToolExecutor<ToolCall>>> {
let mut builder = ExtensionRegistryBuilder::<()>::new();
let goal_api = Arc::new(GoalApi::new());
install_with_backend(
&mut builder,
runtime,
/*metrics_client*/ None,
Weak::new(),
goal_api,
|_| true,
);
let registry = builder.build();
@@ -965,6 +993,7 @@ struct GoalExtensionHarness {
registry: codex_extension_api::ExtensionRegistry<()>,
session_store: ExtensionData,
thread_store: ExtensionData,
goal_api: Arc<GoalApi>,
sink: Arc<RecordingEventSink>,
}
@@ -975,11 +1004,13 @@ impl GoalExtensionHarness {
) -> anyhow::Result<Self> {
let sink = Arc::new(RecordingEventSink::default());
let mut builder = ExtensionRegistryBuilder::<()>::with_event_sink(sink.clone());
let goal_api = Arc::new(GoalApi::new());
install_with_backend(
&mut builder,
runtime,
/*metrics_client*/ None,
Weak::new(),
Arc::clone(&goal_api),
|_| true,
);
let registry = builder.build();
@@ -1001,6 +1032,7 @@ impl GoalExtensionHarness {
registry,
session_store,
thread_store,
goal_api,
sink,
})
}