mirror of
https://github.com/openai/codex.git
synced 2026-05-28 06:55:01 +00:00
137 lines
4.8 KiB
Rust
137 lines
4.8 KiB
Rust
use codex_config::types::EnvironmentVariablePattern;
|
||
use codex_config::types::ShellEnvironmentPolicy;
|
||
use codex_config::types::ShellEnvironmentPolicyInherit;
|
||
use codex_protocol::ThreadId;
|
||
use std::collections::HashMap;
|
||
use std::collections::HashSet;
|
||
|
||
pub const CODEX_THREAD_ID_ENV_VAR: &str = "CODEX_THREAD_ID";
|
||
|
||
/// Construct an environment map based on the rules in the specified policy. The
|
||
/// resulting map can be passed directly to `Command::envs()` after calling
|
||
/// `env_clear()` to ensure no unintended variables are leaked to the spawned
|
||
/// process.
|
||
///
|
||
/// The derivation follows the algorithm documented in the struct-level comment
|
||
/// for [`ShellEnvironmentPolicy`].
|
||
///
|
||
/// `CODEX_THREAD_ID` is injected when a thread id is provided, even when
|
||
/// `include_only` is set.
|
||
pub fn create_env(
|
||
policy: &ShellEnvironmentPolicy,
|
||
thread_id: Option<ThreadId>,
|
||
) -> HashMap<String, String> {
|
||
create_env_from_vars(std::env::vars(), policy, thread_id)
|
||
}
|
||
|
||
fn create_env_from_vars<I>(
|
||
vars: I,
|
||
policy: &ShellEnvironmentPolicy,
|
||
thread_id: Option<ThreadId>,
|
||
) -> HashMap<String, String>
|
||
where
|
||
I: IntoIterator<Item = (String, String)>,
|
||
{
|
||
let mut env_map = populate_env(vars, policy, thread_id);
|
||
|
||
if cfg!(target_os = "windows") {
|
||
// This is a workaround to address the failures we are seeing in the
|
||
// following tests when run via Bazel on Windows:
|
||
//
|
||
// ```
|
||
// suite::shell_command::unicode_output::with_login
|
||
// suite::shell_command::unicode_output::without_login
|
||
// ```
|
||
//
|
||
// Currently, we can only reproduce these failures in CI, which makes
|
||
// iteration times long, so we include this quick fix for now to unblock
|
||
// getting the Windows Bazel build running.
|
||
if !env_map.keys().any(|k| k.eq_ignore_ascii_case("PATHEXT")) {
|
||
env_map.insert("PATHEXT".to_string(), ".COM;.EXE;.BAT;.CMD".to_string());
|
||
}
|
||
}
|
||
env_map
|
||
}
|
||
|
||
const COMMON_CORE_VARS: &[&str] = &["PATH", "SHELL", "TMPDIR", "TEMP", "TMP"];
|
||
|
||
#[cfg(target_os = "windows")]
|
||
const PLATFORM_CORE_VARS: &[&str] = &["PATHEXT", "USERNAME", "USERPROFILE"];
|
||
|
||
#[cfg(unix)]
|
||
const PLATFORM_CORE_VARS: &[&str] = &["HOME", "LANG", "LC_ALL", "LC_CTYPE", "LOGNAME", "USER"];
|
||
|
||
fn populate_env<I>(
|
||
vars: I,
|
||
policy: &ShellEnvironmentPolicy,
|
||
thread_id: Option<ThreadId>,
|
||
) -> HashMap<String, String>
|
||
where
|
||
I: IntoIterator<Item = (String, String)>,
|
||
{
|
||
// Step 1 – determine the starting set of variables based on the
|
||
// `inherit` strategy.
|
||
let mut env_map: HashMap<String, String> = match policy.inherit {
|
||
ShellEnvironmentPolicyInherit::All => vars.into_iter().collect(),
|
||
ShellEnvironmentPolicyInherit::None => HashMap::new(),
|
||
ShellEnvironmentPolicyInherit::Core => {
|
||
let core_vars: HashSet<&str> = COMMON_CORE_VARS
|
||
.iter()
|
||
.copied()
|
||
.chain(PLATFORM_CORE_VARS.iter().copied())
|
||
.collect();
|
||
let is_core_var = |name: &str| {
|
||
if cfg!(target_os = "windows") {
|
||
core_vars
|
||
.iter()
|
||
.any(|allowed| allowed.eq_ignore_ascii_case(name))
|
||
} else {
|
||
core_vars.contains(name)
|
||
}
|
||
};
|
||
vars.into_iter().filter(|(k, _)| is_core_var(k)).collect()
|
||
}
|
||
};
|
||
|
||
// Internal helper – does `name` match **any** pattern in `patterns`?
|
||
let matches_any = |name: &str, patterns: &[EnvironmentVariablePattern]| -> bool {
|
||
patterns.iter().any(|pattern| pattern.matches(name))
|
||
};
|
||
|
||
// Step 2 – Apply the default exclude if not disabled.
|
||
if !policy.ignore_default_excludes {
|
||
let default_excludes = vec![
|
||
EnvironmentVariablePattern::new_case_insensitive("*KEY*"),
|
||
EnvironmentVariablePattern::new_case_insensitive("*SECRET*"),
|
||
EnvironmentVariablePattern::new_case_insensitive("*TOKEN*"),
|
||
];
|
||
env_map.retain(|k, _| !matches_any(k, &default_excludes));
|
||
}
|
||
|
||
// Step 3 – Apply custom excludes.
|
||
if !policy.exclude.is_empty() {
|
||
env_map.retain(|k, _| !matches_any(k, &policy.exclude));
|
||
}
|
||
|
||
// Step 4 – Apply user-provided overrides.
|
||
for (key, val) in &policy.r#set {
|
||
env_map.insert(key.clone(), val.clone());
|
||
}
|
||
|
||
// Step 5 – If include_only is non-empty, keep *only* the matching vars.
|
||
if !policy.include_only.is_empty() {
|
||
env_map.retain(|k, _| matches_any(k, &policy.include_only));
|
||
}
|
||
|
||
// Step 6 – Populate the thread ID environment variable when provided.
|
||
if let Some(thread_id) = thread_id {
|
||
env_map.insert(CODEX_THREAD_ID_ENV_VAR.to_string(), thread_id.to_string());
|
||
}
|
||
|
||
env_map
|
||
}
|
||
|
||
#[cfg(test)]
|
||
#[path = "exec_env_tests.rs"]
|
||
mod tests;
|