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
5 changes: 5 additions & 0 deletions codex-rs/core/config.schema.json
Original file line number Diff line number Diff line change
Expand Up @@ -2859,6 +2859,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,
Expand Down
2 changes: 2 additions & 0 deletions codex-rs/core/src/config/config_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()),
},
),
] {
Expand Down
19 changes: 19 additions & 0 deletions codex-rs/core/src/config/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1089,18 +1089,21 @@ 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<i64>,
pub reminder_message_template: String,
pub guidance_message: Option<String>,
}

impl Default for TokenBudgetConfig {
fn default() -> Self {
Self {
reminder_threshold_tokens: None,
reminder_message_template: DEFAULT_TOKEN_BUDGET_REMINDER_MESSAGE_TEMPLATE.to_string(),
guidance_message: None,
}
}
}
Expand Down Expand Up @@ -2575,9 +2578,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,
}))
}

Expand Down
1 change: 1 addition & 0 deletions codex-rs/core/src/context/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
36 changes: 36 additions & 0 deletions codex-rs/core/src/context/token_budget_context.rs
Original file line number Diff line number Diff line change
@@ -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;

Expand Down Expand Up @@ -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<i64>,
Expand Down
2 changes: 2 additions & 0 deletions codex-rs/core/src/event_mapping.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -38,6 +39,7 @@ const CONTEXTUAL_DEVELOPER_PREFIXES: &[&str] = &[
// Keep recognizing token-budget wrappers persisted by older versions.
"<token_budget>",
CONTEXT_WINDOW_OPEN_TAG,
CONTEXT_WINDOW_GUIDANCE_OPEN_TAG,
"<rollout_budget>",
];

Expand Down
14 changes: 14 additions & 0 deletions codex-rs/core/src/event_mapping_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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();
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 @@ -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
Expand Down Expand Up @@ -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()),
}))
);

Expand Down
10 changes: 10 additions & 0 deletions codex-rs/core/src/session/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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() {
Expand Down
48 changes: 48 additions & 0 deletions codex-rs/core/tests/suite/token_budget.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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(()));
Expand Down
4 changes: 4 additions & 0 deletions codex-rs/features/src/feature_configs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<String>,
/// 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<String>,
}

impl FeatureConfig for TokenBudgetConfigToml {
Expand Down
2 changes: 2 additions & 0 deletions codex-rs/protocol/src/protocol.rs
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,8 @@ pub const REALTIME_CONVERSATION_OPEN_TAG: &str = "<realtime_conversation>";
pub const REALTIME_CONVERSATION_CLOSE_TAG: &str = "</realtime_conversation>";
pub const CONTEXT_WINDOW_OPEN_TAG: &str = "<context_window>";
pub const CONTEXT_WINDOW_CLOSE_TAG: &str = "</context_window>";
pub const CONTEXT_WINDOW_GUIDANCE_OPEN_TAG: &str = "<context_window_guidance>";
pub const CONTEXT_WINDOW_GUIDANCE_CLOSE_TAG: &str = "</context_window_guidance>";
pub const USER_MESSAGE_BEGIN: &str = "## My request for Codex:";

// TODO(anp): Replace `TurnEnvironmentSelection` with `PathUri` once path URIs carry environment
Expand Down
Loading