mirror of
https://github.com/openai/codex.git
synced 2026-10-31 19:28:49 +00:00
rka-oai
dac588f413
## 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.
91 lines
3.1 KiB
Rust
91 lines
3.1 KiB
Rust
use std::sync::Arc;
|
|
|
|
use tokio_util::sync::CancellationToken;
|
|
|
|
use crate::session::TurnInput;
|
|
use crate::session::turn::run_turn;
|
|
use crate::session::turn_context::TurnContext;
|
|
use crate::session_startup_prewarm::SessionStartupPrewarmResolution;
|
|
use crate::state::TaskKind;
|
|
use codex_protocol::protocol::EventMsg;
|
|
use codex_protocol::protocol::TurnStartedEvent;
|
|
use tracing::Instrument;
|
|
use tracing::trace_span;
|
|
|
|
use super::SessionTask;
|
|
use super::SessionTaskContext;
|
|
use super::SessionTaskResult;
|
|
|
|
#[derive(Default)]
|
|
pub(crate) struct RegularTask;
|
|
|
|
impl RegularTask {
|
|
pub(crate) fn new() -> Self {
|
|
Self
|
|
}
|
|
}
|
|
|
|
impl SessionTask for RegularTask {
|
|
fn kind(&self) -> TaskKind {
|
|
TaskKind::Regular
|
|
}
|
|
|
|
fn span_name(&self) -> &'static str {
|
|
"session_task.turn"
|
|
}
|
|
|
|
async fn run(
|
|
self: Arc<Self>,
|
|
session: Arc<SessionTaskContext>,
|
|
ctx: Arc<TurnContext>,
|
|
input: Vec<TurnInput>,
|
|
cancellation_token: CancellationToken,
|
|
) -> SessionTaskResult {
|
|
let sess = session.clone_session();
|
|
let turn_extension_data = session.turn_extension_data();
|
|
let run_turn_span = trace_span!("run_turn");
|
|
// Regular turns emit `TurnStarted` inline so first-turn lifecycle does
|
|
// not wait on startup prewarm resolution.
|
|
let prewarmed_client_session = async {
|
|
let event = EventMsg::TurnStarted(TurnStartedEvent {
|
|
turn_id: ctx.sub_id.clone(),
|
|
trace_id: ctx.trace_id.clone(),
|
|
started_at: ctx.turn_timing_state.started_at_unix_secs().await,
|
|
model_context_window: ctx.model_context_window(),
|
|
collaboration_mode_kind: ctx.collaboration_mode.mode,
|
|
});
|
|
sess.send_event(ctx.as_ref(), event).await;
|
|
sess.set_server_reasoning_included(/*included*/ false).await;
|
|
sess.consume_startup_prewarm_for_regular_turn(&cancellation_token)
|
|
.await
|
|
}
|
|
.instrument(trace_span!("regular_task.prepare_run_turn"))
|
|
.await;
|
|
let prewarmed_client_session = match prewarmed_client_session {
|
|
SessionStartupPrewarmResolution::Cancelled => return Ok(None),
|
|
SessionStartupPrewarmResolution::Unavailable { .. } => None,
|
|
SessionStartupPrewarmResolution::Ready(prewarmed_client_session) => {
|
|
Some(*prewarmed_client_session)
|
|
}
|
|
};
|
|
let mut next_input = input;
|
|
let mut prewarmed_client_session = prewarmed_client_session;
|
|
loop {
|
|
let last_agent_message = run_turn(
|
|
Arc::clone(&sess),
|
|
Arc::clone(&ctx),
|
|
Arc::clone(&turn_extension_data),
|
|
next_input,
|
|
prewarmed_client_session.take(),
|
|
cancellation_token.child_token(),
|
|
)
|
|
.instrument(run_turn_span.clone())
|
|
.await?;
|
|
if !sess.input_queue.has_pending_input(&sess.active_turn).await {
|
|
return Ok(last_agent_message);
|
|
}
|
|
next_input = Vec::new();
|
|
}
|
|
}
|
|
}
|