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::ResponseItem; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::GitInfo; use codex_protocol::protocol::SandboxPolicy; 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(), rollout_path: Some(PathBuf::from("/tmp/rollout.jsonl")), forked_from_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, sandbox_policy: SandboxPolicy::new_read_only_policy(), 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, } } fn user_message(text: impl Into) -> ResponseItem { message("user", ContentItem::InputText { text: text.into() }) } fn assistant_message(text: impl Into) -> 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 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, "\nStartup context from Codex.\n## Current Thread\nhello\n" ); } #[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("\n")); assert!(wrapped.ends_with("\n")); 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()))); }