mirror of
https://github.com/openai/codex.git
synced 2026-05-27 22:44:23 +00:00
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:
@@ -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>,
|
||||
|
||||
@@ -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(());
|
||||
|
||||
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user