[codex] abort turns when rollout budgets expire (token budget 3/3) (#28707)

## Stack

Depends on #28494.

## Description

This PR propagates shared rollout-budget exhaustion through the existing
`CodexErr::TurnAborted` task result.

Each thread records its model usage against the same ledger. Once the
ledger is exhausted, that usage update and all later usage updates
return `TurnAborted`. The task wrapper emits the normal aborted-turn
event and lifecycle instead of completing the turn.

This is intentionally a soft boundary: there is no cross-thread
`Op::Interrupt` fanout. An in-flight thread can finish its current
response before it observes the exhausted ledger, but every thread
aborts at its next usage-accounting boundary.

## Tests

The integration coverage verifies that:

- the response that exhausts the budget aborts its turn;
- a later response also aborts because the shared ledger remains
exhausted; and
- sub-agent usage draws from the same shared ledger; and
- local and remote-v2 compaction abort without retrying or emitting a
generic error.

Local checks:

- `just test -p codex-core
exhausted_budget_aborts_current_and_later_turns`
- `just test -p codex-core subagent_usage_draws_from_the_shared_budget`
- `just test -p codex-core
abort_regular_task_emits_marker_before_turn_aborted`
- `just test -p codex-core
compaction_budget_exhaustion_aborts_without_error_or_retry`
- `just fix -p codex-core`
- `just fmt`
- `git diff --check`

The full workspace test suite was not run locally.
This commit is contained in:
rka-oai
2026-06-19 02:00:01 -07:00
committed by GitHub
parent 7abfcf220b
commit dac588f413
13 changed files with 266 additions and 69 deletions
+133
View File
@@ -4,6 +4,8 @@ use codex_features::Feature;
use codex_model_provider_info::built_in_model_providers;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::Op;
use codex_protocol::protocol::TurnAbortReason;
use codex_protocol::user_input::UserInput;
use core_test_support::responses::ResponsesRequest;
use core_test_support::responses::ev_assistant_message;
use core_test_support::responses::ev_completed;
@@ -20,6 +22,7 @@ use core_test_support::wait_for_event;
use pretty_assertions::assert_eq;
use serde_json::json;
use std::time::Duration;
use test_case::test_case;
use tokio::time::timeout;
const ROLLOUT_BUDGET: RolloutBudgetConfig = RolloutBudgetConfig {
@@ -192,6 +195,136 @@ async fn subagent_usage_draws_from_the_shared_budget() -> Result<()> {
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn exhausted_budget_aborts_current_and_later_turns() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
mount_sse_sequence(
&server,
vec![
sse(vec![
ev_response_created("exhaust-budget"),
ev_completed_with_tokens("exhaust-budget", /*total_tokens*/ 30),
]),
sse(vec![
ev_response_created("already-exhausted"),
ev_completed_with_tokens("already-exhausted", /*total_tokens*/ 1),
]),
],
)
.await;
let test = test_codex()
.with_config(|config| {
config.rollout_budget = Some(RolloutBudgetConfig {
limit_tokens: 30,
reminder_interval_tokens: 10,
..ROLLOUT_BUDGET
});
})
.build(&server)
.await?;
for prompt in ["exhaust the budget", "try another turn"] {
test.codex
.submit(Op::UserInput {
items: vec![UserInput::Text {
text: prompt.to_string(),
text_elements: Vec::new(),
}],
final_output_json_schema: None,
responsesapi_client_metadata: None,
additional_context: Default::default(),
thread_settings: Default::default(),
})
.await?;
let event = wait_for_event(&test.codex, |event| match event {
EventMsg::TurnAborted(_) => true,
EventMsg::TurnComplete(_) => {
panic!("exhausted budget completed the turn instead of aborting")
}
_ => false,
})
.await;
let EventMsg::TurnAborted(abort) = event else {
unreachable!("event filter only accepts TurnAborted")
};
assert_eq!(abort.reason, TurnAbortReason::Interrupted);
}
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[test_case(false ; "local")]
#[test_case(true ; "remote_v2")]
async fn compaction_budget_exhaustion_aborts_without_error_or_retry(remote_v2: bool) -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let compact_response = if remote_v2 {
sse(vec![
json!({
"type": "response.output_item.done",
"item": {
"type": "compaction",
"encrypted_content": "encrypted-summary",
}
}),
ev_completed_with_tokens("compact", /*total_tokens*/ 10),
])
} else {
sse(vec![
ev_response_created("compact"),
ev_assistant_message("compact-summary", "compact summary"),
ev_completed_with_tokens("compact", /*total_tokens*/ 10),
])
};
let responses = mount_sse_sequence(&server, vec![compact_response]).await;
let mut model_provider = built_in_model_providers(/*openai_base_url*/ None)["openai"].clone();
model_provider.base_url = Some(format!("{}/v1", server.uri()));
model_provider.supports_websockets = false;
if !remote_v2 {
model_provider.name = "OpenAI-compatible test provider".to_string();
}
let test = test_codex()
.with_config(move |config| {
config.model_provider = model_provider;
config.rollout_budget = Some(RolloutBudgetConfig {
limit_tokens: 10,
reminder_interval_tokens: 5,
..ROLLOUT_BUDGET
});
if remote_v2 {
config
.features
.enable(Feature::RemoteCompactionV2)
.expect("test config should allow remote compaction v2");
}
})
.build(&server)
.await?;
test.codex.submit(Op::Compact).await?;
let event = wait_for_event(&test.codex, |event| match event {
EventMsg::TurnAborted(_) => true,
EventMsg::Error(error) => panic!("budget exhaustion emitted an error: {}", error.message),
EventMsg::TurnComplete(_) => {
panic!("budget-exhausting compaction completed instead of aborting")
}
_ => false,
})
.await;
let EventMsg::TurnAborted(abort) = event else {
unreachable!("event filter only accepts TurnAborted")
};
assert_eq!(abort.reason, TurnAbortReason::Interrupted);
assert_eq!(responses.requests().len(), 1, "compaction should not retry");
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn restates_the_current_remainder_after_compaction() -> Result<()> {
skip_if_no_network!(Ok(()));