agent/codex
2
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity
Files
4a05d3b282a73dc7baa67f5baae0eaa511ff8d78
codex/codex-rs/core/tests/suite/token_budget.rs
T
pakrym-oai dac5f07403 [codex] Add context remaining tool (#27518) 
## Why

The token budget feature can inject remaining-context notices into
model-visible context, but the model does not have a direct way to ask
for that same remaining-token fragment on demand.

This PR adds a small model tool for the token budget feature so the
model can request the current remaining context window message without
duplicating the fragment format.

## What changed

- Adds a `get_context_remaining` direct-model tool behind
`Feature::TokenBudget`.
- Renders the tool output through `TokenBudgetRemainingContext`,
matching the existing budget message shape.
- Registers the tool alongside `new_context` in the token budget tool
set.
- Adds integration coverage that verifies the tool is exposed and
returns the same `<token_budget>` remaining fragment already present in
context.

## Validation

- `just test -p codex-core token_budget`
2026-06-11 04:17:10 +00:00

450 lines
15 KiB
Rust
Raw Blame History

use anyhow::Result;
use codex_features::Feature;
use codex_model_provider_info::built_in_model_providers;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::Op;
use core_test_support::PathBufExt;
use core_test_support::context_snapshot;
use core_test_support::context_snapshot::ContextSnapshotOptions;
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_sequence;
use core_test_support::responses::sse;
use core_test_support::responses::start_mock_server;
use core_test_support::skip_if_no_network;
use core_test_support::test_codex::local;
use core_test_support::test_codex::test_codex;
use core_test_support::wait_for_event;
use pretty_assertions::assert_eq;
use serde_json::Value;
use serde_json::json;
const CONFIGURED_CONTEXT_WINDOW: i64 = 128_000;
const EFFECTIVE_CONTEXT_WINDOW: i64 = CONFIGURED_CONTEXT_WINDOW * 95 / 100;
fn token_budget_texts(request: &ResponsesRequest) -> Vec<String> {
request
.message_input_texts("developer")
.into_iter()
.filter(|text| text.starts_with("<token_budget>"))
.collect()
}
fn tool_names(request: &ResponsesRequest) -> Vec<String> {
request
.body_json()
.get("tools")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(|tool| tool.get("name").and_then(Value::as_str).map(str::to_string))
.collect()
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn token_budget_context_is_only_emitted_with_full_context() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let responses = mount_sse_sequence(
&server,
vec![
sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]),
sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]),
],
)
.await;
let test = test_codex()
.with_config(|config| {
config.model_context_window = Some(CONFIGURED_CONTEXT_WINDOW);
config
.features
.enable(Feature::TokenBudget)
.expect("test config should allow token budget");
})
.build(&server)
.await?;
test.submit_turn("first turn").await?;
let second_cwd = test.workspace_path("second-cwd");
std::fs::create_dir_all(&second_cwd)?;
test.submit_turn_with_environments("second turn", Some(vec![local(second_cwd.abs())]))
.await?;
let requests = responses.requests();
assert_eq!(requests.len(), 2);
let expected = vec![format!(
"<token_budget>\nCurrent context window 0.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window.\n</token_budget>"
)];
assert_eq!(
token_budget_texts(&requests[0]),
expected,
"initial full context should report context window 0"
);
assert_eq!(
token_budget_texts(&requests[1]),
expected,
"steady-state context update should not advance the context window"
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn token_budget_remaining_context_emits_on_first_threshold_crossing() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let responses = mount_sse_sequence(
&server,
vec![
sse(vec![
ev_response_created("resp-1"),
ev_completed_with_tokens("resp-1", /*total_tokens*/ 2_500),
]),
sse(vec![
ev_response_created("resp-2"),
ev_completed_with_tokens("resp-2", /*total_tokens*/ 3_000),
]),
sse(vec![
ev_response_created("resp-3"),
ev_completed_with_tokens("resp-3", /*total_tokens*/ 5_000),
]),
sse(vec![
ev_response_created("resp-4"),
ev_completed_with_tokens("resp-4", /*total_tokens*/ 8_000),
]),
sse(vec![ev_response_created("resp-5"), ev_completed("resp-5")]),
],
)
.await;
let test = test_codex()
.with_config(|config| {
config.model_context_window = Some(10_000);
config
.features
.enable(Feature::TokenBudget)
.expect("test config should allow token budget");
})
.build(&server)
.await?;
for turn in 1..=5 {
test.submit_turn(&format!("turn {turn}")).await?;
}
let requests = responses.requests();
assert_eq!(requests.len(), 5);
let full_context = "<token_budget>\nCurrent context window 0.\nYou have 9500 tokens left in this context window.\n</token_budget>"
.to_string();
let threshold_25 =
"<token_budget>\nYou have 7000 tokens left in this context window.\n</token_budget>"
.to_string();
let threshold_50 =
"<token_budget>\nYou have 4500 tokens left in this context window.\n</token_budget>"
.to_string();
let threshold_75 =
"<token_budget>\nYou have 1500 tokens left in this context window.\n</token_budget>"
.to_string();
assert_eq!(token_budget_texts(&requests[0]), vec![full_context.clone()]);
assert_eq!(
token_budget_texts(&requests[1]),
vec![full_context.clone(), threshold_25.clone()]
);
assert_eq!(
token_budget_texts(&requests[2]),
vec![full_context.clone(), threshold_25.clone()]
);
assert_eq!(
token_budget_texts(&requests[3]),
vec![
full_context.clone(),
threshold_25.clone(),
threshold_50.clone()
]
);
assert_eq!(
token_budget_texts(&requests[4]),
vec![full_context, threshold_25, threshold_50, threshold_75]
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn get_context_remaining_returns_token_budget_remaining_fragment() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let call_id = "remaining-call";
let responses = mount_sse_sequence(
&server,
vec![
sse(vec![
ev_response_created("resp-1"),
ev_assistant_message("msg-1", "noted"),
ev_completed_with_tokens("resp-1", /*total_tokens*/ 2_500),
]),
sse(vec![
ev_response_created("resp-2"),
ev_function_call(call_id, "get_context_remaining", "{}"),
ev_completed_with_tokens("resp-2", /*total_tokens*/ 2_500),
]),
sse(vec![
ev_response_created("resp-3"),
ev_assistant_message("msg-3", "done"),
ev_completed("resp-3"),
]),
],
)
.await;
let test = test_codex()
.with_config(|config| {
config.model_context_window = Some(10_000);
config
.features
.enable(Feature::TokenBudget)
.expect("test config should allow token budget");
})
.build(&server)
.await?;
test.submit_turn("spend some tokens").await?;
test.submit_turn("check remaining context").await?;
let requests = responses.requests();
assert_eq!(requests.len(), 3);
assert!(
tool_names(&requests[1])
.iter()
.any(|name| name == "get_context_remaining"),
"get_context_remaining should be exposed when token budget is enabled"
);
let full_context = "<token_budget>\nCurrent context window 0.\nYou have 9500 tokens left in this context window.\n</token_budget>"
.to_string();
let remaining_context =
"<token_budget>\nYou have 7000 tokens left in this context window.\n</token_budget>"
.to_string();
assert_eq!(
token_budget_texts(&requests[1]),
vec![full_context, remaining_context.clone()]
);
assert_eq!(
requests[2].function_call_output_content_and_success(call_id),
Some((Some(remaining_context), None))
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn get_context_remaining_returns_unknown_when_window_is_unavailable() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let call_id = "remaining-call";
let responses = mount_sse_sequence(
&server,
vec![
sse(vec![
ev_response_created("resp-1"),
ev_function_call(call_id, "get_context_remaining", "{}"),
ev_completed("resp-1"),
]),
sse(vec![
ev_response_created("resp-2"),
ev_assistant_message("msg-2", "done"),
ev_completed("resp-2"),
]),
],
)
.await;
let test = test_codex()
.with_model_info_override("gpt-5.2", |model_info| {
model_info.context_window = None;
model_info.max_context_window = None;
})
.with_config(|config| {
config.model_context_window = None;
config
.features
.enable(Feature::TokenBudget)
.expect("test config should allow token budget");
})
.build(&server)
.await?;
test.submit_turn("check remaining context").await?;
let requests = responses.requests();
assert_eq!(requests.len(), 2);
assert!(
tool_names(&requests[0])
.iter()
.any(|name| name == "get_context_remaining"),
"get_context_remaining should be exposed when token budget is enabled"
);
assert_eq!(token_budget_texts(&requests[0]), Vec::<String>::new());
assert_eq!(
requests[1].function_call_output_content_and_success(call_id),
Some((
Some(
"<token_budget>\nYou have unknown tokens left in this context window.\n</token_budget>"
.to_string()
),
None,
))
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn token_budget_context_uses_new_window_after_compaction() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let responses = mount_sse_sequence(
&server,
vec![
sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]),
sse(vec![
ev_response_created("resp-compact"),
ev_assistant_message("msg-compact", "compact summary"),
ev_completed("resp-compact"),
]),
sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]),
],
)
.await;
let mut model_provider = built_in_model_providers(/*openai_base_url*/ None)["openai"].clone();
model_provider.name = "OpenAI-compatible test provider".to_string();
model_provider.base_url = Some(format!("{}/v1", server.uri()));
model_provider.supports_websockets = false;
let test = test_codex()
.with_config(move |config| {
config.model_provider = model_provider;
config.model_context_window = Some(CONFIGURED_CONTEXT_WINDOW);
config
.features
.enable(Feature::TokenBudget)
.expect("test config should allow token budget");
})
.build(&server)
.await?;
test.submit_turn("before compact").await?;
test.codex.submit(Op::Compact).await?;
wait_for_event(&test.codex, |event| {
matches!(event, EventMsg::TurnComplete(_))
})
.await;
test.submit_turn("after compact").await?;
let requests = responses.requests();
assert_eq!(requests.len(), 3);
assert_eq!(
token_budget_texts(&requests[2]),
vec![format!(
"<token_budget>\nCurrent context window 1.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window.\n</token_budget>"
)],
"post-compaction full context should report context window 1"
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn new_context_tool_starts_new_window_before_follow_up() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let call_id = "new-window-call";
let continue_call_id = "continue-call";
let continue_args = json!({
"plan": [
{"step": "Continue in the new context window", "status": "in_progress"}
],
})
.to_string();
let responses = mount_sse_sequence(
&server,
vec![
sse(vec![
ev_response_created("resp-1"),
ev_function_call(call_id, "new_context", "{}"),
ev_completed("resp-1"),
]),
sse(vec![
ev_response_created("resp-2"),
ev_function_call(continue_call_id, "update_plan", &continue_args),
ev_completed("resp-2"),
]),
sse(vec![
ev_response_created("resp-3"),
ev_assistant_message("msg-3", "done"),
ev_completed("resp-3"),
]),
],
)
.await;
let test = test_codex()
.with_config(|config| {
config.model_context_window = Some(CONFIGURED_CONTEXT_WINDOW);
config
.features
.enable(Feature::TokenBudget)
.expect("test config should allow token budget");
})
.build(&server)
.await?;
test.submit_turn("request new context window").await?;
let requests = responses.requests();
assert_eq!(requests.len(), 3);
assert!(
tool_names(&requests[0])
.iter()
.any(|name| name == "new_context"),
"new_context should be exposed when token budget is enabled"
);
assert_eq!(
token_budget_texts(&requests[2]),
vec![format!(
"<token_budget>\nCurrent context window 1.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window.\n</token_budget>"
)]
);
assert!(
!requests[2].body_contains_text("request new context window"),
"new_context should drop the prior window history before continuing the turn"
);
assert_eq!(
requests[2].function_call_output_text(continue_call_id),
Some("Plan updated".to_string())
);
insta::assert_snapshot!(
"token_budget_new_context_window_tool_full_context",
context_snapshot::format_labeled_requests_snapshot(
"New context window tool installs fresh full context before the next follow-up request.",
&[("Final Follow-Up Request", &requests[2])],
&ContextSnapshotOptions::default(),
)
);
Ok(())
}
Reference in New Issue View Git Blame Copy Permalink