Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions codex-rs/core/config.schema.json
Original file line number Diff line number Diff line change
Expand Up @@ -3119,6 +3119,14 @@
"maxLength": 2000,
"type": "string"
},
"mode": {
"allOf": [
{
"$ref": "#/definitions/TokenBudgetMode"
}
],
"description": "Select whether context-window metadata identifies the thread or agent name."
},
"reminder_message_template": {
"description": "Reminder template. `{n_remaining}` is replaced with the tokens remaining before auto-compaction.",
"maxLength": 2000,
Expand All @@ -3134,6 +3142,14 @@
},
"type": "object"
},
"TokenBudgetMode": {
"description": "Identity included in the context-window developer message.",
"enum": [
"thread",
"name"
],
"type": "string"
},
"ToolExposureSurface": {
"description": "A model-facing surface on which a tool can be exposed.",
"oneOf": [
Expand Down
9 changes: 9 additions & 0 deletions codex-rs/core/src/config/config_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -609,20 +609,29 @@ async fn load_config_resolves_token_budget_config() -> std::io::Result<()> {
r#"
[features.token_budget]
enabled = true
mode = "thread"
reminder_threshold_tokens = 16000
reminder_message_template = "Custom reminder: {n_remaining} tokens."
guidance_message = "Preserve important state before compaction."
auto_compact_fallback_prompt = " Write notes immediately. "
auto_compact_fallback_buffer_tokens = 8000
"#,
TokenBudgetConfig {
mode: TokenBudgetMode::Thread,
reminder_threshold_tokens: Some(16_000),
reminder_message_template: "Custom reminder: {n_remaining} tokens.".to_string(),
guidance_message: Some("Preserve important state before compaction.".to_string()),
auto_compact_fallback_prompt: Some("Write notes immediately.".to_string()),
auto_compact_fallback_buffer_tokens: Some(8_000),
},
),
(
"[features.token_budget]\nenabled = true\nmode = \"name\"\n",
TokenBudgetConfig {
mode: TokenBudgetMode::Name,
..TokenBudgetConfig::default()
},
),
] {
let codex_home = tempdir()?;
let config_toml = toml::from_str(config_toml).expect("TOML should deserialize");
Expand Down
7 changes: 7 additions & 0 deletions codex-rs/core/src/config/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,7 @@ use codex_features::FeaturesToml;
use codex_features::MultiAgentV2ConfigToml;
use codex_features::NetworkProxyConfigToml;
use codex_features::TokenBudgetConfigToml;
use codex_features::TokenBudgetMode;
use codex_git_utils::resolve_root_git_project_for_trust;
use codex_http_client::HttpClientFactory;
use codex_http_client::OutboundProxyPolicy;
Expand Down Expand Up @@ -1119,6 +1120,7 @@ const AUTO_COMPACT_FALLBACK_PROMPT_MAX_BYTES: usize = 2000;

#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct TokenBudgetConfig {
pub mode: TokenBudgetMode,
pub reminder_threshold_tokens: Option<i64>,
pub reminder_message_template: String,
pub guidance_message: Option<String>,
Expand Down Expand Up @@ -1211,6 +1213,7 @@ impl TokenBudgetConfig {
impl Default for TokenBudgetConfig {
fn default() -> Self {
Self {
mode: TokenBudgetMode::default(),
reminder_threshold_tokens: None,
reminder_message_template: DEFAULT_TOKEN_BUDGET_REMINDER_MESSAGE_TEMPLATE.to_string(),
guidance_message: None,
Expand Down Expand Up @@ -2757,6 +2760,9 @@ fn resolve_token_budget_config(
}

let token_budget_config = token_budget_toml_config(config_toml.features.as_ref());
let mode = token_budget_config
.and_then(|config| config.mode)
.unwrap_or_default();
let reminder_threshold_tokens =
token_budget_config.and_then(|config| config.reminder_threshold_tokens);
let reminder_message_template = token_budget_config
Expand All @@ -2774,6 +2780,7 @@ fn resolve_token_budget_config(
token_budget_config.and_then(|config| config.auto_compact_fallback_buffer_tokens);

let token_budget = TokenBudgetConfig {
mode,
reminder_threshold_tokens,
reminder_message_template,
guidance_message,
Expand Down
11 changes: 9 additions & 2 deletions codex-rs/core/src/context/token_budget_context.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use super::ContextualUserFragment;
use codex_features::TokenBudgetMode;
use codex_protocol::AgentPath;
use codex_protocol::ThreadId;
use codex_protocol::protocol::CONTEXT_WINDOW_CLOSE_TAG;
Expand All @@ -11,6 +12,7 @@ use uuid::Uuid;
pub(crate) struct TokenBudgetContext {
thread_id: ThreadId,
agent_path: AgentPath,
mode: TokenBudgetMode,
first_window_id: Uuid,
previous_window_id: Option<Uuid>,
window_id: Uuid,
Expand All @@ -21,6 +23,7 @@ impl TokenBudgetContext {
pub(crate) fn new(
thread_id: ThreadId,
agent_path: AgentPath,
mode: TokenBudgetMode,
first_window_id: Uuid,
previous_window_id: Option<Uuid>,
window_id: Uuid,
Expand All @@ -29,6 +32,7 @@ impl TokenBudgetContext {
Self {
thread_id,
agent_path,
mode,
first_window_id,
previous_window_id,
window_id,
Expand All @@ -51,11 +55,14 @@ impl ContextualUserFragment for TokenBudgetContext {
}

fn body(&self) -> String {
let agent_path = &self.agent_path;
let identity = match self.mode {
TokenBudgetMode::Thread => format!("Thread id: {}", self.thread_id),
TokenBudgetMode::Name => format!("Agent name: {}", self.agent_path),
};
let first_window_id = self.first_window_id;
let window_id = self.window_id;
let mut lines = vec![
format!("Agent name: {agent_path}"),
identity,
format!("First context window id: {first_window_id}"),
format!("Current context window id: {window_id}"),
];
Expand Down
2 changes: 2 additions & 0 deletions codex-rs/core/src/session/config_lock.rs
Original file line number Diff line number Diff line change
Expand Up @@ -272,6 +272,7 @@ mod tests {
config.multi_agent_v2.subagent_developer_instructions =
Some("Locked subagent developer instructions.".to_string());
config.token_budget = Some(crate::config::TokenBudgetConfig {
mode: codex_features::TokenBudgetMode::Thread,
reminder_threshold_tokens: Some(16_000),
reminder_message_template: "Locked reminder: {n_remaining} tokens.".to_string(),
guidance_message: Some("Locked context-window guidance.".to_string()),
Expand Down Expand Up @@ -371,6 +372,7 @@ mod tests {
features.token_budget,
Some(FeatureToml::Config(TokenBudgetConfigToml {
enabled: Some(true),
mode: Some(codex_features::TokenBudgetMode::Thread),
reminder_threshold_tokens: Some(16_000),
reminder_message_template: Some(
"Locked reminder: {n_remaining} tokens.".to_string()
Expand Down
6 changes: 6 additions & 0 deletions codex-rs/core/src/session/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3502,6 +3502,12 @@ impl Session {
session_source
.get_agent_path()
.unwrap_or_else(codex_protocol::AgentPath::root),
turn_context
.config
.token_budget
.as_ref()
.map(|config| config.mode)
.unwrap_or_default(),
auto_compact_window_ids.first_window_id,
auto_compact_window_ids.previous_window_id,
auto_compact_window_ids.window_id,
Expand Down
7 changes: 6 additions & 1 deletion codex-rs/core/src/session/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9105,7 +9105,12 @@ async fn build_initial_context_adds_multi_agent_v2_subagent_usage_hint_as_develo
let turn_context_mut =
Arc::get_mut(&mut turn_context).expect("thread settings should not be shared");
turn_context_mut.session_source = session_source;
Arc::make_mut(&mut turn_context_mut.config)
let config = Arc::make_mut(&mut turn_context_mut.config);
config.token_budget = Some(crate::config::TokenBudgetConfig {
mode: codex_features::TokenBudgetMode::Name,
..crate::config::TokenBudgetConfig::default()
});
config
.features
.enable(Feature::TokenBudget)
.expect("test config should allow token budget");
Expand Down
18 changes: 13 additions & 5 deletions codex-rs/core/src/session/token_budget.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,11 +13,14 @@ pub(super) fn has_explicit_settings(config: &Config) -> bool {
.get("features")
.and_then(|features| features.get("token_budget"))
.and_then(|token_budget| token_budget.as_table())
.is_some_and(|settings| settings.keys().any(|key| key != "enabled"))
|| config
.token_budget
.as_ref()
.is_some_and(|token_budget| token_budget != &TokenBudgetConfig::default())
.is_some_and(|settings| settings.keys().any(|key| key != "enabled" && key != "mode"))
|| config.token_budget.as_ref().is_some_and(|token_budget| {
token_budget
!= &TokenBudgetConfig {
mode: token_budget.mode,
..TokenBudgetConfig::default()
}
})
}

pub(super) fn apply_model_defaults(config: &mut Config, model_info: &ModelInfo) {
Expand All @@ -34,6 +37,11 @@ pub(super) fn apply_model_defaults(config: &mut Config, model_info: &ModelInfo)
};

let token_budget = TokenBudgetConfig {
mode: config
.token_budget
.as_ref()
.map(|token_budget| token_budget.mode)
.unwrap_or_default(),
reminder_threshold_tokens: Some(model_defaults.reminder_threshold_tokens),
reminder_message_template: model_defaults.reminder_message_template.clone(),
guidance_message: Some(model_defaults.guidance_message.clone()),
Expand Down
25 changes: 21 additions & 4 deletions codex-rs/core/tests/suite/token_budget.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ use codex_extension_api::ExtensionRegistryBuilder;
use codex_features::Feature;
use codex_features::FeatureToml;
use codex_features::TokenBudgetConfigToml;
use codex_features::TokenBudgetMode;
use codex_model_provider_info::built_in_model_providers;
use codex_protocol::config_types::AutoCompactTokenLimitScope;
use codex_protocol::items::TurnItem;
Expand Down Expand Up @@ -74,7 +75,7 @@ fn model_token_budget_config() -> ModelTokenBudgetConfig {
}

fn token_budget_contexts(request: &ResponsesRequest) -> Vec<String> {
let context_window_prefix = format!("{CONTEXT_WINDOW_OPEN_TAG}\nAgent name: ");
let context_window_prefix = format!("{CONTEXT_WINDOW_OPEN_TAG}\n");
request
.message_input_texts("developer")
.into_iter()
Expand All @@ -85,7 +86,7 @@ fn token_budget_contexts(request: &ResponsesRequest) -> Vec<String> {
fn token_budget_window_ids(text: &str, agent_name: &str) -> (String, Option<String>, String) {
let captures = assert_regex_match(
&format!(
r"^{CONTEXT_WINDOW_OPEN_TAG}\nAgent name: {agent_name}\nFirst context window id: ([0-9a-f-]{{36}})\nCurrent context window id: ([0-9a-f-]{{36}})(?:\nPrevious context window id: ([0-9a-f-]{{36}}))?\n{CONTEXT_WINDOW_CLOSE_TAG}$"
r"^{CONTEXT_WINDOW_OPEN_TAG}\n(?:Thread id: [0-9a-f-]{{36}}|Agent name: {agent_name})\nFirst context window id: ([0-9a-f-]{{36}})\nCurrent context window id: ([0-9a-f-]{{36}})(?:\nPrevious context window id: ([0-9a-f-]{{36}}))?\n{CONTEXT_WINDOW_CLOSE_TAG}$"
),
text,
);
Expand Down Expand Up @@ -290,6 +291,13 @@ async fn token_budget_uses_model_message_defaults() -> Result<()> {
let model_defaults = model_token_budget_config();
let expected_guidance = model_defaults.guidance_message.clone();
let test = test_codex()
.with_pre_build_hook(|home| {
std::fs::write(
home.join("config.toml"),
"[features.token_budget]\nenabled = true\nmode = \"name\"\n",
)
.expect("write agent-name token-budget configuration");
})
.with_model_info_override("gpt-5.2", move |model_info| {
model_info
.model_messages
Expand All @@ -310,7 +318,14 @@ async fn token_budget_uses_model_message_defaults() -> Result<()> {
test.submit_turn("inspect model-owned context guidance")
.await?;

let developer_texts = response.single_request().message_input_texts("developer");
let request = response.single_request();
let token_budget_context = token_budget_contexts(&request);
assert_eq!(token_budget_context.len(), 1);
assert!(
token_budget_context[0]
.starts_with(&format!("{CONTEXT_WINDOW_OPEN_TAG}\nAgent name: /root\n"))
);
let developer_texts = request.message_input_texts("developer");
assert!(developer_texts.iter().any(|text| {
text == &format!(
"{CONTEXT_WINDOW_GUIDANCE_OPEN_TAG}\n{expected_guidance}\n{CONTEXT_WINDOW_GUIDANCE_CLOSE_TAG}"
Expand Down Expand Up @@ -414,6 +429,7 @@ async fn token_budget_model_defaults_survive_config_lock_replay() -> Result<()>
.and_then(|features| features.token_budget.as_ref()),
Some(&FeatureToml::Config(TokenBudgetConfigToml {
enabled: Some(true),
mode: Some(TokenBudgetMode::Thread),
reminder_threshold_tokens: Some(6_144),
reminder_message_template: Some(
"Model reminder: {n_remaining} tokens remain.".to_string()
Expand Down Expand Up @@ -727,7 +743,7 @@ async fn token_budget_context_injects_plain_thread_hint_text() -> Result<()> {
assert_eq!(token_budgets.len(), 1);
let captures = assert_regex_match(
&format!(
r"^{CONTEXT_WINDOW_OPEN_TAG}\nAgent name: /root\nFirst context window id: ([0-9a-f-]{{36}})\nCurrent context window id: ([0-9a-f-]{{36}})\nmanual history hint for thread {thread_id}\nunstructured notes/thread_hint fixture result\n{CONTEXT_WINDOW_CLOSE_TAG}$"
r"^{CONTEXT_WINDOW_OPEN_TAG}\nThread id: {thread_id}\nFirst context window id: ([0-9a-f-]{{36}})\nCurrent context window id: ([0-9a-f-]{{36}})\nmanual history hint for thread {thread_id}\nunstructured notes/thread_hint fixture result\n{CONTEXT_WINDOW_CLOSE_TAG}$"
),
&token_budgets[0],
);
Expand Down Expand Up @@ -1492,6 +1508,7 @@ async fn new_context_tool_skips_auto_compact_fallback() -> Result<()> {
.with_config(|config| {
config.model_context_window = Some(10_000);
config.token_budget = Some(TokenBudgetConfig {
mode: TokenBudgetMode::Name,
auto_compact_fallback_prompt: Some(AUTO_COMPACT_FALLBACK_PROMPT.to_string()),
auto_compact_fallback_buffer_tokens: Some(4_000),
..TokenBudgetConfig::default()
Expand Down
12 changes: 12 additions & 0 deletions codex-rs/features/src/feature_configs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -134,11 +134,23 @@ impl FeatureConfig for MultiAgentV2ConfigToml {
}
}

/// Identity included in the context-window developer message.
#[derive(Serialize, Deserialize, Debug, Clone, Copy, Default, PartialEq, Eq, JsonSchema)]
#[serde(rename_all = "snake_case")]
pub enum TokenBudgetMode {
#[default]
Thread,
Name,
}

#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq, JsonSchema)]
#[serde(deny_unknown_fields)]
pub struct TokenBudgetConfigToml {
#[serde(skip_serializing_if = "Option::is_none")]
pub enabled: Option<bool>,
/// Select whether context-window metadata identifies the thread or agent name.
#[serde(skip_serializing_if = "Option::is_none")]
pub mode: Option<TokenBudgetMode>,
/// Number of tokens remaining before auto-compaction when the wrap-up reminder is emitted.
#[serde(skip_serializing_if = "Option::is_none")]
#[schemars(range(min = 1))]
Expand Down
1 change: 1 addition & 0 deletions codex-rs/features/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ pub use feature_configs::NonPrefixedMcpToolNamesConfigToml;
use feature_configs::RemovedAppsMcpPathOverrideConfigToml;
pub use feature_configs::RolloutBudgetConfigToml;
pub use feature_configs::TokenBudgetConfigToml;
pub use feature_configs::TokenBudgetMode;
pub use feature_configs::ToolRegistryConfigToml;
use legacy::LegacyFeatureToggles;
pub use legacy::legacy_feature_keys;
Expand Down
Loading