openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-04 15:08:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: scope context remaining to body window

Browse Source
This commit is contained in:
Michael Bolin
2026-06-23 14:52:05 -07:00
parent a68df3bb0e
commit 5ffef47eb7
4 changed files with 204 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

69
codex-rs/core/src/session/context_window.rs
View File

@@ -10,21 +10,17 @@ pub(crate) struct ContextWindowTokenStatus {
pub(crate) auto_compact_scope_tokens: i64,
pub(crate) auto_compact_scope_limit: i64,
pub(crate) full_context_window_limit: Option<i64>,
pub(crate) tokens_until_compaction: i64,
pub(crate) auto_compact_window_prefill_tokens: Option<i64>,
pub(crate) full_context_window_limit_reached: bool,
pub(crate) token_limit_reached: bool,
pub(crate) context_remaining_tokens: Option<i64>,
}
impl ContextWindowTokenStatus {
pub(crate) fn tokens_until_compaction(&self) -> i64 {
let full_context_remaining = self.full_context_window_limit.map_or(i64::MAX, |limit| {
limit.saturating_sub(self.active_context_tokens)
});
self.auto_compact_scope_limit
.saturating_sub(self.auto_compact_scope_tokens)
.min(full_context_remaining)
.max(0)
}
struct BodyAfterPrefixWindowStatus {
full_context_window_limit: Option<i64>,
auto_compact_window_prefill_tokens: Option<i64>,
has_context_remaining_limit: bool,
}
pub(crate) async fn context_window_token_status(
@@ -32,8 +28,8 @@ pub(crate) async fn context_window_token_status(
turn_context: &TurnContext,
) -> ContextWindowTokenStatus {
let active_context_tokens = sess.get_total_token_usage().await;
let mut auto_compact_window_prefill_tokens = None;
let (auto_compact_scope_tokens, auto_compact_scope_limit, full_context_window_limit) =
let (auto_compact_scope_tokens, auto_compact_scope_limit, body_window) =
match turn_context.config.model_auto_compact_token_limit_scope {
AutoCompactTokenLimitScope::Total => (
active_context_tokens,
@@ -45,32 +41,67 @@ pub(crate) async fn context_window_token_status(
),
AutoCompactTokenLimitScope::BodyAfterPrefix => {
let window = sess.auto_compact_window_snapshot().await;
auto_compact_window_prefill_tokens = window.prefill_input_tokens;
let baseline = window.prefill_input_tokens.unwrap_or(active_context_tokens);
let scope_limit = turn_context
.config
.model_auto_compact_token_limit
.or_else(|| turn_context.model_info.auto_compact_token_limit());
let full_context_window_limit = turn_context.model_context_window();
(
active_context_tokens.saturating_sub(baseline),
turn_context
.config
.model_auto_compact_token_limit
.or_else(|| turn_context.model_info.auto_compact_token_limit())
.unwrap_or(i64::MAX),
turn_context.model_context_window(),
scope_limit.unwrap_or(i64::MAX),
Some(BodyAfterPrefixWindowStatus {
full_context_window_limit,
auto_compact_window_prefill_tokens: window.prefill_input_tokens,
has_context_remaining_limit: scope_limit.is_some()
|| full_context_window_limit.is_some(),
}),
)
}
};
let full_context_window_limit = body_window
.as_ref()
.and_then(|window| window.full_context_window_limit);
let auto_compact_window_prefill_tokens = body_window
.as_ref()
.and_then(|window| window.auto_compact_window_prefill_tokens);
let full_context_window_limit_reached =
full_context_window_limit.is_some_and(|full_context_window_limit| {
active_context_tokens >= full_context_window_limit
});
let token_limit_reached =
auto_compact_scope_tokens >= auto_compact_scope_limit || full_context_window_limit_reached;
let full_context_remaining = full_context_window_limit.map_or(i64::MAX, |limit| {
limit.saturating_sub(active_context_tokens)
});
let tokens_until_compaction = auto_compact_scope_limit
.saturating_sub(auto_compact_scope_tokens)
.min(full_context_remaining)
.max(0);
let context_remaining_tokens = match body_window {
Some(body_window) => body_window
.has_context_remaining_limit
.then_some(tokens_until_compaction),
None => turn_context
.model_context_window()
.map(|limit| limit.saturating_sub(active_context_tokens).max(0)),
};
ContextWindowTokenStatus {
active_context_tokens,
auto_compact_scope_tokens,
auto_compact_scope_limit,
full_context_window_limit,
tokens_until_compaction,
auto_compact_window_prefill_tokens,
full_context_window_limit_reached,
token_limit_reached,
context_remaining_tokens,
}
}

3
codex-rs/core/src/session/turn.rs
View File

@@ -324,11 +324,10 @@ pub(crate) async fn run_turn(
"post sampling token usage"
);
let tokens_until_compaction = token_status.tokens_until_compaction();
super::token_budget::maybe_record(
sess.as_ref(),
turn_context.as_ref(),
tokens_until_compaction,
token_status.tokens_until_compaction,
)
.await;

20
codex-rs/core/src/tools/handlers/get_context_remaining.rs
View File

@@ -75,19 +75,15 @@ impl ToolExecutor<ToolInvocation> for GetContextRemainingHandler {
));
}
let Some(model_context_window) = invocation.turn.model_context_window() else {
return Ok(boxed_tool_output(GetContextRemainingOutput::new(
/*tokens_left*/ None,
)));
};
let active_context_tokens = invocation.session.get_total_token_usage().await.max(0);
let tokens_left = model_context_window
.saturating_sub(active_context_tokens)
.max(0);
let token_status = crate::session::context_window::context_window_token_status(
invocation.session.as_ref(),
invocation.turn.as_ref(),
)
.await;
Ok(boxed_tool_output(GetContextRemainingOutput::new(Some(
tokens_left,
))))
Ok(boxed_tool_output(GetContextRemainingOutput::new(
token_status.context_remaining_tokens,
)))
})
}
}

