diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 32478b81d..25ae84cfb 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -2860,6 +2860,11 @@ "enabled": { "type": "boolean" }, + "guidance_message": { + "description": "Guidance appended to the context-window metadata in a developer message.", + "maxLength": 1000, + "type": "string" + }, "reminder_message_template": { "description": "Reminder template. `{n_remaining}` is replaced with the tokens remaining before auto-compaction.", "maxLength": 1000, diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index e13c91842..4d3a7bcc6 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -475,10 +475,12 @@ async fn load_config_resolves_token_budget_config() -> std::io::Result<()> { enabled = true reminder_threshold_tokens = 16000 reminder_message_template = "Custom reminder: {n_remaining} tokens." +guidance_message = "Preserve important state before compaction." "#, TokenBudgetConfig { 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()), }, ), ] { diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 85ed1ee85..ed565c6ad 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -1089,11 +1089,13 @@ pub(crate) const DEFAULT_TOKEN_BUDGET_REMINDER_MESSAGE_TEMPLATE: &str = concat!( "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." ); const TOKEN_BUDGET_REMINDER_MESSAGE_TEMPLATE_MAX_BYTES: usize = 1000; +const TOKEN_BUDGET_GUIDANCE_MESSAGE_MAX_BYTES: usize = 1000; #[derive(Debug, Clone, PartialEq, Eq, Serialize)] pub struct TokenBudgetConfig { pub reminder_threshold_tokens: Option, pub reminder_message_template: String, + pub guidance_message: Option, } impl Default for TokenBudgetConfig { @@ -1101,6 +1103,7 @@ impl Default for TokenBudgetConfig { Self { reminder_threshold_tokens: None, reminder_message_template: DEFAULT_TOKEN_BUDGET_REMINDER_MESSAGE_TEMPLATE.to_string(), + guidance_message: None, } } } @@ -2578,9 +2581,25 @@ fn resolve_token_budget_config( )); } + let guidance_message = token_budget_config + .and_then(|config| config.guidance_message.clone()) + .filter(|message| !message.trim().is_empty()); + if guidance_message + .as_ref() + .is_some_and(|message| message.len() > TOKEN_BUDGET_GUIDANCE_MESSAGE_MAX_BYTES) + { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!( + "features.token_budget.guidance_message must not exceed {TOKEN_BUDGET_GUIDANCE_MESSAGE_MAX_BYTES} bytes" + ), + )); + } + Ok(Some(TokenBudgetConfig { reminder_threshold_tokens, reminder_message_template, + guidance_message, })) } diff --git a/codex-rs/core/src/context/mod.rs b/codex-rs/core/src/context/mod.rs index 4ba6acf02..da7db0b37 100644 --- a/codex-rs/core/src/context/mod.rs +++ b/codex-rs/core/src/context/mod.rs @@ -71,6 +71,7 @@ pub(crate) use realtime_start_with_instructions::RealtimeStartWithInstructions; pub(crate) use recommended_plugins_instructions::RecommendedPluginsInstructions; pub(crate) use rollout_budget::RolloutBudgetContext; pub(crate) use subagent_notification::SubagentNotification; +pub(crate) use token_budget_context::ContextWindowGuidance; pub(crate) use token_budget_context::TokenBudgetContext; pub(crate) use token_budget_context::TokenBudgetRemainingContext; pub(crate) use token_budget_context::TokenBudgetReminder; diff --git a/codex-rs/core/src/context/token_budget_context.rs b/codex-rs/core/src/context/token_budget_context.rs index 3b48aeaa7..f3eef0b73 100644 --- a/codex-rs/core/src/context/token_budget_context.rs +++ b/codex-rs/core/src/context/token_budget_context.rs @@ -1,6 +1,8 @@ use super::ContextualUserFragment; use codex_protocol::ThreadId; use codex_protocol::protocol::CONTEXT_WINDOW_CLOSE_TAG; +use codex_protocol::protocol::CONTEXT_WINDOW_GUIDANCE_CLOSE_TAG; +use codex_protocol::protocol::CONTEXT_WINDOW_GUIDANCE_OPEN_TAG; use codex_protocol::protocol::CONTEXT_WINDOW_OPEN_TAG; use uuid::Uuid; @@ -63,6 +65,40 @@ impl ContextualUserFragment for TokenBudgetContext { } } +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct ContextWindowGuidance { + message: String, +} + +impl ContextWindowGuidance { + pub(crate) fn new(message: &str) -> Self { + Self { + message: message.to_string(), + } + } +} + +impl ContextualUserFragment for ContextWindowGuidance { + fn role(&self) -> &'static str { + "developer" + } + + fn markers(&self) -> (&'static str, &'static str) { + Self::type_markers() + } + + fn type_markers() -> (&'static str, &'static str) { + ( + CONTEXT_WINDOW_GUIDANCE_OPEN_TAG, + CONTEXT_WINDOW_GUIDANCE_CLOSE_TAG, + ) + } + + fn body(&self) -> String { + format!("\n{}\n", self.message) + } +} + #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct TokenBudgetRemainingContext { tokens_left: Option, diff --git a/codex-rs/core/src/event_mapping.rs b/codex-rs/core/src/event_mapping.rs index 07f06c44b..392203c07 100644 --- a/codex-rs/core/src/event_mapping.rs +++ b/codex-rs/core/src/event_mapping.rs @@ -15,6 +15,7 @@ use codex_protocol::models::is_image_open_tag_text; use codex_protocol::models::is_local_image_close_tag_text; use codex_protocol::models::is_local_image_open_tag_text; use codex_protocol::protocol::COLLABORATION_MODE_OPEN_TAG; +use codex_protocol::protocol::CONTEXT_WINDOW_GUIDANCE_OPEN_TAG; use codex_protocol::protocol::CONTEXT_WINDOW_OPEN_TAG; use codex_protocol::protocol::MULTI_AGENT_MODE_OPEN_TAG; use codex_protocol::protocol::REALTIME_CONVERSATION_OPEN_TAG; @@ -38,6 +39,7 @@ const CONTEXTUAL_DEVELOPER_PREFIXES: &[&str] = &[ // Keep recognizing token-budget wrappers persisted by older versions. "", CONTEXT_WINDOW_OPEN_TAG, + CONTEXT_WINDOW_GUIDANCE_OPEN_TAG, "", ]; diff --git a/codex-rs/core/src/event_mapping_tests.rs b/codex-rs/core/src/event_mapping_tests.rs index d7e999afb..c91034d45 100644 --- a/codex-rs/core/src/event_mapping_tests.rs +++ b/codex-rs/core/src/event_mapping_tests.rs @@ -16,6 +16,8 @@ use codex_protocol::models::ReasoningItemReasoningSummary; use codex_protocol::models::ResponseItem; use codex_protocol::models::WebSearchAction; use codex_protocol::protocol::CONTEXT_WINDOW_CLOSE_TAG; +use codex_protocol::protocol::CONTEXT_WINDOW_GUIDANCE_CLOSE_TAG; +use codex_protocol::protocol::CONTEXT_WINDOW_GUIDANCE_OPEN_TAG; use codex_protocol::protocol::CONTEXT_WINDOW_OPEN_TAG; use codex_protocol::protocol::SKILLS_INSTRUCTIONS_OPEN_TAG; use codex_protocol::user_input::UserInput; @@ -55,6 +57,18 @@ Thread id: 00000000-0000-0000-0000-000000000000 assert!(!has_non_contextual_dev_message_content(&content)); } +#[test] +fn recognizes_context_window_guidance_as_contextual_developer_content() { + let content = vec![ContentItem::InputText { + text: format!( + "{CONTEXT_WINDOW_GUIDANCE_OPEN_TAG}\nPreserve important state.\n{CONTEXT_WINDOW_GUIDANCE_CLOSE_TAG}" + ), + }]; + + assert!(is_contextual_dev_message_content(&content)); + assert!(!has_non_contextual_dev_message_content(&content)); +} + #[test] fn parses_user_message_with_text_and_two_images() { let img1 = "https://example.com/one.png".to_string(); diff --git a/codex-rs/core/src/session/config_lock.rs b/codex-rs/core/src/session/config_lock.rs index ec7fd6de7..b1adf9446 100644 --- a/codex-rs/core/src/session/config_lock.rs +++ b/codex-rs/core/src/session/config_lock.rs @@ -248,6 +248,7 @@ mod tests { config.token_budget = Some(crate::config::TokenBudgetConfig { 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()), }); config .features @@ -340,6 +341,7 @@ mod tests { reminder_message_template: Some( "Locked reminder: {n_remaining} tokens.".to_string() ), + guidance_message: Some("Locked context-window guidance.".to_string()), })) ); diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index daf981258..8a9055aa2 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3339,6 +3339,16 @@ impl Session { ) .render(), ); + if let Some(guidance_message) = turn_context + .config + .token_budget + .as_ref() + .and_then(|config| config.guidance_message.as_deref()) + .filter(|message| !message.trim().is_empty()) + { + developer_sections + .push(crate::context::ContextWindowGuidance::new(guidance_message).render()); + } } for fragment in world_state.render_full() { match fragment.role() { diff --git a/codex-rs/core/tests/suite/token_budget.rs b/codex-rs/core/tests/suite/token_budget.rs index 6d43b9037..18d800850 100644 --- a/codex-rs/core/tests/suite/token_budget.rs +++ b/codex-rs/core/tests/suite/token_budget.rs @@ -7,6 +7,8 @@ use codex_model_provider_info::built_in_model_providers; use codex_protocol::config_types::AutoCompactTokenLimitScope; use codex_protocol::items::TurnItem; use codex_protocol::protocol::CONTEXT_WINDOW_CLOSE_TAG; +use codex_protocol::protocol::CONTEXT_WINDOW_GUIDANCE_CLOSE_TAG; +use codex_protocol::protocol::CONTEXT_WINDOW_GUIDANCE_OPEN_TAG; use codex_protocol::protocol::CONTEXT_WINDOW_OPEN_TAG; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::HookEventName; @@ -211,6 +213,52 @@ async fn token_budget_context_is_only_emitted_with_full_context() -> Result<()> Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn token_budget_guidance_follows_context_window() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let response = mount_sse_sequence( + &server, + vec![sse(vec![ + ev_response_created("resp-1"), + ev_completed("resp-1"), + ])], + ) + .await; + let guidance_message = "Preserve important state before compaction."; + let test = test_codex() + .with_config(move |config| { + config.model_context_window = Some(CONFIGURED_CONTEXT_WINDOW); + config.token_budget = Some(TokenBudgetConfig { + guidance_message: Some(guidance_message.to_string()), + ..TokenBudgetConfig::default() + }); + config + .features + .enable(Feature::TokenBudget) + .expect("test config should allow token budget"); + }) + .build_with_auto_env(&server) + .await?; + + test.submit_turn("inspect context guidance").await?; + + let developer_texts = response.single_request().message_input_texts("developer"); + let context_window_index = developer_texts + .iter() + .position(|text| text.starts_with(CONTEXT_WINDOW_OPEN_TAG)) + .expect("context-window metadata should be present"); + assert_eq!( + developer_texts.get(context_window_index + 1), + Some(&format!( + "{CONTEXT_WINDOW_GUIDANCE_OPEN_TAG}\n{guidance_message}\n{CONTEXT_WINDOW_GUIDANCE_CLOSE_TAG}" + )) + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn token_budget_context_injects_plain_thread_hint_text() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/features/src/feature_configs.rs b/codex-rs/features/src/feature_configs.rs index efb3580a5..a5e9104bc 100644 --- a/codex-rs/features/src/feature_configs.rs +++ b/codex-rs/features/src/feature_configs.rs @@ -88,6 +88,10 @@ pub struct TokenBudgetConfigToml { #[serde(skip_serializing_if = "Option::is_none")] #[schemars(length(min = 1, max = 1000))] pub reminder_message_template: Option, + /// Guidance appended to the context-window metadata in a developer message. + #[serde(skip_serializing_if = "Option::is_none")] + #[schemars(length(max = 1000))] + pub guidance_message: Option, } impl FeatureConfig for TokenBudgetConfigToml { diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 138b36f10..a003203d3 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -114,6 +114,8 @@ pub const REALTIME_CONVERSATION_OPEN_TAG: &str = ""; pub const REALTIME_CONVERSATION_CLOSE_TAG: &str = ""; pub const CONTEXT_WINDOW_OPEN_TAG: &str = ""; pub const CONTEXT_WINDOW_CLOSE_TAG: &str = ""; +pub const CONTEXT_WINDOW_GUIDANCE_OPEN_TAG: &str = ""; +pub const CONTEXT_WINDOW_GUIDANCE_CLOSE_TAG: &str = ""; pub const USER_MESSAGE_BEGIN: &str = "## My request for Codex:"; // TODO(anp): Replace `TurnEnvironmentSelection` with `PathUri` once path URIs carry environment