Expose thread instruction preload API

This commit is contained in:
Adam Perry
2026-06-04 20:53:00 -07:00
parent 878255aa6d
commit 5975615ab4
9 changed files with 276 additions and 148 deletions

View File

@@ -1,35 +1,20 @@
use codex_core::AgentsMdManager;
use codex_core::ThreadManager;
use codex_core::config::Config;
use codex_protocol::error::CodexErr;
use codex_core::load_thread_user_instructions as preload_thread_user_instructions;
use codex_protocol::error::Result as CodexResult;
use codex_protocol::protocol::TurnEnvironmentSelection;
/// Loads AGENTS.md instructions from the primary turn environment into
/// `config` before the app server passes it to the thread manager.
/// Preloads AGENTS.md instructions from the primary selected environment
/// before app-server passes the config to the thread manager.
pub(super) async fn load_thread_user_instructions(
thread_manager: &ThreadManager,
config: &mut Config,
environments: &[TurnEnvironmentSelection],
) -> CodexResult<()> {
let Some(primary_selection) = environments.first() else {
config.user_instructions = None;
return Ok(());
};
let environment = thread_manager
.environment_manager()
.get_environment(&primary_selection.environment_id)
.ok_or_else(|| {
CodexErr::InvalidRequest(format!(
"unknown turn environment id `{}`",
primary_selection.environment_id
))
})?;
let mut warnings = Vec::new();
let user_instructions = AgentsMdManager::new(config)
.load_user_instructions(environment.as_ref(), &mut warnings)
.await;
config.startup_warnings.extend(warnings);
config.user_instructions = user_instructions;
Ok(())
preload_thread_user_instructions(
config,
thread_manager.environment_manager().as_ref(),
environments,
)
.await
}

View File

@@ -24,7 +24,6 @@ pub use codex_config::types::TuiKeymap;
pub use codex_config::types::TuiNotificationSettings;
pub use codex_config::types::TuiPetAnchor;
pub use codex_config::types::UriBasedFileOpener;
pub use codex_core::AgentsMdManager;
pub use codex_core::CodexThread;
pub use codex_core::ForkSnapshot;
pub use codex_core::LoadedAgentsMd;
@@ -43,6 +42,7 @@ pub use codex_core::config::TerminalResizeReflowConfig;
pub use codex_core::config::ThreadStoreConfig;
pub use codex_core::config::find_codex_home;
pub use codex_core::init_state_db;
pub use codex_core::load_thread_user_instructions;
pub use codex_core::resolve_installation_id;
pub use codex_core::skills::SkillsManager;
pub use codex_core::thread_store_from_config;

View File

@@ -71,14 +71,17 @@ async fn apply_role_to_config_inner(
}
let preserve_current_provider = role_layer_toml.get("model_provider").is_none();
let preserve_current_service_tier = role_layer_toml.get("service_tier").is_none();
let user_instructions = config.user_instructions.clone();
*config = reload::build_next_config(
let mut next_config = reload::build_next_config(
config,
role_layer_toml,
preserve_current_provider,
preserve_current_service_tier,
)
.await?;
next_config.user_instructions = user_instructions;
*config = next_config;
Ok(())
}

View File