145
codex-rs/core/tests/suite/token_budget.rs
View File

@@ -4,6 +4,7 @@ use codex_config::types::McpServerTransportConfig;
use codex_core::config::TokenBudgetConfig;
use codex_features::Feature;
use codex_model_provider_info::built_in_model_providers;
use codex_protocol::config_types::AutoCompactTokenLimitScope;
use codex_protocol::protocol::CONTEXT_WINDOW_CLOSE_TAG;
use codex_protocol::protocol::CONTEXT_WINDOW_OPEN_TAG;
use codex_protocol::protocol::EventMsg;
@@ -79,6 +80,22 @@ fn tool_names(request: &ResponsesRequest) -> Vec<String> {
.collect()
}
fn ev_completed_with_usage(id: &str, input_tokens: i64, output_tokens: i64) -> Value {
json!({
"type": "response.completed",
"response": {
"id": id,
"usage": {
"input_tokens": input_tokens,
"input_tokens_details": null,
"output_tokens": output_tokens,
"output_tokens_details": null,
"total_tokens": input_tokens + output_tokens
}
}
})
}
#[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(()));
@@ -264,6 +281,70 @@ async fn token_budget_reminder_emits_after_crossing_compaction_threshold() -> Re
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn token_budget_reminder_uses_body_after_prefix_window() -> 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*/ 8_000),
]),
sse(vec![
ev_response_created("resp-2"),
ev_completed_with_tokens("resp-2", /*total_tokens*/ 8_600),
]),
sse(vec![ev_response_created("resp-3"), ev_completed("resp-3")]),
],
)
.await;
let test = test_codex()
.with_config(|config| {
config.model_context_window = Some(10_000);
config.model_auto_compact_token_limit = Some(1_000);
config.model_auto_compact_token_limit_scope =
AutoCompactTokenLimitScope::BodyAfterPrefix;
config.token_budget = Some(TokenBudgetConfig {
reminder_threshold_tokens: Some(600),
..TokenBudgetConfig::default()
});
config
.features
.enable(Feature::TokenBudget)
.expect("test config should allow token budget");
})
.build(&server)
.await?;
test.submit_turn("establish prefix").await?;
test.submit_turn("grow body").await?;
test.submit_turn("observe reminder").await?;
let requests = responses.requests();
assert_eq!(requests.len(), 3);
let reminder = "Your context window is nearly exhausted (only 400 tokens remaining) and will be automatically reset for you soon. Once reset, message items in current context window will be cleared in the new window, but notes and history items will be persistent across windows.";
assert!(
requests[1]
.message_input_texts("developer")
.into_iter()
.all(|text| text != reminder),
"first-window prefix should not count against the body-after-prefix reminder threshold"
);
assert_eq!(
requests[2]
.message_input_texts("developer")
.into_iter()
.filter(|text| text == reminder)
.count(),
1
);
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(()));
@@ -327,6 +408,70 @@ async fn get_context_remaining_returns_token_budget_remaining_fragment() -> Resu
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn get_context_remaining_uses_body_after_prefix_window() -> 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_usage(
"resp-1", /*input_tokens*/ 2_000, /*output_tokens*/ 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.model_auto_compact_token_limit = Some(7_000);
config.model_auto_compact_token_limit_scope =
AutoCompactTokenLimitScope::BodyAfterPrefix;
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 remaining_context = "You have 6500 tokens left in this context window.".to_string();
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(()));