diff --git a/codex-rs/core/src/rollout_budget.rs b/codex-rs/core/src/rollout_budget.rs index 570b1b7c8f..4170935844 100644 --- a/codex-rs/core/src/rollout_budget.rs +++ b/codex-rs/core/src/rollout_budget.rs @@ -49,13 +49,13 @@ impl RolloutBudget { let Some(mut state) = self.lock() else { return false; }; - let charge_prefill = state.threads_with_usage.insert(thread_id); + let skip_prefill = state.threads_with_usage.insert(thread_id); state.weighted_tokens_used += usage.output_tokens.max(0) as f64 * state.config.sampling_token_weight - + if charge_prefill { - usage.non_cached_input() as f64 * state.config.prefill_token_weight - } else { + + if skip_prefill { 0.0 + } else { + usage.non_cached_input() as f64 * state.config.prefill_token_weight }; state.weighted_tokens_used >= state.config.limit_tokens as f64 } diff --git a/codex-rs/core/tests/suite/rollout_budget.rs b/codex-rs/core/tests/suite/rollout_budget.rs index 44abdf9456..58c66af3e3 100644 --- a/codex-rs/core/tests/suite/rollout_budget.rs +++ b/codex-rs/core/tests/suite/rollout_budget.rs @@ -9,7 +9,6 @@ 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; -use core_test_support::responses::ev_completed_with_tokens; use core_test_support::responses::ev_function_call; use core_test_support::responses::ev_response_created; use core_test_support::responses::mount_sse_once_match; @@ -48,6 +47,22 @@ fn rollout_budget_message(remaining_tokens: i64) -> String { ) } +fn ev_completed_with_output_tokens(id: &str, output_tokens: i64) -> serde_json::Value { + json!({ + "type": "response.completed", + "response": { + "id": id, + "usage": { + "input_tokens": 0, + "input_tokens_details": null, + "output_tokens": output_tokens, + "output_tokens_details": null, + "total_tokens": output_tokens + } + } + }) +} + fn wire_request_contains(request: &wiremock::Request, text: &str) -> bool { std::str::from_utf8(&request.body).is_ok_and(|body| body.contains(text)) } @@ -156,7 +171,7 @@ async fn subagent_usage_draws_from_the_shared_budget() -> Result<()> { sse(vec![ ev_response_created("root-1"), ev_function_call(SPAWN_CALL_ID, "spawn_agent", &spawn_args), - ev_completed_with_tokens("root-1", /*total_tokens*/ 10), + ev_completed_with_output_tokens("root-1", /*output_tokens*/ 10), ]), ) .await; @@ -165,7 +180,7 @@ async fn subagent_usage_draws_from_the_shared_budget() -> Result<()> { |request: &wiremock::Request| wire_request_contains(request, "\"type\":\"agent_message\""), sse(vec![ ev_response_created("child-1"), - ev_completed_with_tokens("child-1", /*total_tokens*/ 30), + ev_completed_with_output_tokens("child-1", /*output_tokens*/ 30), ]), ) .await; @@ -177,7 +192,7 @@ async fn subagent_usage_draws_from_the_shared_budget() -> Result<()> { }, sse(vec![ ev_response_created("root-2"), - ev_completed_with_tokens("root-2", /*total_tokens*/ 10), + ev_completed_with_output_tokens("root-2", /*output_tokens*/ 10), ]), ) .await; @@ -232,11 +247,11 @@ async fn exhausted_budget_aborts_current_and_later_turns() -> Result<()> { vec![ sse(vec![ ev_response_created("exhaust-budget"), - ev_completed_with_tokens("exhaust-budget", /*total_tokens*/ 30), + ev_completed_with_output_tokens("exhaust-budget", /*output_tokens*/ 30), ]), sse(vec![ ev_response_created("already-exhausted"), - ev_completed_with_tokens("already-exhausted", /*total_tokens*/ 1), + ev_completed_with_output_tokens("already-exhausted", /*output_tokens*/ 1), ]), ], ) @@ -299,13 +314,13 @@ async fn compaction_budget_exhaustion_aborts_without_error_or_retry(remote_v2: b "encrypted_content": "encrypted-summary", } }), - ev_completed_with_tokens("compact", /*total_tokens*/ 10), + ev_completed_with_output_tokens("compact", /*output_tokens*/ 10), ]) } else { sse(vec![ ev_response_created("compact"), ev_assistant_message("compact-summary", "compact summary"), - ev_completed_with_tokens("compact", /*total_tokens*/ 10), + ev_completed_with_output_tokens("compact", /*output_tokens*/ 10), ]) }; let responses = mount_sse_sequence(&server, vec![compact_response]).await; @@ -362,12 +377,12 @@ async fn restates_the_current_remainder_after_compaction() -> Result<()> { vec![ sse(vec![ ev_response_created("resp-1"), - ev_completed_with_tokens("resp-1", /*total_tokens*/ 20), + ev_completed_with_output_tokens("resp-1", /*output_tokens*/ 20), ]), sse(vec![ ev_response_created("resp-compact"), ev_assistant_message("msg-compact", "compact summary"), - ev_completed_with_tokens("resp-compact", /*total_tokens*/ 10), + ev_completed_with_output_tokens("resp-compact", /*output_tokens*/ 10), ]), sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), ], @@ -427,7 +442,7 @@ async fn restates_the_current_remainder_after_rollback() -> Result<()> { vec![ sse(vec![ ev_response_created("resp-1"), - ev_completed_with_tokens("resp-1", /*total_tokens*/ 30), + ev_completed_with_output_tokens("resp-1", /*output_tokens*/ 30), ]), sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), ],