@@ -1,4 +1,5 @@
use super::*;
use crate::LoadedAgentsMd;
use crate::SkillsManager;
use crate::config::ConfigBuilder;
use crate::skills_load_input_from_config;
@@ -212,6 +213,35 @@ async fn apply_role_preserves_unspecified_keys() {
);
}
#[tokio::test]
async fn apply_role_preserves_preloaded_user_instructions() {
let (home, mut config) = test_config_with_cli_overrides(Vec::new()).await;
config.user_instructions = Some(LoadedAgentsMd::from_text_for_testing(
"preloaded instructions",
));
let expected = config.user_instructions.clone();
let role_path = write_role_config(
&home,
"instruction-role.toml",
"developer_instructions = \"Stay focused\"",
)
.await;
config.agent_roles.insert(
"custom".to_string(),
AgentRoleConfig {
description: None,
config_file: Some(role_path),
nickname_candidates: None,
},
);
apply_role_to_config(&mut config, Some("custom"))
.await
.expect("custom role should apply");
assert_eq!(config.user_instructions, expected);
}
#[tokio::test]
async fn apply_role_reports_explicit_service_tier() {
let (home, mut config) = test_config_with_cli_overrides(Vec::new()).await;

View File

@@ -22,9 +22,13 @@ use codex_config::default_project_root_markers;
use codex_config::merge_toml_values;
use codex_config::project_root_markers_from_config;
use codex_exec_server::Environment;
use codex_exec_server::EnvironmentManager;
use codex_exec_server::ExecutorFileSystem;
use codex_features::Feature;
use codex_prompts::HIERARCHICAL_AGENTS_MESSAGE;
use codex_protocol::error::CodexErr;
use codex_protocol::error::Result as CodexResult;
use codex_protocol::protocol::TurnEnvironmentSelection;
use codex_utils_absolute_path::AbsolutePathBuf;
use std::io;
use toml::Value as TomlValue;
@@ -39,6 +43,38 @@ pub const LOCAL_AGENTS_MD_FILENAME: &str = "AGENTS.override.md";
/// concatenated with the following separator.
const AGENTS_MD_SEPARATOR: &str = "\n\n--- project-doc ---\n\n";
/// Preloads AGENTS.md instructions for thread construction from the primary
/// selected environment.
///
/// Call this after the selected environments have been materialized and before
/// passing `config` to a thread start, resume, or fork operation. When no
/// environment is selected, this clears [`Config::user_instructions`].
pub async fn load_thread_user_instructions(
config: &mut Config,
environment_manager: &EnvironmentManager,
environments: &[TurnEnvironmentSelection],
) -> CodexResult<()> {
let Some(primary_selection) = environments.first() else {
config.user_instructions = None;
return Ok(());
};
let environment = environment_manager
.get_environment(&primary_selection.environment_id)
.ok_or_else(|| {
CodexErr::InvalidRequest(format!(
"unknown turn environment id `{}`",
primary_selection.environment_id
))
})?;
let mut warnings = Vec::new();
let user_instructions = AgentsMdManager::new(config)
.load_user_instructions(environment.as_ref(), &mut warnings)
.await;
config.startup_warnings.extend(warnings);
config.user_instructions = user_instructions;
Ok(())
}
/// Resolves AGENTS.md files into model-visible user instructions and source
/// paths.
pub struct AgentsMdManager<'a> {
@@ -102,10 +138,15 @@ impl<'a> AgentsMdManager<'a> {
) -> Option<LoadedAgentsMd> {
let agents_md_docs = self.read_agents_md(fs, startup_warnings).await;
let mut loaded = self.config.user_instructions.clone().unwrap_or_default();
let mut loaded = self
.config
.user_instructions
.as_ref()
.map(LoadedAgentsMd::user_instructions)
.unwrap_or_default();
match agents_md_docs {
Ok(Some(docs)) => loaded.entries.extend(docs.entries),
Ok(Some(docs)) => loaded.project_entries = docs.project_entries,
Ok(None) => {}
Err(e) => {
error!("error trying to find AGENTS.md docs: {e:#}");
@@ -113,9 +154,8 @@ impl<'a> AgentsMdManager<'a> {
};
if self.config.features.enabled(Feature::ChildAgentsMd) {
loaded.entries.push(InstructionEntry {
loaded.internal_entries.push(InternalInstructionEntry {
contents: HIERARCHICAL_AGENTS_MESSAGE.to_string(),
provenance: InstructionProvenance::Internal,
});
}
@@ -181,9 +221,9 @@ impl<'a> AgentsMdManager<'a> {
let text = String::from_utf8_lossy(&data).to_string();
if !text.trim().is_empty() {
loaded.entries.push(InstructionEntry {
loaded.project_entries.push(ProjectInstructionEntry {
contents: text,
provenance: InstructionProvenance::Project(p),
source: p,
});
remaining = remaining.saturating_sub(data.len() as u64);
}
@@ -312,8 +352,9 @@ impl<'a> AgentsMdManager<'a> {
/// guidance.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct LoadedAgentsMd {
/// Ordered instructions and their provenance.
entries: Vec<InstructionEntry>,
user_entries: Vec<UserInstructionEntry>,
project_entries: Vec<ProjectInstructionEntry>,
internal_entries: Vec<InternalInstructionEntry>,
}
impl LoadedAgentsMd {
@@ -323,10 +364,11 @@ impl LoadedAgentsMd {
return Self::default();
}
Self {
entries: vec![InstructionEntry {
user_entries: vec![UserInstructionEntry {
contents,
provenance: InstructionProvenance::User(path),
source: Some(path),
}],
..Default::default()
}
}
@@ -340,80 +382,93 @@ impl LoadedAgentsMd {
return Self::default();
}
Self {
entries: vec![InstructionEntry {
user_entries: vec![UserInstructionEntry {
contents,
provenance: InstructionProvenance::Internal,
source: None,
}],
..Default::default()
}
}
fn is_empty(&self) -> bool {
self.entries
self.user_entries
.iter()
.all(|entry| entry.contents.trim().is_empty())
&& self
.project_entries
.iter()
.all(|entry| entry.contents.trim().is_empty())
&& self
.internal_entries
.iter()
.all(|entry| entry.contents.trim().is_empty())
}
fn user_instructions(&self) -> Self {
Self {
user_entries: self.user_entries.clone(),
..Default::default()
}
}
/// Returns the concatenated model-visible instruction text.
pub fn text(&self) -> String {
let mut output = String::new();
let mut previous_provenance: Option<&InstructionProvenance> = None;
for entry in &self.entries {
if let Some(previous_provenance) = previous_provenance {
// The project-doc marker tells the model where workspace-scoped
// instructions begin, so it is only needed on the transition
// from user or internal instructions to project instructions.
let separator = match (previous_provenance, &entry.provenance) {
(
InstructionProvenance::User(_) | InstructionProvenance::Internal,
InstructionProvenance::Project(_),
) => AGENTS_MD_SEPARATOR,
_ => "\n\n",
};
output.push_str(separator);
for entry in &self.user_entries {
append_instruction(&mut output, &entry.contents, "\n\n");
}
let mut appended_project_entry = false;
for entry in &self.project_entries {
let separator = if appended_project_entry || output.is_empty() {
"\n\n"
} else {
AGENTS_MD_SEPARATOR
};
if append_instruction(&mut output, &entry.contents, separator) {
appended_project_entry = true;
}
output.push_str(&entry.contents);
previous_provenance = Some(&entry.provenance);
}
for entry in &self.internal_entries {
append_instruction(&mut output, &entry.contents, "\n\n");
}
output
}
/// Returns the AGENTS.md files that supplied instruction entries.
pub fn sources(&self) -> impl Iterator<Item = &AbsolutePathBuf> {
self.entries
self.user_entries
.iter()
.filter_map(|entry| entry.provenance.path())
.filter_map(|entry| entry.source.as_ref())
.chain(self.project_entries.iter().map(|entry| &entry.source))
}
}
/// One model-visible instruction and its provenance.
#[derive(Clone, Debug, PartialEq, Eq)]
struct InstructionEntry {
/// Model-visible instruction text.
struct UserInstructionEntry {
contents: String,
/// Origin of the instruction.
provenance: InstructionProvenance,
source: Option<AbsolutePathBuf>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
enum InstructionProvenance {
/// User-level instructions, normally loaded from CODEX_HOME.
User(AbsolutePathBuf),
/// Workspace instructions discovered from project AGENTS.md files.
Project(AbsolutePathBuf),
/// Instructions without a file source, including internally defined guidance.
Internal,
struct ProjectInstructionEntry {
contents: String,
source: AbsolutePathBuf,
}
impl InstructionProvenance {
fn path(&self) -> Option<&AbsolutePathBuf> {
match self {
Self::User(path) | Self::Project(path) => Some(path),
Self::Internal => None,
}
#[derive(Clone, Debug, PartialEq, Eq)]
struct InternalInstructionEntry {
contents: String,
}
fn append_instruction(output: &mut String, contents: &str, separator: &str) -> bool {
if contents.trim().is_empty() {
return false;
}
if !output.is_empty() {
output.push_str(separator);
}
output.push_str(contents);
true
}
fn warn_invalid_utf8(

View File

@@ -1,7 +1,9 @@
use super::*;
use crate::config::ConfigBuilder;
use codex_exec_server::EnvironmentManager;
use codex_exec_server::LOCAL_FS;
use codex_features::Feature;
use codex_protocol::protocol::TurnEnvironmentSelection;
use codex_utils_absolute_path::AbsolutePathBuf;
use core_test_support::PathBufExt;
use core_test_support::TempDirExt;
@@ -153,16 +155,16 @@ fn empty_loaded_instructions_are_empty() {
#[test]
fn loaded_instructions_with_only_empty_or_whitespace_entries_are_empty() {
let empty = LoadedAgentsMd {
entries: vec![InstructionEntry {
internal_entries: vec![InternalInstructionEntry {
contents: String::new(),
provenance: InstructionProvenance::Internal,
}],
..Default::default()
};
let whitespace = LoadedAgentsMd {
entries: vec![InstructionEntry {
internal_entries: vec![InternalInstructionEntry {
contents: " \n\t".to_string(),
provenance: InstructionProvenance::Internal,
}],
..Default::default()
};
assert!(empty.is_empty());
@@ -353,6 +355,79 @@ async fn keeps_existing_instructions_when_doc_missing() {
assert_eq!(res, Some(INSTRUCTIONS.to_string()));
}
#[tokio::test]
async fn resolving_again_replaces_project_instructions() {
let tmp = tempfile::tempdir().expect("tempdir");
let first_project = tmp.path().join("first");
let second_project = tmp.path().join("second");
fs::create_dir_all(&first_project).unwrap();
fs::create_dir_all(&second_project).unwrap();
fs::write(first_project.join("AGENTS.md"), "first project").unwrap();
fs::write(second_project.join("AGENTS.md"), "second project").unwrap();
let mut config = make_config(&tmp, /*limit*/ 4096, Some("user instructions")).await;
config.cwd = first_project.abs();
let mut warnings = Vec::new();
config.user_instructions = AgentsMdManager::new(&config)
.user_instructions_with_fs(LOCAL_FS.as_ref(), &mut warnings)
.await;
config.cwd = second_project.abs();
let loaded = AgentsMdManager::new(&config)
.user_instructions_with_fs(LOCAL_FS.as_ref(), &mut warnings)
.await
.expect("instructions expected");
let user_agents = config.codex_home.join(DEFAULT_AGENTS_MD_FILENAME);
let second_agents = AbsolutePathBuf::try_from(
dunce::canonicalize(second_project.join("AGENTS.md"))
.expect("canonical second project doc path"),
)
.expect("absolute second project doc path");
assert_eq!(
loaded,
LoadedAgentsMd {
user_entries: vec![UserInstructionEntry {
contents: "user instructions".to_string(),
source: Some(user_agents),
}],
project_entries: vec![ProjectInstructionEntry {
contents: "second project".to_string(),
source: second_agents,
}],
internal_entries: Vec::new(),
}
);
}
#[tokio::test]
async fn load_thread_user_instructions_uses_primary_environment() {
let tmp = tempfile::tempdir().expect("tempdir");
fs::write(tmp.path().join("AGENTS.md"), "project instructions").unwrap();
let mut config = make_config(&tmp, /*limit*/ 4096, Some("user instructions")).await;
let environment_manager = EnvironmentManager::default_for_tests();
let environment_id = environment_manager
.default_environment_id()
.expect("default environment")
.to_string();
let environments = vec![TurnEnvironmentSelection {
environment_id,
cwd: config.cwd.clone(),
}];
load_thread_user_instructions(&mut config, &environment_manager, &environments)
.await
.expect("load thread instructions");
assert_eq!(
config
.user_instructions
.expect("instructions expected")
.text(),
format!("user instructions{AGENTS_MD_SEPARATOR}project instructions")
);
}
/// When both the repository root and the working directory contain
/// AGENTS.md files, their contents are concatenated from root to cwd.
#[tokio::test]
@@ -385,16 +460,17 @@ async fn concatenates_root_and_cwd_docs() {
let root_agents = repo.path().join("AGENTS.md").abs();
let crate_agents = cfg.cwd.join("AGENTS.md");
let expected = LoadedAgentsMd {
entries: vec![
InstructionEntry {
project_entries: vec![
ProjectInstructionEntry {
contents: "root doc".to_string(),
provenance: InstructionProvenance::Project(root_agents.clone()),
source: root_agents.clone(),
},
InstructionEntry {
ProjectInstructionEntry {
contents: "crate doc".to_string(),
provenance: InstructionProvenance::Project(crate_agents.clone()),
source: crate_agents.clone(),
},
],
..Default::default()
};
assert_eq!(loaded, expected);
@@ -468,16 +544,14 @@ async fn child_agents_message_after_global_instructions_uses_plain_separator() {
.expect("instructions expected");
let global_agents = cfg.codex_home.join(DEFAULT_AGENTS_MD_FILENAME);
let expected = LoadedAgentsMd {
entries: vec![
InstructionEntry {
contents: "global doc".to_string(),
provenance: InstructionProvenance::User(global_agents),
},
InstructionEntry {
contents: HIERARCHICAL_AGENTS_MESSAGE.to_string(),
provenance: InstructionProvenance::Internal,
},
],
user_entries: vec![UserInstructionEntry {
contents: "global doc".to_string(),
source: Some(global_agents),
}],
internal_entries: vec![InternalInstructionEntry {
contents: HIERARCHICAL_AGENTS_MESSAGE.to_string(),
}],
..Default::default()
};
assert_eq!(loaded, expected);
@@ -505,16 +579,15 @@ async fn instruction_sources_include_global_before_agents_md_docs() {
let project_agents = cfg.cwd.join("AGENTS.md");
let expected = LoadedAgentsMd {
entries: vec![
InstructionEntry {
contents: "global doc".to_string(),
provenance: InstructionProvenance::User(global_agents.clone()),
},
InstructionEntry {
contents: "project doc".to_string(),
provenance: InstructionProvenance::Project(project_agents.clone()),
},
],
user_entries: vec![UserInstructionEntry {
contents: "global doc".to_string(),
source: Some(global_agents.clone()),
}],
project_entries: vec![ProjectInstructionEntry {
contents: "project doc".to_string(),
source: project_agents.clone(),
}],
internal_entries: Vec::new(),
};
assert_eq!(loaded, expected);
assert_eq!(
@@ -546,20 +619,17 @@ async fn child_agents_message_after_project_docs_is_not_an_instruction_source()
let project_agents = cfg.cwd.join("AGENTS.md");
let expected = LoadedAgentsMd {
entries: vec![
InstructionEntry {
contents: "global doc".to_string(),
provenance: InstructionProvenance::User(global_agents.clone()),
},
InstructionEntry {
contents: "project doc".to_string(),
provenance: InstructionProvenance::Project(project_agents.clone()),
},
InstructionEntry {
contents: HIERARCHICAL_AGENTS_MESSAGE.to_string(),
provenance: InstructionProvenance::Internal,
},
],
user_entries: vec![UserInstructionEntry {
contents: "global doc".to_string(),
source: Some(global_agents.clone()),
}],
project_entries: vec![ProjectInstructionEntry {
contents: "project doc".to_string(),
source: project_agents.clone(),
}],
internal_entries: vec![InternalInstructionEntry {
contents: HIERARCHICAL_AGENTS_MESSAGE.to_string(),
}],
};
assert_eq!(loaded, expected);
assert_eq!(

View File

@@ -131,6 +131,7 @@ pub use agents_md::AgentsMdManager;
pub use agents_md::DEFAULT_AGENTS_MD_FILENAME;
pub use agents_md::LOCAL_AGENTS_MD_FILENAME;
pub use agents_md::LoadedAgentsMd;
pub use agents_md::load_thread_user_instructions;
mod rollout;
pub(crate) mod safety;
mod session_rollout_init_error;

View File

@@ -1,4 +1,3 @@
use codex_core::AgentsMdManager;
use codex_core::CodexThread;
use codex_core::ModelClient;
use codex_core::NewThread;
@@ -8,6 +7,7 @@ use codex_core::StartThreadOptions;
use codex_core::ThreadManager;
use codex_core::config::Config;
use codex_core::content_items_to_text;
use codex_core::load_thread_user_instructions;
use codex_core::resolve_installation_id;
use codex_features::Feature;
use codex_login::AuthManager;
@@ -237,20 +237,12 @@ impl MemoryStartupContext {
let environments = self
.thread_manager
.default_environment_selections(&config.cwd);
let mut warnings = Vec::new();
config.user_instructions = match environments.first().and_then(|selection| {
self.thread_manager
.environment_manager()
.get_environment(&selection.environment_id)
}) {
Some(environment) => {
AgentsMdManager::new(&config)
.load_user_instructions(environment.as_ref(), &mut warnings)
.await
}
None => None,
};
config.startup_warnings.extend(warnings);
load_thread_user_instructions(
&mut config,
self.thread_manager.environment_manager().as_ref(),
&environments,
)
.await?;
let NewThread {
thread_id, thread, ..
} = self

View File

@@ -9,7 +9,6 @@ use anyhow::Context;
use anyhow::bail;
use clap::Parser;
use codex_core_api::AbsolutePathBuf;
use codex_core_api::AgentsMdManager;
use codex_core_api::AltScreenMode;
use codex_core_api::ApprovalsReviewer;
use codex_core_api::Arg0DispatchPaths;
@@ -59,6 +58,7 @@ use codex_core_api::empty_extension_registry;
use codex_core_api::find_codex_home;
use codex_core_api::init_state_db;
use codex_core_api::item_event_to_server_notification;
use codex_core_api::load_thread_user_instructions;
use codex_core_api::resolve_installation_id;
use codex_core_api::set_default_originator;
use codex_core_api::thread_store_from_config;
@@ -121,21 +121,11 @@ async fn run_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> {
.await?,
);
let installation_id = resolve_installation_id(&config.codex_home).await?;
let mut warnings = Vec::new();
config.user_instructions = match environment_manager.default_environment() {
Some(environment) => {
AgentsMdManager::new(&config)
.load_user_instructions(environment.as_ref(), &mut warnings)
.await
}
None => None,
};
config.startup_warnings.extend(warnings);
let thread_manager = ThreadManager::new(
&config,
auth_manager,
SessionSource::Exec,
environment_manager,
Arc::clone(&environment_manager),
empty_extension_registry(),
/*analytics_events_client*/ None,
Arc::clone(&thread_store),
@@ -143,6 +133,8 @@ async fn run_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> {
installation_id,
/*attestation_provider*/ None,
);
let environments = thread_manager.default_environment_selections(&config.cwd);
load_thread_user_instructions(&mut config, environment_manager.as_ref(), &environments).await?;
let NewThread {
thread_id, thread, ..