Wire app-server goal RPCs through GoalApi

This commit is contained in:
jif-oai
2026-05-29 16:20:06 +01:00
parent f76b1da0a6
commit 9d4a88c100
8 changed files with 46 additions and 528 deletions

1
codex-rs/Cargo.lock generated
View File

@@ -1926,6 +1926,7 @@ dependencies = [
"codex-file-search",
"codex-file-watcher",
"codex-git-utils",
"codex-goal-extension",
"codex-guardian",
"codex-hooks",
"codex-image-generation-extension",

View File

@@ -44,6 +44,7 @@ codex-features = { workspace = true }
codex-guardian = { workspace = true }
codex-git-utils = { workspace = true }
codex-file-watcher = { workspace = true }
codex-goal-extension = { workspace = true }
codex-hooks = { workspace = true }
codex-otel = { workspace = true }
codex-plugin = { workspace = true }

View File

@@ -274,8 +274,6 @@ use codex_config::loader::project_trust_key;
use codex_config::types::McpServerTransportConfig;
use codex_core::CodexThread;
use codex_core::CodexThreadSettingsOverrides;
use codex_core::ExternalGoalPreviousStatus;
use codex_core::ExternalGoalSet;
use codex_core::ForkSnapshot;
use codex_core::NewThread;
#[cfg(test)]

View File

@@ -1,5 +1,9 @@
use super::*;
use codex_protocol::protocol::validate_thread_goal_objective;
use codex_goal_extension::GoalApi;
use codex_goal_extension::GoalApiError;
use codex_goal_extension::GoalObjectiveUpdate;
use codex_goal_extension::GoalSetRequest;
use codex_goal_extension::GoalTokenBudgetUpdate;
#[derive(Clone)]
pub(crate) struct ThreadGoalRequestProcessor {
@@ -8,6 +12,7 @@ pub(crate) struct ThreadGoalRequestProcessor {
config: Arc<Config>,
thread_state_manager: ThreadStateManager,
state_db: Option<StateDbHandle>,
goal_api: Arc<GoalApi>,
}
impl ThreadGoalRequestProcessor {
@@ -24,6 +29,7 @@ impl ThreadGoalRequestProcessor {
config,
thread_state_manager,
state_db,
goal_api: Arc::new(GoalApi::new()),
}
}
@@ -134,106 +140,25 @@ impl ThreadGoalRequestProcessor {
let thread_state = thread_state.lock().await;
thread_state.listener_command_tx()
};
let status = params.status.map(thread_goal_status_to_state);
let objective = params.objective.as_deref().map(str::trim);
if let Some(objective) = objective {
validate_thread_goal_objective(objective).map_err(invalid_request)?;
}
if objective.is_some() || params.token_budget.is_some() {
validate_goal_budget(params.token_budget.flatten()).map_err(invalid_request)?;
}
if let Some(thread) = running_thread.as_ref() {
thread.prepare_external_goal_mutation().await;
}
let should_set_thread_preview = objective.is_some();
let (goal, previous_status) = (if let Some(objective) = objective {
let existing_goal = state_db
.thread_goals()
.get_thread_goal(thread_id)
.await
.map_err(|err| invalid_request(err.to_string()))?;
if let Some(goal) = existing_goal.as_ref() {
let previous_status = ExternalGoalPreviousStatus::from(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(goal.goal_id.clone()),
},
)
.await
.and_then(|goal| {
goal.ok_or_else(|| {
anyhow::anyhow!(
"cannot update goal for thread {thread_id}: no goal exists"
)
})
})
.map(|goal| (goal, previous_status))
} else {
let previous_status = ExternalGoalPreviousStatus::NewGoal;
state_db
.thread_goals()
.replace_thread_goal(
thread_id,
objective,
status.unwrap_or(codex_state::ThreadGoalStatus::Active),
params.token_budget.flatten(),
)
.await
.map(|goal| (goal, previous_status))
}
} else {
let existing_goal = state_db
.thread_goals()
.get_thread_goal(thread_id)
.await
.map_err(|err| invalid_request(err.to_string()))?;
let Some(existing_goal) = existing_goal else {
return Err(invalid_request(format!(
"cannot update goal for thread {thread_id}: no goal exists"
)));
};
let previous_status = ExternalGoalPreviousStatus::from(&existing_goal);
state_db
.thread_goals()
.update_thread_goal(
let outcome = self
.goal_api
.set_thread_goal(
&state_db,
GoalSetRequest {
thread_id,
codex_state::GoalUpdate {
objective: None,
status,
token_budget: params.token_budget,
expected_goal_id: None,
},
)
.await
.and_then(|goal| {
goal.ok_or_else(|| {
anyhow::anyhow!("cannot update goal for thread {thread_id}: no goal exists")
})
})
.map(|goal| (goal, previous_status))
})
.map_err(|err| invalid_request(err.to_string()))?;
if should_set_thread_preview
&& let Err(err) = state_db
.set_thread_preview_if_empty(thread_id, goal.objective.as_str())
.await
{
warn!("failed to set empty thread preview from goal objective for {thread_id}: {err}");
}
let external_goal_set = ExternalGoalSet {
goal: goal.clone(),
previous_status,
};
let goal = api_thread_goal_from_state(goal);
objective: params
.objective
.as_deref()
.map_or(GoalObjectiveUpdate::Keep, GoalObjectiveUpdate::Set),
status: params.status.map(ThreadGoalStatus::to_core),
token_budget: params
.token_budget
.map_or(GoalTokenBudgetUpdate::Keep, GoalTokenBudgetUpdate::Set),
},
)
.await
.map_err(goal_api_error)?;
let goal: ThreadGoal = outcome.goal.clone().into();
self.outgoing
.send_response(
request_id.clone(),
@@ -242,9 +167,7 @@ impl ThreadGoalRequestProcessor {
.await;
self.emit_thread_goal_updated_ordered(thread_id, goal, listener_command_tx)
.await;
if let Some(thread) = running_thread.as_ref() {
thread.apply_external_goal_set(external_goal_set).await;
}
outcome.apply_runtime_effects(&self.goal_api).await;
Ok(())
}
@@ -258,12 +181,12 @@ impl ThreadGoalRequestProcessor {
let thread_id = parse_thread_id_for_request(params.thread_id.as_str())?;
let state_db = self.state_db_for_materialized_thread(thread_id).await?;
let goal = state_db
.thread_goals()
.get_thread_goal(thread_id)
let goal = self
.goal_api
.get_thread_goal(&state_db, thread_id)
.await
.map_err(|err| internal_error(format!("failed to read thread goal: {err}")))?
.map(api_thread_goal_from_state);
.map_err(goal_api_error)?
.map(ThreadGoal::from);
Ok(ThreadGoalGetResponse { goal })
}
@@ -307,24 +230,16 @@ impl ThreadGoalRequestProcessor {
)
.await;
if let Some(thread) = running_thread.as_ref() {
thread.prepare_external_goal_mutation().await;
}
let listener_command_tx = {
let thread_state = self.thread_state_manager.thread_state(thread_id).await;
let thread_state = thread_state.lock().await;
thread_state.listener_command_tx()
};
let cleared = state_db
.thread_goals()
.delete_thread_goal(thread_id)
let cleared = self
.goal_api
.clear_thread_goal(&state_db, thread_id)
.await
.map_err(|err| internal_error(format!("failed to clear thread goal: {err}")))?;
if cleared && let Some(thread) = running_thread.as_ref() {
thread.apply_external_goal_clear().await;
}
.map_err(goal_api_error)?;
self.outgoing
.send_response(request_id, ThreadGoalClearResponse { cleared })
@@ -449,26 +364,6 @@ impl ThreadGoalRequestProcessor {
}
}
fn validate_goal_budget(value: Option<i64>) -> Result<(), String> {
if let Some(value) = value
&& value <= 0
{
return Err("goal budgets must be positive when provided".to_string());
}
Ok(())
}
fn thread_goal_status_to_state(status: ThreadGoalStatus) -> codex_state::ThreadGoalStatus {
match status {
ThreadGoalStatus::Active => codex_state::ThreadGoalStatus::Active,
ThreadGoalStatus::Paused => codex_state::ThreadGoalStatus::Paused,
ThreadGoalStatus::Blocked => codex_state::ThreadGoalStatus::Blocked,
ThreadGoalStatus::UsageLimited => codex_state::ThreadGoalStatus::UsageLimited,
ThreadGoalStatus::BudgetLimited => codex_state::ThreadGoalStatus::BudgetLimited,
ThreadGoalStatus::Complete => codex_state::ThreadGoalStatus::Complete,
}
}
fn thread_goal_status_from_state(status: codex_state::ThreadGoalStatus) -> ThreadGoalStatus {
match status {
codex_state::ThreadGoalStatus::Active => ThreadGoalStatus::Active,
@@ -497,3 +392,10 @@ fn parse_thread_id_for_request(thread_id: &str) -> Result<ThreadId, JSONRPCError
ThreadId::from_string(thread_id)
.map_err(|err| invalid_request(format!("invalid thread id: {err}")))
}
fn goal_api_error(err: GoalApiError) -> JSONRPCErrorError {
match err {
GoalApiError::InvalidRequest(message) => invalid_request(message),
GoalApiError::Internal(message) => internal_error(message),
}
}

View File

@@ -1,6 +1,5 @@
use crate::agent::AgentStatus;
use crate::config::ConstraintResult;
use crate::goals::ExternalGoalSet;
use crate::goals::GoalRuntimeEvent;
use crate::session::Codex;
use crate::session::SessionSettingsUpdate;
@@ -178,39 +177,6 @@ impl CodexThread {
.await
}
pub async fn prepare_external_goal_mutation(&self) {
if let Err(err) = self
.codex
.session
.goal_runtime_apply(GoalRuntimeEvent::ExternalMutationStarting)
.await
{
tracing::warn!("failed to prepare external goal mutation: {err}");
}
}
pub async fn apply_external_goal_set(&self, external_set: ExternalGoalSet) {
if let Err(err) = self
.codex
.session
.goal_runtime_apply(GoalRuntimeEvent::ExternalSet { external_set })
.await
{
tracing::warn!("failed to apply external goal status runtime effects: {err}");
}
}
pub async fn apply_external_goal_clear(&self) {
if let Err(err) = self
.codex
.session
.goal_runtime_apply(GoalRuntimeEvent::ExternalClear)
.await
{
tracing::warn!("failed to apply external goal clear runtime effects: {err}");
}
}
#[doc(hidden)]
pub async fn ensure_rollout_materialized(&self) {
self.codex.session.ensure_rollout_materialized().await;

View File

@@ -73,6 +73,7 @@ static BUDGET_LIMIT_PROMPT_TEMPLATE: LazyLock<Template> =
},
);
#[cfg(test)]
static OBJECTIVE_UPDATED_PROMPT_TEMPLATE: LazyLock<Template> = LazyLock::new(|| {
match Template::parse(include_str!("../templates/goals/objective_updated.md")) {
Ok(template) => template,
@@ -94,44 +95,6 @@ enum TerminalMetricEmission {
Suppress,
}
/// Describes whether an external goal mutation created a new logical goal or
/// updated an existing one.
#[derive(Clone)]
pub enum ExternalGoalPreviousStatus {
NewGoal,
Existing(ExternalGoalPreviousGoal),
}
#[derive(Clone)]
pub struct ExternalGoalPreviousGoal {
goal_id: String,
status: codex_state::ThreadGoalStatus,
objective: String,
}
impl From<&codex_state::ThreadGoal> for ExternalGoalPreviousStatus {
fn from(goal: &codex_state::ThreadGoal) -> Self {
Self::Existing(ExternalGoalPreviousGoal::from(goal))
}
}
impl From<&codex_state::ThreadGoal> for ExternalGoalPreviousGoal {
fn from(goal: &codex_state::ThreadGoal) -> Self {
Self {
goal_id: goal.goal_id.clone(),
status: goal.status,
objective: goal.objective.clone(),
}
}
}
/// Runtime effects for an externally persisted goal mutation.
#[derive(Clone)]
pub struct ExternalGoalSet {
pub goal: codex_state::ThreadGoal,
pub previous_status: ExternalGoalPreviousStatus,
}
/// Runtime lifecycle events that can affect goal accounting, scheduling, or
/// model-visible steering.
///
@@ -160,11 +123,6 @@ pub(crate) enum GoalRuntimeEvent<'a> {
UsageLimitReached {
turn_context: &'a TurnContext,
},
ExternalMutationStarting,
ExternalSet {
external_set: ExternalGoalSet,
},
ExternalClear,
ThreadResumed,
}
@@ -396,22 +354,6 @@ impl Session {
.await?;
Ok(())
}),
GoalRuntimeEvent::ExternalMutationStarting => Box::pin(async move {
if let Err(err) = self.account_thread_goal_before_external_mutation().await {
tracing::warn!(
"failed to account thread goal progress before external mutation: {err}"
);
}
Ok(())
}),
GoalRuntimeEvent::ExternalSet { external_set } => Box::pin(async move {
self.apply_external_thread_goal_status(external_set).await;
Ok(())
}),
GoalRuntimeEvent::ExternalClear => Box::pin(async move {
self.clear_stopped_thread_goal_runtime_state().await;
Ok(())
}),
GoalRuntimeEvent::ThreadResumed => Box::pin(async move {
self.restore_thread_goal_runtime_after_resume().await?;
Ok(())
@@ -650,65 +592,6 @@ impl Session {
Ok(goal)
}
async fn apply_external_thread_goal_status(self: &Arc<Self>, external_set: ExternalGoalSet) {
let ExternalGoalSet {
goal,
previous_status,
} = external_set;
let previous_goal = match previous_status {
ExternalGoalPreviousStatus::NewGoal => None,
ExternalGoalPreviousStatus::Existing(goal) => Some(goal),
};
let replaced_existing_goal = previous_goal
.as_ref()
.is_some_and(|previous_goal| previous_goal.goal_id != goal.goal_id);
if previous_goal.is_none() || replaced_existing_goal {
self.emit_goal_created_metric();
}
let objective_changed = previous_goal
.as_ref()
.is_some_and(|previous_goal| previous_goal.objective != goal.objective);
let previous_status = previous_goal
.as_ref()
.and_then(|previous_goal| (!replaced_existing_goal).then_some(previous_goal.status));
self.emit_goal_resumed_metric_if_status_changed(previous_status, goal.status);
self.emit_goal_terminal_metrics_if_status_changed(previous_status, &goal);
let goal_for_steering = objective_changed.then(|| protocol_goal_from_state(goal.clone()));
let goal_id = goal.goal_id;
let status = goal.status;
match status {
codex_state::ThreadGoalStatus::Active => {
let turn_id = self
.active_turn_context()
.await
.map(|turn_context| turn_context.sub_id.clone());
let current_token_usage = self.total_token_usage().await.unwrap_or_default();
self.mark_active_goal_accounting(goal_id, turn_id, current_token_usage)
.await;
if let Some(goal) = goal_for_steering {
let item = goal_context_input_item(objective_updated_prompt(&goal));
if self.inject_if_running(vec![item]).await.is_err() {
tracing::debug!(
"skipping objective-updated goal steering because no turn is active"
);
}
}
self.maybe_continue_goal_if_idle_runtime().await;
}
codex_state::ThreadGoalStatus::BudgetLimited => {
if self.active_turn_context().await.is_none() {
self.clear_stopped_thread_goal_runtime_state().await;
}
}
codex_state::ThreadGoalStatus::Paused
| codex_state::ThreadGoalStatus::Blocked
| codex_state::ThreadGoalStatus::UsageLimited
| codex_state::ThreadGoalStatus::Complete => {
self.clear_stopped_thread_goal_runtime_state().await;
}
}
}
async fn clear_stopped_thread_goal_runtime_state(&self) {
*self.goal_runtime.budget_limit_reported_goal_id.lock().await = None;
let mut accounting = self.goal_runtime.accounting.lock().await;
@@ -832,14 +715,6 @@ impl Session {
}))
}
async fn active_turn_context(&self) -> Option<Arc<TurnContext>> {
let active = self.active_turn.lock().await;
active
.as_ref()
.and_then(|active_turn| active_turn.task.as_ref())
.map(|task| Arc::clone(&task.turn_context))
}
async fn mark_thread_goal_turn_started(
&self,
turn_context: &TurnContext,
@@ -1083,29 +958,6 @@ impl Session {
Ok(())
}
async fn account_thread_goal_before_external_mutation(&self) -> anyhow::Result<()> {
if let Some(turn_context) = self.active_turn_context().await {
return self
.account_thread_goal_progress(
turn_context.as_ref(),
BudgetLimitSteering::Suppressed,
TerminalMetricEmission::Emit,
)
.await;
}
let Some(state_db) = self.state_db_for_thread_goals().await? else {
return Ok(());
};
self.account_thread_goal_wall_clock_usage(
&state_db,
codex_state::GoalAccountingMode::ActiveOnly,
TerminalMetricEmission::Suppress,
)
.await?;
Ok(())
}
async fn account_thread_goal_wall_clock_usage(
&self,
state_db: &StateDbHandle,
@@ -1563,6 +1415,7 @@ fn budget_limit_prompt(goal: &ThreadGoal) -> String {
}
}
#[cfg(test)]
fn objective_updated_prompt(goal: &ThreadGoal) -> String {
let token_budget = goal
.token_budget

View File

@@ -38,8 +38,6 @@ mod exec_policy;
#[cfg(test)]
mod git_info_tests;
mod goals;
pub use goals::ExternalGoalPreviousStatus;
pub use goals::ExternalGoalSet;
mod guardian;
mod hook_runtime;
mod installation_id;

View File

@@ -57,8 +57,6 @@ use codex_protocol::request_permissions::RequestPermissionProfile;
use tracing::Span;
use crate::goals::CreateGoalRequest;
use crate::goals::ExternalGoalPreviousStatus;
use crate::goals::ExternalGoalSet;
use crate::goals::GoalRuntimeEvent;
use crate::goals::SetGoalRequest;
use crate::rollout::recorder::RolloutRecorder;
@@ -9115,205 +9113,6 @@ async fn usage_limit_runtime_stops_active_goal_and_prevents_idle_continuation()
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn external_goal_mutation_accounts_active_turn_before_status_change() -> anyhow::Result<()> {
let (sess, tc, _rx, _codex_home) = make_goal_session_and_context_with_rx().await;
sess.set_thread_goal(
tc.as_ref(),
SetGoalRequest {
objective: Some("Keep improving the benchmark".to_string()),
status: None,
token_budget: None,
},
)
.await?;
sess.spawn_task(
Arc::clone(&tc),
Vec::new(),
NeverEndingTask {
kind: TaskKind::Regular,
listen_to_cancellation_token: false,
},
)
.await;
set_total_token_usage(&sess, post_goal_token_usage()).await;
sess.goal_runtime_apply(GoalRuntimeEvent::ExternalMutationStarting)
.await?;
let state_db = goal_test_state_db(sess.as_ref()).await?;
let goal = state_db
.thread_goals()
.get_thread_goal(sess.conversation_id)
.await?
.expect("goal should remain persisted");
assert_eq!(70, goal.tokens_used);
let previous_goal = goal.clone();
let goal_id = goal.goal_id.clone();
let updated_goal = state_db
.thread_goals()
.update_thread_goal(
sess.conversation_id,
codex_state::GoalUpdate {
objective: None,
status: Some(codex_state::ThreadGoalStatus::Complete),
token_budget: None,
expected_goal_id: Some(goal_id),
},
)
.await?
.expect("goal status update should succeed");
sess.goal_runtime_apply(GoalRuntimeEvent::ExternalSet {
external_set: ExternalGoalSet {
goal: updated_goal,
previous_status: ExternalGoalPreviousStatus::from(&previous_goal),
},
})
.await?;
assert!(sess.active_turn.lock().await.is_some());
let goal = state_db
.thread_goals()
.get_thread_goal(sess.conversation_id)
.await?
.expect("goal should remain persisted");
assert_eq!(codex_state::ThreadGoalStatus::Complete, goal.status);
assert_eq!(70, goal.tokens_used);
sess.abort_all_tasks(TurnAbortReason::Replaced).await;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn external_objective_change_steers_active_turn() -> anyhow::Result<()> {
let (sess, tc, _rx, _codex_home) = make_goal_session_and_context_with_rx().await;
sess.spawn_task(
Arc::clone(&tc),
Vec::new(),
NeverEndingTask {
kind: TaskKind::Regular,
listen_to_cancellation_token: false,
},
)
.await;
let state_db = goal_test_state_db(sess.as_ref()).await?;
let old_goal = state_db
.thread_goals()
.replace_thread_goal(
sess.conversation_id,
"Keep improving the benchmark",
codex_state::ThreadGoalStatus::Active,
/*token_budget*/ Some(10_000),
)
.await?;
let new_goal = state_db
.thread_goals()
.replace_thread_goal(
sess.conversation_id,
"Write a concise benchmark summary",
codex_state::ThreadGoalStatus::Active,
/*token_budget*/ Some(10_000),
)
.await?;
sess.goal_runtime_apply(GoalRuntimeEvent::ExternalSet {
external_set: ExternalGoalSet {
goal: new_goal,
previous_status: ExternalGoalPreviousStatus::from(&old_goal),
},
})
.await?;
let pending_input = sess.input_queue.get_pending_input(&sess.active_turn).await;
assert!(
pending_input.iter().any(|item| {
matches!(
item,
TurnInput::ResponseItem(ResponseItem::Message { role, content, .. })
if role == "user"
&& content.iter().any(|content| matches!(
content,
ContentItem::InputText { text }
if text.starts_with("<codex_internal_context source=\"goal\">")
&& text.trim_end().ends_with("</codex_internal_context>")
&& text.contains("The active thread goal objective was edited")
&& text.contains("Write a concise benchmark summary")
))
)
}),
"expected objective-updated steering prompt in pending input: {pending_input:?}"
);
sess.abort_all_tasks(TurnAbortReason::Replaced).await;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn external_active_goal_set_marks_current_turn_for_accounting() -> anyhow::Result<()> {
let (sess, tc, _rx, _codex_home) = make_goal_session_and_context_with_rx().await;
sess.spawn_task(
Arc::clone(&tc),
Vec::new(),
NeverEndingTask {
kind: TaskKind::Regular,
listen_to_cancellation_token: false,
},
)
.await;
set_total_token_usage(&sess, post_goal_token_usage()).await;
let state_db = goal_test_state_db(sess.as_ref()).await?;
let goal = state_db
.thread_goals()
.replace_thread_goal(
sess.conversation_id,
"Keep improving the benchmark",
codex_state::ThreadGoalStatus::Active,
/*token_budget*/ None,
)
.await?;
sess.goal_runtime_apply(GoalRuntimeEvent::ExternalSet {
external_set: ExternalGoalSet {
goal,
previous_status: ExternalGoalPreviousStatus::NewGoal,
},
})
.await?;
set_total_token_usage(
&sess,
TokenUsage {
input_tokens: 65,
cached_input_tokens: 10,
output_tokens: 40,
reasoning_output_tokens: 5,
total_tokens: 110,
},
)
.await;
sess.goal_runtime_apply(GoalRuntimeEvent::ToolCompleted {
turn_context: tc.as_ref(),
tool_name: "shell_command",
})
.await?;
let goal = state_db
.thread_goals()
.get_thread_goal(sess.conversation_id)
.await?
.expect("goal should remain persisted");
assert_eq!(codex_state::ThreadGoalStatus::Active, goal.status);
assert_eq!(25, goal.tokens_used);
sess.abort_all_tasks(TurnAbortReason::Replaced).await;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn completed_goal_accounts_current_turn_tokens_before_tool_response() -> anyhow::Result<()> {
let server = start_mock_server().await;