mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
[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:
@@ -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(()));
|
||||
|
||||
Reference in New Issue
Block a user