Codex Handbook
core/src/realtime_context.rs 577 lines
use crate::compact::content_items_to_text;use crate::event_mapping::is_contextual_user_message_content;use crate::session::session::Session;use chrono::Utc;use codex_exec_server::LOCAL_FS;use codex_git_utils::resolve_root_git_project_for_trust;use codex_protocol::models::ResponseItem;use codex_protocol::models::plaintext_agent_message_content;use codex_thread_store::ListThreadsParams;use codex_thread_store::SortDirection;use codex_thread_store::StoredThread;use codex_thread_store::ThreadSortKey;use codex_utils_absolute_path::AbsolutePathBuf;use codex_utils_output_truncation::TruncationPolicy;use codex_utils_output_truncation::truncate_text;use dirs::home_dir;use std::cmp::Reverse;use std::collections::HashMap;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;use tracing::info;use tracing::warn;const STARTUP_CONTEXT_HEADER: &str = "Startup context from Codex.\nThis is background context about recent work and machine/workspace layout. It may be incomplete or stale. Use it to inform responses, and do not repeat it back unless relevant.";const STARTUP_CONTEXT_OPEN_TAG: &str = "<startup_context>";const STARTUP_CONTEXT_CLOSE_TAG: &str = "</startup_context>";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;pub(crate) const REALTIME_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;const MAX_OTHER_CWD_ASKS: usize = 5;const MAX_ASK_CHARS: usize = 240;const TREE_MAX_DEPTH: usize = 2;const DIR_ENTRY_LIMIT: usize = 20;const APPROX_BYTES_PER_TOKEN: usize = 4;const NOISY_DIR_NAMES: &[&str] = &[    ".git",    ".next",    ".pytest_cache",    ".ruff_cache",    "__pycache__",    "build",    "dist",    "node_modules",    "out",    "target",];pub(crate) async fn build_realtime_startup_context(    sess: &Session,    budget_tokens: usize,) -> Option<String> {    let config = sess.get_config().await;    let cwd = config.cwd.clone();    let history = sess.clone_history().await;    let current_thread_section = build_current_thread_section(history.raw_items());    let recent_threads = load_recent_threads(sess).await;    let recent_work_section = build_recent_work_section(&cwd, &recent_threads).await;    let workspace_section = build_workspace_section_with_user_root(&cwd, home_dir()).await;    if current_thread_section.is_none()        && recent_work_section.is_none()        && workspace_section.is_none()    {        debug!("realtime startup context unavailable; skipping injection");        return None;    }    let mut parts = vec![STARTUP_CONTEXT_HEADER.to_string()];    let has_current_thread_section = current_thread_section.is_some();    let has_recent_work_section = recent_work_section.is_some();    let has_workspace_section = workspace_section.is_some();    if let Some(section) = format_section(        "Current Thread",        current_thread_section,        CURRENT_THREAD_SECTION_TOKEN_BUDGET,    ) {        parts.push(section);    }    if let Some(section) = format_section(        "Recent Work",        recent_work_section,        RECENT_WORK_SECTION_TOKEN_BUDGET,    ) {        parts.push(section);    }    if let Some(section) = format_section(        "Machine / Workspace Map",        workspace_section,        WORKSPACE_SECTION_TOKEN_BUDGET,    ) {        parts.push(section);    }    if let Some(section) = format_section(        "Notes",        Some("Built at realtime startup from the current thread history, local thread metadata, and a bounded local workspace scan. This excludes repo memory instructions, AGENTS files, project-doc prompt blends, and memory summaries.".to_string()),        NOTES_SECTION_TOKEN_BUDGET,    ) {        parts.push(section);    }    let context = format_startup_context_blob(&parts.join("\n\n"));    debug!(        approx_tokens = approx_token_count(&context),        requested_budget_tokens = budget_tokens,        bytes = context.len(),        has_current_thread_section,        has_recent_work_section,        has_workspace_section,        "built realtime startup context"    );    info!("realtime startup context: {context}");    Some(context)}async fn load_recent_threads(sess: &Session) -> Vec<StoredThread> {    match sess        .services        .thread_store        .list_threads(ListThreadsParams {            page_size: MAX_RECENT_THREADS,            cursor: None,            sort_key: ThreadSortKey::UpdatedAt,            sort_direction: SortDirection::Desc,            allowed_sources: Vec::new(),            model_providers: None,            cwd_filters: None,            parent_thread_id: None,            archived: false,            search_term: None,            use_state_db_only: false,        })        .await    {        Ok(page) => page.items,        Err(err) => {            warn!("failed to load realtime startup threads from thread store: {err}");            Vec::new()        }    }}async fn build_recent_work_section(    cwd: &AbsolutePathBuf,    recent_threads: &[StoredThread],) -> Option<String> {    let mut groups: HashMap<PathBuf, Vec<&StoredThread>> = HashMap::new();    for entry in recent_threads {        let group = match AbsolutePathBuf::from_absolute_path(entry.cwd.as_path()) {            Ok(entry_cwd) => resolve_root_git_project_for_trust(LOCAL_FS.as_ref(), &entry_cwd)                .await                .map(AbsolutePathBuf::into_path_buf)                .unwrap_or_else(|| entry.cwd.clone()),            Err(_) => entry.cwd.clone(),        };        groups.entry(group).or_default().push(entry);    }    let current_group = resolve_root_git_project_for_trust(LOCAL_FS.as_ref(), cwd)        .await        .map(AbsolutePathBuf::into_path_buf)        .unwrap_or_else(|| cwd.clone().into_path_buf());    let mut groups = groups.into_iter().collect::<Vec<_>>();    groups.sort_by(|(left_group, left_entries), (right_group, right_entries)| {        let left_latest = left_entries            .iter()            .map(|entry| entry.updated_at)            .max()            .unwrap_or_else(Utc::now);        let right_latest = right_entries            .iter()            .map(|entry| entry.updated_at)            .max()            .unwrap_or_else(Utc::now);        (            *left_group != current_group,            Reverse(left_latest),            left_group.as_os_str(),        )            .cmp(&(                *right_group != current_group,                Reverse(right_latest),                right_group.as_os_str(),            ))    });    let mut sections = Vec::new();    for (group, mut entries) in groups.into_iter().take(MAX_RECENT_WORK_GROUPS) {        entries.sort_by_key(|entry| Reverse(entry.updated_at));        if let Some(section) = format_thread_group(&current_group, &group, entries).await {            sections.push(section);        }    }    (!sections.is_empty()).then(|| sections.join("\n\n"))}fn build_current_thread_section(items: &[ResponseItem]) -> Option<String> {    let mut turns = Vec::new();    let mut current_user = Vec::new();    let mut current_assistant = Vec::new();    for item in items {        match item {            ResponseItem::Message { role, content, .. } if role == "user" => {                if is_contextual_user_message_content(content) {                    continue;                }                let Some(text) = content_items_to_text(content)                    .map(|text| text.trim().to_string())                    .filter(|text| !text.is_empty())                else {                    continue;                };                if !current_user.is_empty() || !current_assistant.is_empty() {                    turns.push((take(&mut current_user), take(&mut current_assistant)));                }                current_user.push(text);            }            ResponseItem::Message { role, content, .. } if role == "assistant" => {                let Some(text) = content_items_to_text(content)                    .map(|text| text.trim().to_string())                    .filter(|text| !text.is_empty())                else {                    continue;                };                if current_user.is_empty() && current_assistant.is_empty() {                    continue;                }                current_assistant.push(text);            }            ResponseItem::AgentMessage {                author, content, ..            } => {                let Some(text) = plaintext_agent_message_content(content) else {                    continue;                };                if current_user.is_empty() && current_assistant.is_empty() {                    continue;                }                current_assistant.push(format!("Agent message from {author}:\n{text}"));            }            _ => {}        }    }    if !current_user.is_empty() || !current_assistant.is_empty() {        turns.push((current_user, current_assistant));    }    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;    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 {            turn_lines.push(format!("### Previous turn {index}"));        }        if !user_messages.is_empty() {            turn_lines.push("User:".to_string());            turn_lines.push(user_messages.join("\n\n"));        }        if !assistant_messages.is_empty() {            turn_lines.push(String::new());            turn_lines.push("Assistant:".to_string());            turn_lines.push(assistant_messages.join("\n\n"));        }        let turn_budget = REALTIME_TURN_TOKEN_BUDGET.min(remaining_budget);        let turn_text = turn_lines.join("\n");        let turn_text = truncate_realtime_text_to_token_budget(&turn_text, turn_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;    }    (retained_turn_count > 0).then(|| lines.join("\n"))}pub(crate) fn truncate_realtime_text_to_token_budget(text: &str, budget_tokens: usize) -> String {    let mut truncation_budget = budget_tokens;    loop {        let candidate = truncate_text(text, TruncationPolicy::Tokens(truncation_budget));        let candidate_tokens = approx_token_count(&candidate);        if candidate_tokens <= budget_tokens {            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(budget_tokens);        let next_budget = truncation_budget.saturating_sub(excess_tokens.max(1));        if next_budget == 0 {            let candidate = truncate_text(text, TruncationPolicy::Tokens(0));            if approx_token_count(&candidate) <= budget_tokens {                break candidate;            }            break String::new();        }        truncation_budget = next_budget;    }}async fn build_workspace_section_with_user_root(    cwd: &AbsolutePathBuf,    user_root: Option<PathBuf>,) -> Option<String> {    let cwd_path = cwd.as_path();    let git_root = resolve_root_git_project_for_trust(LOCAL_FS.as_ref(), cwd).await;    let cwd_tree = render_tree(cwd_path);    let git_root_tree = git_root        .as_ref()        .filter(|git_root| git_root.as_path() != cwd_path)        .and_then(|git_root| render_tree(git_root.as_path()));    let user_root_tree = user_root        .as_ref()        .filter(|user_root| user_root.as_path() != cwd_path)        .filter(|user_root| {            git_root                .as_ref()                .is_none_or(|git_root| git_root.as_path() != user_root.as_path())        })        .and_then(|user_root| render_tree(user_root));    if cwd_tree.is_none() && git_root.is_none() && user_root_tree.is_none() {        return None;    }    let mut lines = vec![        format!("Current working directory: {}", cwd_path.display()),        format!("Working directory name: {}", file_name_string(cwd_path)),    ];    if let Some(git_root) = &git_root {        lines.push(format!("Git root: {}", git_root.display()));        lines.push(format!("Git project: {}", file_name_string(git_root)));    }    if let Some(user_root) = &user_root {        lines.push(format!("User root: {}", user_root.display()));    }    if let Some(tree) = cwd_tree {        lines.push(String::new());        lines.push("Working directory tree:".to_string());        lines.extend(tree);    }    if let Some(tree) = git_root_tree {        lines.push(String::new());        lines.push("Git root tree:".to_string());        lines.extend(tree);    }    if let Some(tree) = user_root_tree {        lines.push(String::new());        lines.push("User root tree:".to_string());        lines.extend(tree);    }    Some(lines.join("\n"))}fn render_tree(root: &Path) -> Option<Vec<String>> {    if !root.is_dir() {        return None;    }    let mut lines = Vec::new();    collect_tree_lines(root, /*depth*/ 0, &mut lines);    (!lines.is_empty()).then_some(lines)}fn collect_tree_lines(dir: &Path, depth: usize, lines: &mut Vec<String>) {    if depth >= TREE_MAX_DEPTH {        return;    }    let entries = match read_sorted_entries(dir) {        Ok(entries) => entries,        Err(_) => return,    };    let total_entries = entries.len();    for entry in entries.into_iter().take(DIR_ENTRY_LIMIT) {        let Ok(file_type) = entry.file_type() else {            continue;        };        let name = file_name_string(&entry.path());        let indent = "  ".repeat(depth);        let suffix = if file_type.is_dir() { "/" } else { "" };        lines.push(format!("{indent}- {name}{suffix}"));        if file_type.is_dir() {            collect_tree_lines(&entry.path(), depth + 1, lines);        }    }    if total_entries > DIR_ENTRY_LIMIT {        lines.push(format!(            "{}- ... {} more entries",            "  ".repeat(depth),            total_entries - DIR_ENTRY_LIMIT        ));    }}fn read_sorted_entries(dir: &Path) -> io::Result<Vec<DirEntry>> {    let mut entries = std::fs::read_dir(dir)?        .filter_map(Result::ok)        .filter(|entry| !is_noisy_name(&entry.file_name()))        .collect::<Vec<_>>();    entries.sort_by(|left, right| {        let left_is_dir = left            .file_type()            .map(|file_type| file_type.is_dir())            .unwrap_or(false);        let right_is_dir = right            .file_type()            .map(|file_type| file_type.is_dir())            .unwrap_or(false);        (!left_is_dir, file_name_string(&left.path()))            .cmp(&(!right_is_dir, file_name_string(&right.path())))    });    Ok(entries)}fn is_noisy_name(name: &OsStr) -> bool {    let name = name.to_string_lossy();    name.starts_with('.') || NOISY_DIR_NAMES.iter().any(|noisy| *noisy == name)}fn format_section(title: &str, body: Option<String>, budget_tokens: usize) -> Option<String> {    let body = body?;    let body = body.trim();    if body.is_empty() {        return None;    }    let heading = format!("## {title}\n");    let body_budget = budget_tokens.saturating_sub(approx_token_count(&heading));    if body_budget == 0 {        return None;    }    let body = truncate_realtime_text_to_token_budget(body, body_budget);    if body.is_empty() {        return None;    }    Some(format!("{heading}{body}"))}fn format_startup_context_blob(body: &str) -> String {    format!("{STARTUP_CONTEXT_OPEN_TAG}\n{body}\n{STARTUP_CONTEXT_CLOSE_TAG}")}async fn format_thread_group(    current_group: &Path,    group: &Path,    entries: Vec<&StoredThread>,) -> Option<String> {    let latest = entries.first()?;    let group_label =        if let Ok(latest_cwd) = AbsolutePathBuf::from_absolute_path(latest.cwd.as_path()) {            if resolve_root_git_project_for_trust(LOCAL_FS.as_ref(), &latest_cwd)                .await                .is_some()            {                format!("### Git repo: {}", group.display())            } else {                format!("### Directory: {}", group.display())            }        } else {            format!("### Directory: {}", group.display())        };    let mut lines = vec![        group_label,        format!("Recent sessions: {}", entries.len()),        format!("Latest activity: {}", latest.updated_at.to_rfc3339()),    ];    if let Some(git_branch) = latest        .git_info        .as_ref()        .and_then(|git| git.branch.as_deref())        .filter(|git_branch| !git_branch.is_empty())    {        lines.push(format!("Latest branch: {git_branch}"));    }    lines.push(String::new());    lines.push("User asks:".to_string());    let mut seen = HashSet::new();    let max_asks = if group == current_group {        MAX_CURRENT_CWD_ASKS    } else {        MAX_OTHER_CWD_ASKS    };    for entry in entries {        let Some(first_user_message) = entry.first_user_message.as_deref() else {            continue;        };        let ask = first_user_message            .split_whitespace()            .collect::<Vec<_>>()            .join(" ");        let dedupe_key = format!("{}:{ask}", entry.cwd.display());        if ask.is_empty() || !seen.insert(dedupe_key) {            continue;        }        let ask = if ask.chars().count() > MAX_ASK_CHARS {            format!(                "{}...",                ask.chars()                    .take(MAX_ASK_CHARS.saturating_sub(3))                    .collect::<String>()            )        } else {            ask        };        lines.push(format!("- {}: {ask}", entry.cwd.display()));        if seen.len() == max_asks {            break;        }    }    (lines.len() > 5).then(|| lines.join("\n"))}fn file_name_string(path: &Path) -> String {    path.file_name()        .and_then(OsStr::to_str)        .map(str::to_owned)        .unwrap_or_else(|| path.display().to_string())}fn approx_token_count(text: &str) -> usize {    text.len().div_ceil(APPROX_BYTES_PER_TOKEN)}#[cfg(test)]#[path = "realtime_context_tests.rs"]mod tests;