use super::CURRENT_THREAD_SECTION_TOKEN_BUDGET;use super::NOTES_SECTION_TOKEN_BUDGET;use super::RECENT_WORK_SECTION_TOKEN_BUDGET;use super::STARTUP_CONTEXT_HEADER;use super::WORKSPACE_SECTION_TOKEN_BUDGET;use super::build_current_thread_section;use super::build_recent_work_section;use super::build_workspace_section_with_user_root;use super::format_section;use super::format_startup_context_blob;use chrono::TimeZone;use chrono::Utc;use codex_git_utils::GitSha;use codex_protocol::ThreadId;use codex_protocol::models::ContentItem;use codex_protocol::models::PermissionProfile;use codex_protocol::models::ResponseItem;use codex_protocol::protocol::AskForApproval;use codex_protocol::protocol::GitInfo;use codex_protocol::protocol::SessionSource;use codex_thread_store::StoredThread;use core_test_support::PathBufExt;use core_test_support::PathExt;use pretty_assertions::assert_eq;use std::fs;use std::path::PathBuf;use std::process::Command;use tempfile::TempDir;fn stored_thread(cwd: &str, title: &str, first_user_message: &str) -> StoredThread { StoredThread { thread_id: ThreadId::new(), extra_config: None, rollout_path: Some(PathBuf::from("/tmp/rollout.jsonl")), forked_from_id: None, parent_thread_id: None, preview: first_user_message.to_string(), name: (!title.is_empty()).then(|| title.to_string()), model_provider: "test-provider".to_string(), model: Some("gpt-5.2".to_string()), reasoning_effort: None, created_at: Utc .timestamp_opt(1_709_251_100, 0) .single() .expect("valid timestamp"), updated_at: Utc .timestamp_opt(1_709_251_200, 0) .single() .expect("valid timestamp"), archived_at: None, cwd: PathBuf::from(cwd), cli_version: "test".to_string(), source: SessionSource::Cli, thread_source: None, agent_nickname: None, agent_role: None, agent_path: None, git_info: Some(GitInfo { commit_hash: Some(GitSha::new("abcdef")), branch: Some("main".to_string()), repository_url: None, }), approval_mode: AskForApproval::Never, permission_profile: PermissionProfile::read_only(), token_usage: None, first_user_message: Some(first_user_message.to_string()), history: None, }}fn message(role: &str, content: ContentItem) -> ResponseItem { ResponseItem::Message { id: None, role: role.to_string(), content: vec![content], phase: None, metadata: 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 turnUser:user turn 4Assistant:assistant turn 4### Previous turn 1User:user turn 3Assistant:assistant turn 3### Previous turn 2User:user turn 2Assistant:assistant turn 2### Previous turn 3User:user turn 1Assistant: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 startup_context_blob_is_wrapped_in_tags_without_final_truncation() { let body = "Startup context from Codex.\n## Current Thread\nhello"; let wrapped = format_startup_context_blob(body); assert_eq!( wrapped, "<startup_context>\nStartup context from Codex.\n## Current Thread\nhello\n</startup_context>" );}#[test]fn fixed_section_budgets_apply_per_section_without_total_blob_truncation() { let body = [ STARTUP_CONTEXT_HEADER.to_string(), format_section( "Current Thread", Some("current thread ".repeat(2_000)), CURRENT_THREAD_SECTION_TOKEN_BUDGET, ) .expect("current thread section"), format_section( "Recent Work", Some("recent work ".repeat(3_000)), RECENT_WORK_SECTION_TOKEN_BUDGET, ) .expect("recent work section"), format_section( "Machine / Workspace Map", Some("workspace map ".repeat(2_500)), WORKSPACE_SECTION_TOKEN_BUDGET, ) .expect("workspace section"), format_section( "Notes", Some("notes ".repeat(500)), NOTES_SECTION_TOKEN_BUDGET, ) .expect("notes section"), ] .join("\n\n"); let wrapped = format_startup_context_blob(&body); assert!(wrapped.starts_with("<startup_context>\n")); assert!(wrapped.ends_with("\n</startup_context>")); assert!(wrapped.contains("tokens truncated")); assert!(wrapped.contains("## Current Thread")); assert!(wrapped.contains("## Recent Work")); assert!(wrapped.contains("## Machine / Workspace Map")); assert!(wrapped.contains("## Notes"));}#[tokio::test]async fn workspace_section_requires_meaningful_structure() { let cwd = TempDir::new().expect("tempdir"); assert_eq!( build_workspace_section_with_user_root(&cwd.path().abs(), /*user_root*/ None).await, None );}#[tokio::test]async fn workspace_section_includes_tree_when_entries_exist() { let cwd = TempDir::new().expect("tempdir"); fs::create_dir(cwd.path().join("docs")).expect("create docs dir"); fs::write(cwd.path().join("README.md"), "hello").expect("write readme"); let section = build_workspace_section_with_user_root(&cwd.path().abs(), /*user_root*/ None) .await .expect("workspace section"); assert!(section.contains("Working directory tree:")); assert!(section.contains("- docs/")); assert!(section.contains("- README.md"));}#[tokio::test]async fn workspace_section_includes_user_root_tree_when_distinct() { let root = TempDir::new().expect("tempdir"); let cwd = root.path().join("cwd"); let git_root = root.path().join("git"); let user_root = root.path().join("home"); fs::create_dir_all(cwd.join("docs")).expect("create cwd docs dir"); fs::write(cwd.join("README.md"), "hello").expect("write cwd readme"); fs::create_dir_all(git_root.join(".git")).expect("create git dir"); fs::write(git_root.join("Cargo.toml"), "[workspace]").expect("write git root marker"); fs::create_dir_all(user_root.join("code")).expect("create user root child"); fs::write(user_root.join(".zshrc"), "export TEST=1").expect("write home file"); let section = build_workspace_section_with_user_root(&cwd.abs(), Some(user_root)) .await .expect("workspace section"); assert!(section.contains("User root tree:")); assert!(section.contains("- code/")); assert!(!section.contains("- .zshrc"));}#[tokio::test]async fn recent_work_section_groups_threads_by_cwd() { let root = TempDir::new().expect("tempdir"); let repo = root.path().join("repo"); let workspace_a = repo.join("workspace-a"); let workspace_b = repo.join("workspace-b"); let outside = root.path().join("outside"); fs::create_dir(&repo).expect("create repo dir"); Command::new("git") .env("GIT_CONFIG_GLOBAL", "/dev/null") .env("GIT_CONFIG_NOSYSTEM", "1") .args(["init"]) .current_dir(&repo) .output() .expect("git init"); fs::create_dir_all(&workspace_a).expect("create workspace a"); fs::create_dir_all(&workspace_b).expect("create workspace b"); fs::create_dir_all(&outside).expect("create outside dir"); let recent_threads = vec![ stored_thread( workspace_a.to_string_lossy().as_ref(), "Investigate realtime startup context", "Log the startup context before sending it", ), stored_thread( workspace_b.to_string_lossy().as_ref(), "Trim websocket startup payload", "Remove memories from the realtime startup context", ), stored_thread(outside.to_string_lossy().as_ref(), "", "Inspect flaky test"), ]; let current_cwd = workspace_a; let repo = repo.abs(); let section = build_recent_work_section(¤t_cwd.abs(), &recent_threads) .await .expect("recent work section"); assert!(section.contains(&format!("### Git repo: {}", repo.display()))); assert!(section.contains("Recent sessions: 2")); assert!(section.contains("User asks:")); assert!(section.contains(&format!( "- {}: Log the startup context before sending it", current_cwd.display() ))); assert!(section.contains(&format!("### Directory: {}", outside.display()))); assert!(section.contains(&format!("- {}: Inspect flaky test", outside.display())));}