github/codex
4
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-19 16:17:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

Account for token usage from incomplete responses (#53670)

Browse Source 
## Why

Terminal `response.incomplete` errors discard reported token usage. Counting that usage as a new context baseline can also double-count input and trigger unnecessary compaction on the next turn.

## What changed

- Carry the response ID and valid reported usage with incomplete-response errors.
- Record usage for sampling and foreground compaction in cumulative totals, rollout records, budgets, and token-usage contributors without replacing the context baseline. Preserve the original terminal error, including when accounting exhausts the budget or cancellation races the error.
- Charge incomplete background compaction requests to the rollout budget without publishing foreground token usage.
- Ignore missing or malformed usage rather than estimating it.

## Testing

Extend coverage for SSE and WebSocket incomplete responses, rollout usage records, compaction accounting, deferred token-count events, and token-usage contributors. Add regressions for unnecessary compaction on the next turn and background compaction publishing foreground usage.

GitOrigin-RevId: bf1ee6df3912052ab7c19e1c6af8ad3dc4b07675
This commit is contained in:
jif
2026-10-13 13:36:24 +00:00
committed by copyberry
parent a68a7a6ab0
commit 816fa702ec
12 changed files with 383 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/codex-api/src/error.rs
+3 -2
View File
@@ -1,6 +1,7 @@
use crate::rate_limits::RateLimitError;
use codex_client::TransportError;
use codex_http_client::RetryAfter;
use codex_protocol::error::ResponseIncompleteError;
use codex_protocol::protocol::MisalignmentErrorDetails;
use http::StatusCode;
use serde_json::Value;
@@ -16,8 +17,8 @@ pub enum ApiError {
Stream(String),
#[error("stream error: Incomplete response returned, reason: content_filter")]
ContentFilter,
#[error("Incomplete response returned, reason: {0}")]
ResponseIncomplete(String),
#[error(transparent)]
ResponseIncomplete(Box<ResponseIncompleteError>),
#[error("context window exceeded")]
ContextWindowExceeded,
#[error("quota exceeded")]
codex-rs/codex-api/src/sse/responses.rs
+17 -1
View File
@@ -11,6 +11,7 @@ use crate::telemetry::SseTelemetry;
use codex_client::ByteStream;
use codex_client::StreamResponse;
use codex_protocol::ResponseUsageMetadata;
use codex_protocol::error::ResponseIncompleteError;
use codex_protocol::models::ResponseItem;
use codex_protocol::protocol::ModelVerification;
use codex_protocol::protocol::TokenUsage;
@@ -429,8 +430,23 @@ pub fn process_responses_event(
return Err(ResponsesEventError::Api(ApiError::ContentFilter));
}
if reason != "interrupted" {
let response = event.response.as_ref();
let response_id = response
.and_then(|response| response.get("id"))
.and_then(Value::as_str)
.map(str::to_owned);
let token_usage = response
.and_then(|response| response.get("usage"))
.and_then(|usage| {
serde_json::from_value::<ResponseCompletedUsage>(usage.clone()).ok()
})
.map(Into::into);
return Err(ResponsesEventError::Api(ApiError::ResponseIncomplete(
reason.to_string(),
Box::new(ResponseIncompleteError {
reason: reason.to_string(),
response_id,
token_usage,
}),
)));
}
}
codex-rs/core/src/compact.rs
+12 -1
View File
@@ -866,7 +866,18 @@ async fn drain_to_completed(
});
}
Ok(_) => continue,
Err(e) => return Err(e),
Err(e) => {
if let CodexErrorDetails::ResponseIncomplete(incomplete) = e.details() {
sess.record_incomplete_response_usage(
turn_context,
&turn_context.initial_settings,
incomplete,
)
.await;
sess.send_token_count_event(turn_context).await;
}
return Err(e);
}
}
}
}
codex-rs/core/src/compact_remote_v2.rs
+56
View File
@@ -394,6 +394,27 @@ async fn collect_compaction_output(
},
};
let Some(event) = event else { break };
if let Err(err) = &event
&& let CodexErrorDetails::ResponseIncomplete(incomplete) = err.details()
{
match request_kind {
CompactionRequestKind::Background(_) => {
if let Some(usage) = incomplete.token_usage.as_ref() {
// Match completed background requests without replacing the terminal error.
let _ = sess.record_rollout_budget_usage(usage).await;
}
}
CompactionRequestKind::Foreground => {
sess.record_incomplete_response_usage(
turn_context,
&turn_context.initial_settings,
incomplete,
)
.await;
sess.send_token_count_event(turn_context).await;
}
}
}
match event? {
ResponseEvent::OutputItemDone(item) => {
output_item_count += 1;
@@ -1219,6 +1240,41 @@ mod tests {
assert_eq!(truncated, vec![newest]);
}
#[tokio::test]
async fn collect_background_incomplete_does_not_publish_foreground_usage() {
let stream = response_stream(vec![Err(CodexErr::new(
CodexErrorDetails::ResponseIncomplete(Box::new(
codex_protocol::error::ResponseIncompleteError {
reason: "max_output_tokens".to_string(),
response_id: Some("resp-background".to_string()),
token_usage: Some(TokenUsage {
input_tokens: 100,
output_tokens: 20,
total_tokens: 120,
..Default::default()
}),
},
)),
))]);
let (sess, turn_context, rx) =
crate::session::tests::make_session_and_context_with_rx().await;
let cancellation = CancellationToken::new();
cancellation.cancel();
let result = collect_compaction_output(
&sess,
&turn_context,
stream,
CompactionRequestKind::Background(&cancellation),
)
.await;
assert!(matches!(result, Err(ref err)
if matches!(err.details(), CodexErrorDetails::ResponseIncomplete(_))));
assert_eq!(sess.token_usage_info().await, None);
assert!(rx.is_empty());
}
#[tokio::test]
async fn collect_compaction_output_accepts_additional_output_items() {
let compaction = ResponseItem::Compaction {
codex-rs/core/src/session/mod.rs
+81 -9
View File
@@ -233,6 +233,7 @@ use codex_model_provider::create_model_provider;
use codex_model_provider_info::ModelProviderInfo;
use codex_protocol::error::CodexErr;
use codex_protocol::error::CodexErrorDetails;
use codex_protocol::error::ResponseIncompleteError;
use codex_protocol::error::Result as CodexResult;
#[cfg(test)]
use codex_protocol::exec_output::StreamOutput;
@@ -408,6 +409,13 @@ use codex_utils_absolute_path::AbsolutePathBuf;
#[cfg(test)]
use codex_utils_stream_parser::ProposedPlanSegment;
/// Whether reported usage updates the installed context baseline or only spending.
#[derive(Clone, Copy)]
pub(crate) enum TokenUsageContext {
Sampling,
SpendingOnly,
}
/// Queue and lifecycle endpoints for a running [`Session`].
///
/// Runtime state lives on `Session`; keeping these endpoints separate lets all
@@ -4938,13 +4946,44 @@ impl Session {
Ok(Arc::new(world_state))
}
/// Record reported spending without changing the context baseline or estimating usage.
pub(crate) async fn record_incomplete_response_usage(
&self,
turn_context: &TurnContext,
settings: &ResolvedStepSettings,
incomplete: &ResponseIncompleteError,
) {
if let Some(response_id) = &incomplete.response_id {
self.record_response_token_usage(
turn_context,
response_id,
incomplete.token_usage.as_ref(),
)
.await;
}
// Accounting may exhaust the budget; keep the original incomplete-response reason.
let _ = self
.record_token_usage_info(
turn_context,
settings,
incomplete.token_usage.as_ref(),
TokenUsageContext::SpendingOnly,
)
.await;
}
pub(crate) async fn update_token_usage_info(
&self,
turn_context: &TurnContext,
token_usage: Option<&TokenUsage>,
) -> CodexResult<()> {
let result = self
.record_token_usage_info(turn_context, &turn_context.initial_settings, token_usage)
.record_token_usage_info(
turn_context,
&turn_context.initial_settings,
token_usage,
TokenUsageContext::Sampling,
)
.await;
self.send_token_count_event(turn_context).await;
result
@@ -4966,6 +5005,16 @@ impl Session {
}),
)
.await;
self.record_response_token_usage(turn_context, response_id, usage)
.await;
}
pub(crate) async fn record_response_token_usage(
&self,
turn_context: &TurnContext,
response_id: &str,
usage: Option<&TokenUsage>,
) {
let Some(usage) = usage else {
return;
};
@@ -4989,19 +5038,42 @@ impl Session {
turn_context: &TurnContext,
settings: &ResolvedStepSettings,
token_usage: Option<&TokenUsage>,
context: TokenUsageContext,
) -> CodexResult<()> {
if let Some(token_usage) = token_usage {
let token_info = {
let mut state = self.state.lock().await;
state
.update_token_info_from_usage(token_usage, turn_context.model_context_window());
if matches!(
turn_context.config.model_auto_compact_token_limit_scope,
AutoCompactTokenLimitScope::BodyAfterPrefix
) {
state.ensure_auto_compact_window_server_prefill_from_usage(token_usage);
match context {
TokenUsageContext::Sampling => {
state.update_token_info_from_usage(
token_usage,
turn_context.model_context_window(),
);
if matches!(
turn_context.config.model_auto_compact_token_limit_scope,
AutoCompactTokenLimitScope::BodyAfterPrefix
) {
state.ensure_auto_compact_window_server_prefill_from_usage(token_usage);
}
state.token_info()
}
TokenUsageContext::SpendingOnly => {
let mut info = state.token_info().unwrap_or(TokenUsageInfo {
total_token_usage: TokenUsage::default(),
last_token_usage: TokenUsage::default(),
model_context_window: turn_context.model_context_window(),
});
info.total_token_usage.add_assign(token_usage);
state.set_token_info(Some(info.clone()));
// Contributors receive the reported usage; stored context stays unchanged.
info.last_token_usage = token_usage.clone();
info.model_context_window = turn_context
.model_context_window()
.or(info.model_context_window);
Some(info)
}
}
state.token_info()
};
let turn_state = self
.input_queue
codex-rs/core/src/session/tests.rs
+19 -1
View File
@@ -3525,8 +3525,12 @@ async fn recompute_token_usage_updates_model_context_window() {
assert_eq!(actual.model_context_window, Some(128_000));
}
#[test_case(crate::session::TokenUsageContext::Sampling; "sampling")]
#[test_case(crate::session::TokenUsageContext::SpendingOnly; "spending_only")]
#[tokio::test]
async fn record_token_usage_info_notifies_extension_contributors() {
async fn record_token_usage_info_notifies_extension_contributors(
context: crate::session::TokenUsageContext,
) {
struct SessionTokenUsageMarker;
struct ThreadTokenUsageMarker;
@@ -3608,6 +3612,7 @@ async fn record_token_usage_info_notifies_extension_contributors() {
&turn_context,
&turn_context.initial_settings,
Some(&first_usage),
crate::session::TokenUsageContext::Sampling,
)
.await
.expect("first usage should be recorded");
@@ -3616,12 +3621,21 @@ async fn record_token_usage_info_notifies_extension_contributors() {
&turn_context,
&turn_context.initial_settings,
Some(&second_usage),
context,
)
.await
.expect("second usage should be recorded");
let mut expected_total_usage = first_usage.clone();
expected_total_usage.add_assign(&second_usage);
let expected_session_info = TokenUsageInfo {
total_token_usage: expected_total_usage.clone(),
last_token_usage: match context {
crate::session::TokenUsageContext::Sampling => second_usage.clone(),
crate::session::TokenUsageContext::SpendingOnly => first_usage.clone(),
},
model_context_window: turn_context.model_context_window(),
};
let expected = vec![
RecordedTokenUsage {
session_level_id: session.session_id().to_string(),
@@ -3654,6 +3668,10 @@ async fn record_token_usage_info_notifies_extension_contributors() {
.drain(..)
.collect::<Vec<_>>();
assert_eq!(expected, actual);
assert_eq!(
session.token_usage_info().await,
Some(expected_session_info)
);
}
#[tokio::test]
codex-rs/core/src/session/turn.rs
+22 -4
View File
@@ -2669,9 +2669,8 @@ async fn try_run_sampling_request(
.or_cancel(&preempt)
.or_cancel(&cancellation_token)
.await;
if cancellation_token.is_cancelled() {
break Err(CodexErr::TurnAborted);
}
// Account for an already-received terminal error even if cancellation raced it.
// Cancellation still stops ordinary stream events before they can run tools.
let event = match event {
Ok(Ok(event)) => event,
Ok(Err(_)) => {
@@ -2694,9 +2693,27 @@ async fn try_run_sampling_request(
Err(_) => break Err(CodexErr::TurnAborted),
};
if cancellation_token.is_cancelled()
&& !matches!(event.as_ref(), Some(Err(err))
if matches!(err.details(), CodexErrorDetails::ResponseIncomplete(_)))
{
break Err(CodexErr::TurnAborted);
}
let event = match event {
Some(Ok(event)) => event,
Some(Err(err)) => break Err(err),
Some(Err(err)) => {
if let CodexErrorDetails::ResponseIncomplete(incomplete) = err.details() {
sess.record_incomplete_response_usage(
&turn_context,
&step_context.settings,
incomplete,
)
.await;
should_emit_token_count |= incomplete.token_usage.is_some();
}
break Err(err);
}
None => {
break Err(CodexErr::Stream(
"stream closed before response.completed".into(),
@@ -3020,6 +3037,7 @@ async fn try_run_sampling_request(
&turn_context,
&step_context.settings,
token_usage.as_ref(),
super::TokenUsageContext::Sampling,
)
.await;
should_emit_token_count = true;
codex-rs/core/src/tasks/mod_tests.rs
+1
View File
@@ -225,6 +225,7 @@ async fn turn_completion_metrics_follow_model_switch(scenario: UsageScenario) {
&turn_context,
&response_context.initial_settings,
Some(&response_usage),
crate::session::TokenUsageContext::Sampling,
)
.await
.expect("response usage should be recorded");
codex-rs/core/tests/suite/client.rs
+130 -5
View File
@@ -3799,10 +3799,17 @@ async fn incomplete_response_emits_content_filter_error_message() -> anyhow::Res
Ok(())
}
#[test_case::test_case(false; "sse")]
#[test_case::test_case(true; "websocket")]
#[test_case::test_case(false, json!({"input_tokens": 120, "output_tokens": 30, "total_tokens": 150}); "sse_usage")]
#[test_case::test_case(true, json!({"input_tokens": 120, "output_tokens": 30, "total_tokens": 150}); "websocket_usage")]
#[test_case::test_case(false, serde_json::Value::Null; "sse_missing_usage")]
#[test_case::test_case(true, serde_json::Value::Null; "websocket_missing_usage")]
#[test_case::test_case(false, json!({"malformed": true}); "sse_malformed_usage")]
#[test_case::test_case(true, json!({"malformed": true}); "websocket_malformed_usage")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn incomplete_response_is_terminal(use_websocket: bool) -> anyhow::Result<()> {
async fn incomplete_response_is_terminal(
use_websocket: bool,
reported_usage: serde_json::Value,
) -> anyhow::Result<()> {
skip_if_no_network!(Ok(()));
let incomplete = vec![
ev_response_created("resp-incomplete"),
@@ -3811,7 +3818,8 @@ async fn incomplete_response_is_terminal(use_websocket: bool) -> anyhow::Result<
"type": "response.incomplete",
"response": {
"id": "resp-incomplete",
"incomplete_details": {"reason": "max_output_tokens"}
"incomplete_details": {"reason": "max_output_tokens"},
"usage": reported_usage
}
}),
];
@@ -3859,7 +3867,44 @@ async fn incomplete_response_is_terminal(use_websocket: bool) -> anyhow::Result<
matches!(event, EventMsg::TurnComplete(_))
})
.await;
insta::assert_snapshot!(errors.join("\n"), @"Incomplete response returned, reason: max_output_tokens");
insta::allow_duplicates! {
insta::assert_snapshot!(errors.join("\n"), @"Incomplete response returned, reason: max_output_tokens");
}
let expected_usage = if reported_usage.get("input_tokens").is_some() {
codex_protocol::protocol::TokenUsage {
input_tokens: 120,
output_tokens: 30,
total_tokens: 150,
..Default::default()
}
} else {
codex_protocol::protocol::TokenUsage::default()
};
assert_eq!(
test.codex
.token_usage_info()
.await
.map(|info| info.total_token_usage)
.unwrap_or_default(),
expected_usage
);
test.codex.flush_rollout().await?;
let records = std::fs::read_to_string(test.codex.rollout_path().expect("rollout path"))?
.lines()
.filter_map(|line| codex_rollout::parse_rollout_line(line).ok())
.filter_map(|line| match line.item {
RolloutItem::TokenUsageRecord(record) => Some(record.usage),
_ => None,
})
.collect::<Vec<_>>();
assert_eq!(
records,
if reported_usage.get("input_tokens").is_some() {
vec![expected_usage]
} else {
vec![]
}
);
if let Some(server) = websocket_server {
assert_eq!(
server
@@ -3876,6 +3921,86 @@ async fn incomplete_response_is_terminal(use_websocket: bool) -> anyhow::Result<
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn incomplete_response_without_output_does_not_compact_the_next_turn() -> anyhow::Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let test = test_codex()
.with_config(|config| {
config.model_auto_compact_token_limit = Some(128_000);
config.model_post_turn_compact_threshold_percent = 0;
})
.build_with_auto_env(&server)
.await?;
let initial = mount_sse_once(
&server,
sse(vec![
ev_assistant_message("initial", "Initial reply."),
ev_completed_with_tokens("initial", /*total_tokens*/ 40_000),
]),
)
.await;
test.submit_turn("hello").await?;
initial.single_request();
// An unfinished parallel call can report usage without emitting an output item.
// This input fits the window, but counting it again on top of reported usage does not.
let incomplete = mount_sse_once(
&server,
sse(vec![json!({
"type": "response.incomplete",
"response": {
"id": "incomplete",
"incomplete_details": {"reason": "max_output_tokens"},
"usage": {"input_tokens": 100_000, "output_tokens": 16, "total_tokens": 100_016}
}
})]),
)
.await;
test.submit_turn(&"x".repeat(240_000)).await?;
incomplete.single_request();
assert_eq!(
test.codex
.token_usage_info()
.await
.unwrap()
.total_token_usage
.total_tokens,
140_016,
);
let continuation = mount_sse_once(
&server,
sse(vec![
ev_assistant_message("continued", "Continued."),
ev_completed("continued"),
]),
)
.await;
test.codex
.start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
text: "continue".into(),
text_elements: Vec::new(),
}]))
.await?;
wait_for_event(&test.codex, |event| {
assert!(
!matches!(event, EventMsg::Error(_) | EventMsg::StreamError(_)),
"continuation failed: {event:?}"
);
assert!(
!matches!(event, EventMsg::ItemStarted(event)
if matches!(&event.item, codex_protocol::items::TurnItem::ContextCompaction(_))),
"the incomplete request's input must only count once toward compaction"
);
matches!(event, EventMsg::TurnComplete(_))
})
.await;
continuation.single_request();
Ok(())
}
/// We try to avoid setting env vars in tests because std::env::set_var() is
/// process-wide and unsafe. Though for this test, we want to simulate the
/// presence of an environment variable that the provider will read for auth, so
codex-rs/core/tests/suite/compact.rs
+20 -2
View File
@@ -1151,8 +1151,12 @@ async fn reasoning_effort_override_remote_v2_compaction_resets_pinned_effort(
Ok(())
}
#[test_case::test_case(json!({"input_tokens": 120_000, "output_tokens": 30_000, "total_tokens": 150_000}); "reported_usage")]
#[test_case::test_case(Value::Null; "missing_usage")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn incomplete_response_during_compaction_is_terminal() -> Result<()> {
async fn incomplete_response_during_compaction_is_terminal_and_counted(
reported_usage: Value,
) -> Result<()> {
skip_if_no_network!(Ok(()));
for (provider, expected_error) in [
(
@@ -1178,7 +1182,8 @@ async fn incomplete_response_during_compaction_is_terminal() -> Result<()> {
"type": "response.incomplete",
"response": {
"id": "incomplete-compaction",
"incomplete_details": {"reason": "max_output_tokens"}
"incomplete_details": {"reason": "max_output_tokens"},
"usage": reported_usage
}
})]),
],
@@ -1243,6 +1248,19 @@ async fn incomplete_response_during_compaction_is_terminal() -> Result<()> {
})
.await;
assert_eq!(errors, vec![expected_error]);
assert_eq!(
test.codex
.token_usage_info()
.await
.expect("partial usage")
.total_token_usage
.total_tokens,
if reported_usage.is_null() {
120_000
} else {
270_000
},
);
assert_eq!(
response_mock
.requests()
codex-rs/core/tests/suite/request_user_input.rs
+9 -2
View File
@@ -235,8 +235,12 @@ fn ev_rate_limits() -> Value {
})
}
#[test_case::test_case("response.completed"; "completed")]
#[test_case::test_case("response.incomplete"; "incomplete")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn request_user_input_interrupt_emits_deferred_token_count() -> anyhow::Result<()> {
async fn request_user_input_interrupt_emits_deferred_token_count(
terminal_type: &str,
) -> anyhow::Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
@@ -264,10 +268,13 @@ async fn request_user_input_interrupt_emits_deferred_token_count() -> anyhow::Re
})
.to_string();
let mut terminal = ev_completed_with_tokens("resp-interrupt", /*total_tokens*/ 77);
terminal["type"] = json!(terminal_type);
terminal["response"]["incomplete_details"] = json!({"reason": "max_output_tokens"});
let response = sse(vec![
ev_response_created("resp-interrupt"),
ev_function_call(call_id, "request_user_input", &request_args),
ev_completed_with_tokens("resp-interrupt", /*total_tokens*/ 77),
terminal,
]);
responses::mount_sse_once(&server, response).await;
codex-rs/protocol/src/error.rs
+13 -2
View File
@@ -10,6 +10,7 @@ use crate::protocol::ErrorEvent;
use crate::protocol::MisalignmentErrorDetails;
use crate::protocol::RateLimitReachedType;
use crate::protocol::RateLimitSnapshot;
use crate::protocol::TokenUsage;
use crate::protocol::TruncationPolicy;
use chrono::DateTime;
use chrono::Datelike;
@@ -80,6 +81,16 @@ pub struct CodexErr {
agent_context: Option<AgentErrorContext>,
}
/// A terminal response that stopped before the model finished, with any reported usage.
/// Content-filter retries and explicit interruptions have separate handling.
#[derive(Debug, Error, PartialEq, Eq)]
#[error("Incomplete response returned, reason: {reason}")]
pub struct ResponseIncompleteError {
pub reason: String,
pub response_id: Option<String>,
pub token_usage: Option<TokenUsage>,
}
/// The semantic category and diagnostic payload for a [`CodexErr`].
#[derive(Error, Debug, EnumDiscriminants)]
#[strum_discriminants(name(CodexErrKind))]
@@ -105,8 +116,8 @@ pub enum CodexErrorDetails {
"stream disconnected before completion: Incomplete response returned, reason: content_filter"
)]
ContentFilter,
#[error("Incomplete response returned, reason: {0}")]
ResponseIncomplete(String),
#[error(transparent)]
ResponseIncomplete(Box<ResponseIncompleteError>),
/// A retryable upstream rate limit received inside the response stream.
#[error("rate limit exceeded: {0}")]
RateLimitExceeded(String),