feat: handle goal usage limits in goal extension (#24628)

## Why

The extracted goal runtime needs a host-callable path for turns that
stop because the workspace usage limit is reached. In that case, any
in-turn goal progress should be accounted before the goal becomes
terminal, and active goal accounting must be cleared so later
tool-finish or turn-stop handling does not keep charging usage to a
stopped goal.

## What changed

- Adds `GoalRuntimeHandle::usage_limit_active_goal_for_turn`, which
accounts current active-goal progress, marks the active or
budget-limited thread goal as `UsageLimited`, records terminal metrics
when the status changes, clears active goal accounting, and emits the
updated goal event.
- Covers both active and budget-limited goals in
`ext/goal/tests/goal_extension_backend.rs`, including the invariant that
later token/tool events do not add usage after the goal has been
usage-limited.

## Testing

- Added
`usage_limit_active_goal_accounts_progress_and_clears_accounting`.
- Added `usage_limit_budget_limited_goal_accounts_remaining_progress`.
This commit is contained in:
jif-oai
2026-05-27 13:00:06 +02:00
committed by GitHub
parent ca2e343f65
commit bc005029bd
3 changed files with 303 additions and 1 deletions

View File

@@ -83,6 +83,17 @@ impl GoalAccountingState {
self.inner().current_turn_id.clone()
}
pub(crate) fn turn_is_current_active_goal(&self, turn_id: &str) -> bool {
let inner = self.inner();
if inner.current_turn_id.as_deref() != Some(turn_id) {
return false;
}
let Some(turn) = inner.turns.get(turn_id) else {
return false;
};
turn.account_tokens && turn.active_goal_id.is_some()
}
pub(crate) fn record_token_usage(
&self,
turn_id: impl Into<String>,

View File

@@ -192,6 +192,54 @@ impl GoalRuntimeHandle {
Ok(())
}
pub async fn usage_limit_active_goal_for_turn(&self, turn_id: &str) -> Result<(), String> {
if !self.is_enabled() {
return Ok(());
}
if !self
.inner
.accounting_state
.turn_is_current_active_goal(turn_id)
{
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,
)
.await?;
let previous_status = self
.current_goal_status_for_metrics(/*expected_goal_id*/ None)
.await?;
let Some(goal) = self
.inner
.state_dbs
.thread_goals()
.usage_limit_active_thread_goal(self.thread_id())
.await
.map_err(|err| err.to_string())?
else {
return Ok(());
};
self.inner
.metrics
.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,
);
Ok(())
}
pub async fn restore_after_resume(&self) -> Result<(), String> {
if !self.is_enabled() {
return Ok(());

View File

@@ -363,6 +363,243 @@ async fn budget_limited_goal_keeps_accounting_after_later_tool_finish() -> anyho
Ok(())
}
#[tokio::test]
async fn usage_limit_active_goal_accounts_progress_and_clears_accounting() -> anyhow::Result<()> {
let runtime = test_runtime().await?;
let thread_id = test_thread_id()?;
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?;
harness.start_turn("turn-1", &TokenUsage::default()).await;
let tools = harness.tools();
let create_tool = tool_by_name(&tools, "create_goal");
create_tool
.handle(tool_call(
"create_goal",
"call-create-goal",
json!({ "objective": "ship goal extension backend" }),
))
.await?;
harness.sink.clear();
harness
.record_token_usage(
"turn-1",
&token_usage(
/*input_tokens*/ 20, /*cached_input_tokens*/ 5, /*output_tokens*/ 8,
/*reasoning_output_tokens*/ 2, /*total_tokens*/ 30,
),
)
.await;
harness
.runtime_handle()
.usage_limit_active_goal_for_turn("turn-1")
.await
.map_err(anyhow::Error::msg)?;
let goal = runtime
.thread_goals()
.get_thread_goal(thread_id)
.await?
.ok_or_else(|| anyhow::anyhow!("goal should exist"))?;
assert_eq!(23, goal.tokens_used);
assert_eq!(codex_state::ThreadGoalStatus::UsageLimited, goal.status);
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,
},
],
harness.sink.goal_events()
);
harness
.record_token_usage(
"turn-1",
&token_usage(
/*input_tokens*/ 50, /*cached_input_tokens*/ 5,
/*output_tokens*/ 20, /*reasoning_output_tokens*/ 0,
/*total_tokens*/ 70,
),
)
.await;
harness
.notify_tool_finish("turn-1", "call-shell-after-usage-limit", "shell")
.await;
harness.stop_turn("turn-1").await;
let goal = runtime
.thread_goals()
.get_thread_goal(thread_id)
.await?
.ok_or_else(|| anyhow::anyhow!("goal should exist"))?;
assert_eq!(23, goal.tokens_used);
assert_eq!(codex_state::ThreadGoalStatus::UsageLimited, goal.status);
Ok(())
}
#[tokio::test]
async fn usage_limit_budget_limited_goal_accounts_remaining_progress() -> anyhow::Result<()> {
let runtime = test_runtime().await?;
let thread_id = test_thread_id()?;
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?;
harness.start_turn("turn-1", &TokenUsage::default()).await;
let tools = harness.tools();
let create_tool = tool_by_name(&tools, "create_goal");
create_tool
.handle(tool_call(
"create_goal",
"call-create-goal",
json!({
"objective": "ship goal extension backend",
"token_budget": 25,
}),
))
.await?;
harness
.record_token_usage(
"turn-1",
&token_usage(
/*input_tokens*/ 20, /*cached_input_tokens*/ 5,
/*output_tokens*/ 10, /*reasoning_output_tokens*/ 0,
/*total_tokens*/ 30,
),
)
.await;
harness
.notify_tool_finish("turn-1", "call-shell", "shell")
.await;
harness.sink.clear();
harness
.record_token_usage(
"turn-1",
&token_usage(
/*input_tokens*/ 24, /*cached_input_tokens*/ 5,
/*output_tokens*/ 16, /*reasoning_output_tokens*/ 0,
/*total_tokens*/ 40,
),
)
.await;
harness
.runtime_handle()
.usage_limit_active_goal_for_turn("turn-1")
.await
.map_err(anyhow::Error::msg)?;
let goal = runtime
.thread_goals()
.get_thread_goal(thread_id)
.await?
.ok_or_else(|| anyhow::anyhow!("goal should exist"))?;
assert_eq!(35, goal.tokens_used);
assert_eq!(codex_state::ThreadGoalStatus::UsageLimited, goal.status);
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,
},
],
harness.sink.goal_events()
);
Ok(())
}
#[tokio::test]
async fn usage_limit_plan_turn_does_not_stop_goal() -> anyhow::Result<()> {
let runtime = test_runtime().await?;
let thread_id = test_thread_id()?;
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?;
let tools = harness.tools();
let create_tool = tool_by_name(&tools, "create_goal");
create_tool
.handle(tool_call(
"create_goal",
"call-create-goal",
json!({ "objective": "ship goal extension backend" }),
))
.await?;
harness
.start_turn_with_mode("turn-plan", ModeKind::Plan, &TokenUsage::default())
.await;
harness.sink.clear();
harness
.runtime_handle()
.usage_limit_active_goal_for_turn("turn-plan")
.await
.map_err(anyhow::Error::msg)?;
let goal = runtime
.thread_goals()
.get_thread_goal(thread_id)
.await?
.ok_or_else(|| anyhow::anyhow!("goal should exist"))?;
assert_eq!(codex_state::ThreadGoalStatus::Active, goal.status);
assert_eq!(Vec::<CapturedGoalEvent>::new(), harness.sink.goal_events());
Ok(())
}
#[tokio::test]
async fn usage_limit_stale_turn_does_not_stop_current_goal() -> anyhow::Result<()> {
let runtime = test_runtime().await?;
let thread_id = test_thread_id()?;
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?;
harness.start_turn("turn-1", &TokenUsage::default()).await;
let tools = harness.tools();
let create_tool = tool_by_name(&tools, "create_goal");
create_tool
.handle(tool_call(
"create_goal",
"call-create-goal",
json!({ "objective": "ship goal extension backend" }),
))
.await?;
harness.stop_turn("turn-1").await;
harness.start_turn("turn-2", &TokenUsage::default()).await;
harness.sink.clear();
harness
.runtime_handle()
.usage_limit_active_goal_for_turn("turn-1")
.await
.map_err(anyhow::Error::msg)?;
let goal = runtime
.thread_goals()
.get_thread_goal(thread_id)
.await?
.ok_or_else(|| anyhow::anyhow!("goal should exist"))?;
assert_eq!(codex_state::ThreadGoalStatus::Active, goal.status);
assert_eq!(Vec::<CapturedGoalEvent>::new(), harness.sink.goal_events());
Ok(())
}
#[tokio::test]
async fn update_goal_can_block_and_accounts_final_progress() -> anyhow::Result<()> {
let runtime = test_runtime().await?;
@@ -719,8 +956,14 @@ impl GoalExtensionHarness {
}
async fn start_turn(&self, turn_id: &str, usage: &TokenUsage) {
self.start_turn_with_mode(turn_id, ModeKind::Default, usage)
.await;
}
async fn start_turn_with_mode(&self, turn_id: &str, mode: ModeKind, usage: &TokenUsage) {
let turn_store = ExtensionData::new(turn_id);
let collaboration_mode = default_collaboration_mode();
let mut collaboration_mode = default_collaboration_mode();
collaboration_mode.mode = mode;
for contributor in self.registry.turn_lifecycle_contributors() {
contributor
.on_turn_start(TurnStartInput {