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

Budget realtime current thread context (#17519)

Browse Source 
Select Current Thread startup context by budget from newest turns, cap
each rendered turn at 300 approximate tokens, and add formatter plus
integration snapshot coverage.
This commit is contained in:
Ahmed Ibrahim
2026-04-12 11:59:09 -07:00
committed by GitHub
Unverified
parent 1288bb60a1
commit 4db60d5d8b
4 changed files with 389 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/realtime_context.rs
+56 -26
View File
@@ -15,6 +15,7 @@ use std::collections::HashSet;
use std::ffi::OsStr;
use std::fs::DirEntry;
use std::io;
use std::mem::take;
use std::path::Path;
use std::path::PathBuf;
use tracing::debug;
@@ -26,7 +27,7 @@ const CURRENT_THREAD_SECTION_TOKEN_BUDGET: usize = 1_200;
const RECENT_WORK_SECTION_TOKEN_BUDGET: usize = 2_200;
const WORKSPACE_SECTION_TOKEN_BUDGET: usize = 1_600;
const NOTES_SECTION_TOKEN_BUDGET: usize = 300;
const MAX_CURRENT_THREAD_TURNS: usize = 2;
const CURRENT_THREAD_TURN_TOKEN_BUDGET: usize = 300;
const MAX_RECENT_THREADS: usize = 40;
const MAX_RECENT_WORK_GROUPS: usize = 8;
const MAX_CURRENT_CWD_ASKS: usize = 8;
@@ -204,10 +205,7 @@ fn build_current_thread_section(items: &[ResponseItem]) -> Option<String> {
continue;
};
if !current_user.is_empty() || !current_assistant.is_empty() {
turns.push((
std::mem::take(&mut current_user),
std::mem::take(&mut current_assistant),
));
turns.push((take(&mut current_user), take(&mut current_assistant)));
}
current_user.push(text);
}
@@ -231,43 +229,75 @@ fn build_current_thread_section(items: &[ResponseItem]) -> Option<String> {
turns.push((current_user, current_assistant));
}
let retained_turns = turns
.into_iter()
.rev()
.take(MAX_CURRENT_THREAD_TURNS)
.collect::<Vec<_>>()
.into_iter()
.rev()
.collect::<Vec<_>>();
if retained_turns.is_empty() {
if turns.is_empty() {
return None;
}
let mut lines = vec![
"Most recent user/assistant turns from this exact thread. Use them for continuity when responding.".to_string(),
];
let mut remaining_budget =
CURRENT_THREAD_SECTION_TOKEN_BUDGET.saturating_sub(approx_token_count(&lines.join("\n")));
let mut retained_turn_count = 0;
let retained_turn_count = retained_turns.len();
for (index, (user_messages, assistant_messages)) in retained_turns.into_iter().enumerate() {
lines.push(String::new());
if retained_turn_count == 1 || index + 1 == retained_turn_count {
lines.push("### Latest turn".to_string());
for (index, (user_messages, assistant_messages)) in turns.into_iter().rev().enumerate() {
if remaining_budget == 0 {
break;
}
let mut turn_lines = Vec::new();
if index == 0 {
turn_lines.push("### Latest turn".to_string());
} else {
lines.push(format!("### Prior turn {}", index + 1));
turn_lines.push(format!("### Previous turn {index}"));
}
if !user_messages.is_empty() {
lines.push("User:".to_string());
lines.push(user_messages.join("\n\n"));
turn_lines.push("User:".to_string());
turn_lines.push(user_messages.join("\n\n"));
}
if !assistant_messages.is_empty() {
lines.push(String::new());
lines.push("Assistant:".to_string());
lines.push(assistant_messages.join("\n\n"));
turn_lines.push(String::new());
turn_lines.push("Assistant:".to_string());
turn_lines.push(assistant_messages.join("\n\n"));
}
let turn_budget = CURRENT_THREAD_TURN_TOKEN_BUDGET.min(remaining_budget);
let turn_text = turn_lines.join("\n");
let mut truncation_budget = turn_budget;
let turn_text = loop {
let candidate = truncate_text(&turn_text, TruncationPolicy::Tokens(truncation_budget));
let candidate_tokens = approx_token_count(&candidate);
if candidate_tokens <= turn_budget {
break candidate;
}
// The shared truncator adds its marker after choosing preserved
// content, so tighten the content budget until the rendered turn
// itself fits the per-turn cap.
let excess_tokens = candidate_tokens.saturating_sub(turn_budget);
let next_budget = truncation_budget.saturating_sub(excess_tokens.max(1));
if next_budget == 0 {
let candidate = truncate_text(&turn_text, TruncationPolicy::Tokens(0));
if approx_token_count(&candidate) <= turn_budget {
break candidate;
}
break String::new();
}
truncation_budget = next_budget;
};
let turn_tokens = approx_token_count(&turn_text);
if turn_tokens == 0 {
continue;
}
lines.push(String::new());
lines.push(turn_text);
remaining_budget = remaining_budget.saturating_sub(turn_tokens);
retained_turn_count += 1;
}
Some(lines.join("\n"))
(retained_turn_count > 0).then(|| lines.join("\n"))
}
fn build_workspace_section_with_user_root(
codex-rs/core/src/realtime_context_tests.rs
+117
View File
@@ -1,8 +1,11 @@
use super::build_current_thread_section;
use super::build_recent_work_section;
use super::build_workspace_section_with_user_root;
use chrono::TimeZone;
use chrono::Utc;
use codex_protocol::ThreadId;
use codex_protocol::models::ContentItem;
use codex_protocol::models::ResponseItem;
use codex_state::ThreadMetadata;
use pretty_assertions::assert_eq;
use std::fs;
@@ -43,6 +46,120 @@ fn thread_metadata(cwd: &str, title: &str, first_user_message: &str) -> ThreadMe
}
}
fn message(role: &str, content: ContentItem) -> ResponseItem {
ResponseItem::Message {
id: None,
role: role.to_string(),
content: vec![content],
end_turn: None,
phase: None,
}
}
fn user_message(text: impl Into<String>) -> ResponseItem {
message("user", ContentItem::InputText { text: text.into() })
}
fn assistant_message(text: impl Into<String>) -> ResponseItem {
message("assistant", ContentItem::OutputText { text: text.into() })
}
fn long_turn_text(index: usize) -> String {
format!(
"turn-{index}-start {} turn-{index}-middle {} turn-{index}-end",
"head filler ".repeat(160),
"tail filler ".repeat(240),
)
}
#[test]
fn current_thread_section_includes_short_turns_newest_first_until_budget() {
let items = vec![
user_message("user turn 1"),
assistant_message("assistant turn 1"),
user_message("user turn 2"),
assistant_message("assistant turn 2"),
user_message("user turn 3"),
assistant_message("assistant turn 3"),
user_message("user turn 4"),
assistant_message("assistant turn 4"),
];
assert_eq!(
build_current_thread_section(&items),
Some(
r#"Most recent user/assistant turns from this exact thread. Use them for continuity when responding.
### Latest turn
User:
user turn 4
Assistant:
assistant turn 4
### Previous turn 1
User:
user turn 3
Assistant:
assistant turn 3
### Previous turn 2
User:
user turn 2
Assistant:
assistant turn 2
### Previous turn 3
User:
user turn 1
Assistant:
assistant turn 1"#
.to_string()
)
);
}
#[test]
fn current_thread_turn_truncation_preserves_start_and_end() {
let items = vec![user_message(long_turn_text(/*index*/ 0))];
let section = build_current_thread_section(&items).expect("current thread section");
assert_eq!(
(
section.contains("turn-0-start"),
section.contains("turn-0-middle"),
section.contains("turn-0-end"),
section.contains("tokens truncated"),
),
(true, false, true, true),
);
}
#[test]
fn current_thread_section_keeps_latest_turns_when_history_exceeds_budget() {
let mut items = Vec::new();
for index in 1..=8 {
items.push(user_message(long_turn_text(index)));
items.push(assistant_message(format!("assistant turn {index}")));
}
let section = build_current_thread_section(&items).expect("current thread section");
assert_eq!(
(
section.contains("turn-8-start"),
section.contains("turn-8-end"),
section.contains("### Previous turn 2"),
section.contains("turn-1-start"),
section.contains("turn-1-end"),
),
(true, true, true, false, false),
);
}
#[test]
fn workspace_section_requires_meaningful_structure() {
let cwd = TempDir::new().expect("tempdir");