[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 authored and GitHub committed 2026-06-19 02:00:01 -07:00
1 parent 7abfcf220b
commit dac588f413
13 files changed
+266 -69

No files matched your search

+46 -21
View File
@@ -54,6 +54,8 @@ use codex_protocol::protocol::TurnCompleteEvent;
use codex_protocol::protocol::WarningEvent;
use codex_features::Feature;
use codex_protocol::error::CodexErr;
use codex_protocol::error::Result as CodexResult;
use codex_protocol::models::ContentItem;
pub(crate) use compact::CompactTask;
pub(crate) use regular::RegularTask;
@@ -65,6 +67,8 @@ pub(crate) use user_shell::execute_user_shell_command;
const GRACEFULL_INTERRUPTION_TIMEOUT_MS: u64 = 100;
const TASK_COMPACT_METRIC: &str = "codex.task.compact";
pub(crate) type SessionTaskResult = CodexResult<Option<String>>;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum InterruptedTurnHistoryMarker {
Disabled,
@@ -222,14 +226,16 @@ pub(crate) trait SessionTask: Send + Sync + 'static {
/// provided `cancellation_token` is cancelled when the session requests an
/// abort; implementers should watch for it and terminate quickly once it
/// fires. Returning [`Some`] yields a final message that
/// [`Session::on_task_finished`] will emit to the client.
/// [`Session::on_task_finished`] will emit to the client. Returning
/// [`CodexErr::TurnAborted`] completes the task through the aborted-turn
/// lifecycle instead.
fn run(
self: Arc<Self>,
session: Arc<SessionTaskContext>,
ctx: Arc<TurnContext>,
input: Vec<TurnInput>,
cancellation_token: CancellationToken,
) -> impl std::future::Future<Output = Option<String>> + Send;
) -> impl std::future::Future<Output = SessionTaskResult> + Send;
/// Gives the task a chance to perform cleanup after an abort.
///
@@ -258,7 +264,7 @@ pub(crate) trait AnySessionTask: Send + Sync + 'static {
ctx: Arc<TurnContext>,
input: Vec<TurnInput>,
cancellation_token: CancellationToken,
) -> BoxFuture<'static, Option<String>>;
) -> BoxFuture<'static, SessionTaskResult>;
fn abort<'a>(
&'a self,
@@ -285,7 +291,7 @@ where
ctx: Arc<TurnContext>,
input: Vec<TurnInput>,
cancellation_token: CancellationToken,
) -> BoxFuture<'static, Option<String>> {
) -> BoxFuture<'static, SessionTaskResult> {
Box::pin(SessionTask::run(
self,
session,
@@ -395,7 +401,7 @@ impl Session {
let handle = tokio::spawn(
async move {
let ctx_for_finish = Arc::clone(&ctx);
let last_agent_message = task_for_run
let task_result = task_for_run
.run(
Arc::clone(&session_ctx),
ctx,
@@ -418,8 +424,8 @@ impl Session {
.await;
}
if !task_cancellation_token.is_cancelled() {
// Emit completion uniformly from spawn site so all tasks share the same lifecycle.
sess.on_task_finished(Arc::clone(&ctx_for_finish), last_agent_message)
// Finish uniformly from the spawn site so all tasks share the same lifecycle.
sess.on_task_finished(Arc::clone(&ctx_for_finish), task_result)
.await;
}
done_clone.notify_waiters();
@@ -557,8 +563,16 @@ impl Session {
pub async fn on_task_finished(
self: &Arc<Self>,
turn_context: Arc<TurnContext>,
last_agent_message: Option<String>,
task_result: SessionTaskResult,
) {
let (last_agent_message, abort_reason) = match task_result {
Ok(last_agent_message) => (last_agent_message, None),
Err(CodexErr::TurnAborted) => (None, Some(TurnAbortReason::Interrupted)),
Err(err) => {
warn!(%err, "session task returned an unexpected error");
(None, None)
}
};
turn_context
.turn_metadata_state
.cancel_git_enrichment_task();
@@ -730,25 +744,36 @@ impl Session {
.turn_timing_state
.completed_at_and_duration_ms()
.await;
let time_to_first_token_ms = turn_context
.turn_timing_state
.time_to_first_token_ms()
.await;
self.services
.analytics_events_client
.track_turn_profile(TurnProfileFact {
turn_id: turn_context.sub_id.clone(),
profile: turn_context.turn_timing_state.complete_profile(),
});
self.emit_turn_stop_lifecycle(turn_context.extension_data.as_ref())
.await;
let event = EventMsg::TurnComplete(TurnCompleteEvent {
turn_id: turn_context.sub_id.clone(),
last_agent_message,
completed_at,
duration_ms,
time_to_first_token_ms,
});
let event = if let Some(reason) = abort_reason {
self.emit_turn_abort_lifecycle(reason.clone(), turn_context.extension_data.as_ref())
.await;
EventMsg::TurnAborted(TurnAbortedEvent {
turn_id: Some(turn_context.sub_id.clone()),
reason,
completed_at,
duration_ms,
})
} else {
let time_to_first_token_ms = turn_context
.turn_timing_state
.time_to_first_token_ms()
.await;
self.emit_turn_stop_lifecycle(turn_context.extension_data.as_ref())
.await;
EventMsg::TurnComplete(TurnCompleteEvent {
turn_id: turn_context.sub_id.clone(),
last_agent_message,
completed_at,
duration_ms,
time_to_first_token_ms,
})
};
self.send_event(turn_context.as_ref(), event).await;
self.services
.guardian_rejection_circuit_breaker