lywsvip/warp
1
0
Fork 0
mirror of https://github.com/warpdotdev/warp.git synced 2026-05-06 23:32:51 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fix secret precedence in Oz CLI.

Browse Source
This commit is contained in:
Abhishek Pandya
2026-05-05 11:30:29 -05:00
parent c1ebde1fde
commit 0f07a67ee3
4 changed files with 383 additions and 82 deletions
Download Patch File Download Diff File Expand all files Collapse all files

185
app/src/ai/agent_sdk/driver.rs
View File

@@ -534,78 +534,7 @@ impl AgentDriver {
_ => None,
});
// Build environment variables from secrets for the terminal session.
// Do not override env vars that are already set to a non-empty value in the current
// process. This ensures that worker-injected credentials (e.g. harness auth secrets)
// and user-provided env vars (e.g. on self-hosted workers) take precedence over
// generic managed secrets.
let mut env_vars = HashMap::with_capacity(secrets.len() + 1);
for (name, secret) in &secrets {
let (env_name, env_value) = match secret {
ManagedSecretValue::RawValue { value } => (name.as_str(), value.as_str()),
ManagedSecretValue::AnthropicApiKey { api_key } => {
("ANTHROPIC_API_KEY", api_key.as_str())
}
ManagedSecretValue::AnthropicBedrockAccessKey {
aws_access_key_id,
aws_secret_access_key,
aws_session_token,
aws_region,
} => {
// Inject env vars needed for Claude Code Bedrock access key authentication.
// AWS_SESSION_TOKEN is only injected when the user provided one (i.e. for
// temporary/STS credentials).
let mut vars = vec![
("AWS_ACCESS_KEY_ID", aws_access_key_id.as_str()),
("AWS_SECRET_ACCESS_KEY", aws_secret_access_key.as_str()),
("CLAUDE_CODE_USE_BEDROCK", "1"),
("AWS_REGION", aws_region.as_str()),
];
if let Some(token) = aws_session_token.as_deref() {
vars.push(("AWS_SESSION_TOKEN", token));
}
for (env_name, env_value) in vars {
if std::env::var(env_name).is_ok_and(|v| !v.is_empty()) {
log::warn!(
"Skipping managed secret {env_name}: already set in environment"
);
continue;
}
env_vars.insert(OsString::from(env_name), OsString::from(env_value));
}
continue; // Skip the single-var insert below since we handled all vars inline.
}
ManagedSecretValue::AnthropicBedrockApiKey {
aws_bearer_token_bedrock,
aws_region,
} => {
// Inject all three env vars needed for Claude Code Bedrock authentication.
let vars = [
(
"AWS_BEARER_TOKEN_BEDROCK",
aws_bearer_token_bedrock.as_str(),
),
("CLAUDE_CODE_USE_BEDROCK", "1"),
("AWS_REGION", aws_region.as_str()),
];
for (env_name, env_value) in vars {
if std::env::var(env_name).is_ok_and(|v| !v.is_empty()) {
log::warn!(
"Skipping managed secret {env_name}: already set in environment"
);
continue;
}
env_vars.insert(OsString::from(env_name), OsString::from(env_value));
}
continue; // Skip the single-var insert below since we handled all vars inline.
}
};
if std::env::var(env_name).is_ok_and(|v| !v.is_empty()) {
log::warn!("Skipping managed secret {env_name}: already set in environment");
continue;
}
env_vars.insert(OsString::from(env_name), OsString::from(env_value));
}
let mut env_vars = build_secret_env_vars(&secrets);
// Inject cloud provider env vars.
cloud_provider::collect_env_vars(&cloud_providers, &mut env_vars)?;
@@ -2374,6 +2303,118 @@ impl AgentDriver {
}
}
/// Build the env-var map for the agent terminal session from managed secrets.
///
/// Precedence order:
/// 1. Worker-injected process env (already non-empty in `std::env`). Never overridden.
/// 2. Typed auth secrets (`AnthropicApiKey`, `AnthropicBedrock*`). Inserted atomically:
/// if any one env var for a typed secret is already worker-injected, the entire
/// secret is skipped.
/// 3. Generic `RawValue` secrets. Skipped on collision with either of the above.
fn build_secret_env_vars(
secrets: &HashMap<String, ManagedSecretValue>,
) -> HashMap<OsString, OsString> {
let mut env_vars = HashMap::with_capacity(secrets.len() + 1);
// Phase 1: Record which env-var names are claimed by typed auth secrets.
let typed_env_names = typed_secret_env_names(secrets);
// Phase 2: Insert typed auth secrets atomically.
for (name, secret) in secrets {
let entries = typed_secret_entries(secret);
if entries.is_empty() {
continue;
}
if let Some((conflict, _)) = entries
.iter()
.find(|(env_name, _)| std::env::var(env_name).is_ok_and(|v| !v.is_empty()))
{
log::warn!(
"Skipping auth secret '{name}' ({:?}): '{conflict}' is already set \
in the process environment",
secret.secret_type(),
);
continue;
}
for (env_name, env_value) in entries {
env_vars.insert(OsString::from(env_name), OsString::from(env_value));
}
}
// Phase 3: Insert generic RawValue secrets, skipping any that collide
// with worker-injected env vars or typed-secret-claimed names.
for (name, secret) in secrets {
let ManagedSecretValue::RawValue { value } = secret else {
continue;
};
let env_name = name.as_str();
if std::env::var(env_name).is_ok_and(|v| !v.is_empty()) {
log::warn!("Skipping managed secret {env_name}: already set in environment");
continue;
}
if typed_env_names.contains(env_name) {
log::warn!("Skipping generic secret '{env_name}': overridden by a typed auth secret");
continue;
}
env_vars.insert(OsString::from(env_name), OsString::from(value.as_str()));
}
env_vars
}
/// The env-var names that any typed auth secret in `secrets` will populate.
/// Used for phase-3 collision detection and by the suffix resolver.
fn typed_secret_env_names(secrets: &HashMap<String, ManagedSecretValue>) -> HashSet<&'static str> {
let mut names = HashSet::new();
for secret in secrets.values() {
for (env_name, _) in typed_secret_entries(secret) {
names.insert(env_name);
}
}
names
}
fn typed_secret_entries(secret: &ManagedSecretValue) -> Vec<(&'static str, &str)> {
match secret {
ManagedSecretValue::RawValue { .. } => vec![],
ManagedSecretValue::AnthropicApiKey { api_key } => {
vec![("ANTHROPIC_API_KEY", api_key.as_str())]
}
ManagedSecretValue::AnthropicBedrockApiKey {
aws_bearer_token_bedrock,
aws_region,
} => vec![
(
"AWS_BEARER_TOKEN_BEDROCK",
aws_bearer_token_bedrock.as_str(),
),
("CLAUDE_CODE_USE_BEDROCK", "1"),
("AWS_REGION", aws_region.as_str()),
],
ManagedSecretValue::AnthropicBedrockAccessKey {
aws_access_key_id,
aws_secret_access_key,
aws_session_token,
aws_region,
} => {
let mut entries = vec![
("AWS_ACCESS_KEY_ID", aws_access_key_id.as_str()),
("AWS_SECRET_ACCESS_KEY", aws_secret_access_key.as_str()),
("CLAUDE_CODE_USE_BEDROCK", "1"),
("AWS_REGION", aws_region.as_str()),
];
if let Some(token) = aws_session_token.as_deref() {
entries.push(("AWS_SESSION_TOKEN", token));
}
entries
}
}
}
impl Entity for AgentDriver {
type Event = ();
}

16
app/src/ai/agent_sdk/driver/harness/claude_code.rs
View File

@@ -657,21 +657,23 @@ struct ClaudeSettings {
fn resolve_anthropic_api_key_suffix(
secrets: &HashMap<String, ManagedSecretValue>,
) -> Option<String> {
// First, check for an AnthropicApiKey variant anywhere in the secrets map,
// since the secret name doesn't necessarily match the env var.
// 1. Worker-injected process env wins.
if let Ok(key) = std::env::var(ANTHROPIC_API_KEY_ENV) {
if !key.is_empty() {
return suffix_of(&key).map(str::to_owned);
}
}
// 2. Typed AnthropicApiKey secret. The secret name does not have to match the env var.
for secret in secrets.values() {
if let ManagedSecretValue::AnthropicApiKey { api_key } = secret {
return suffix_of(api_key).map(str::to_owned);
}
}
// Then check for a RawValue stored under the env var name.
// 3. Generic RawValue secret stored under the env var name.
if let Some(ManagedSecretValue::RawValue { value }) = secrets.get(ANTHROPIC_API_KEY_ENV) {
return suffix_of(value).map(str::to_owned);
}
// Fall back to the environment variable, which a user may have set separately in the env.
std::env::var(ANTHROPIC_API_KEY_ENV)
.ok()
.and_then(|k| suffix_of(&k).map(str::to_owned))
None
}
fn suffix_of(key: &str) -> Option<&str> {

70
app/src/ai/agent_sdk/driver/harness/claude_code_tests.rs
View File

@@ -582,7 +582,9 @@ fn prepare_claude_config_none_suffix_preserves_existing_responses() {
}
#[test]
#[serial_test::serial]
fn resolve_suffix_from_raw_value_secret() {
std::env::remove_var(ANTHROPIC_API_KEY_ENV);
let key = "sk-ant-api03-abcdefghij1234567890ABCDEFGHIJ1234567890abcdefghij1234567890QLWn-dUnuwQ-hIhDiAAA";
let secrets = HashMap::from([(
"ANTHROPIC_API_KEY".to_string(),
@@ -593,7 +595,9 @@ fn resolve_suffix_from_raw_value_secret() {
}
#[test]
#[serial_test::serial]
fn resolve_suffix_from_anthropic_api_key_secret() {
std::env::remove_var(ANTHROPIC_API_KEY_ENV);
let key = "sk-ant-api03-abcdefghij1234567890ABCDEFGHIJ1234567890abcdefghij1234567890QLWn-dUnuwQ-hIhDiAAA";
let secrets = HashMap::from([(
"ANTHROPIC_API_KEY".to_string(),
@@ -604,7 +608,9 @@ fn resolve_suffix_from_anthropic_api_key_secret() {
}
#[test]
#[serial_test::serial]
fn resolve_suffix_from_anthropic_api_key_with_different_secret_name() {
std::env::remove_var(ANTHROPIC_API_KEY_ENV);
let key = "sk-ant-api03-abcdefghij1234567890ABCDEFGHIJ1234567890abcdefghij1234567890QLWn-dUnuwQ-hIhDiAAA";
// Secret name doesn't match the env var, but the AnthropicApiKey variant
// should still be found by iterating all secrets.
@@ -617,7 +623,9 @@ fn resolve_suffix_from_anthropic_api_key_with_different_secret_name() {
}
#[test]
#[serial_test::serial]
fn resolve_suffix_prefers_anthropic_api_key_variant_over_raw_value() {
std::env::remove_var(ANTHROPIC_API_KEY_ENV);
let anthropic_key = "sk-ant-api03-AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA-anthropic-suffix";
let raw_key = "sk-ant-api03-BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB-raw-suffix";
let secrets = HashMap::from([
@@ -631,12 +639,14 @@ fn resolve_suffix_prefers_anthropic_api_key_variant_over_raw_value() {
),
]);
let suffix = resolve_anthropic_api_key_suffix(&secrets);
// AnthropicApiKey variant should be preferred.
// AnthropicApiKey variant should be preferred over RawValue.
assert_eq!(suffix.as_deref(), Some("AAA-anthropic-suffix"));
}
#[test]
#[serial_test::serial]
fn resolve_suffix_returns_none_for_short_key() {
std::env::remove_var(ANTHROPIC_API_KEY_ENV);
let secrets = HashMap::from([(
"ANTHROPIC_API_KEY".to_string(),
ManagedSecretValue::raw_value("short"),
@@ -725,7 +735,9 @@ fn prepare_local_wake_command_rehydrates_transcript_with_self_managed_listener()
}
#[test]
#[serial_test::serial]
fn resolve_suffix_returns_none_for_short_anthropic_api_key() {
std::env::remove_var(ANTHROPIC_API_KEY_ENV);
let secrets = HashMap::from([(
"ANTHROPIC_API_KEY".to_string(),
ManagedSecretValue::anthropic_api_key("short"),
@@ -733,6 +745,62 @@ fn resolve_suffix_returns_none_for_short_anthropic_api_key() {
assert_eq!(resolve_anthropic_api_key_suffix(&secrets), None);
}
// ── Suffix resolver precedence tests ────────────────────────────────────
#[test]
#[serial_test::serial]
fn suffix_uses_worker_injected_env_when_present() {
// Key must be >= 20 chars. Last 20: "WWWW-worker--suffix!"
let worker_key = "sk-ant-api03-WWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWW-worker-suffix!";
let typed_key = "sk-ant-api03-TTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTT-typed-suffix!";
std::env::set_var(ANTHROPIC_API_KEY_ENV, worker_key);
let secrets = HashMap::from([(
"my-auth".to_string(),
ManagedSecretValue::anthropic_api_key(typed_key),
)]);
let suffix = resolve_anthropic_api_key_suffix(&secrets);
let expected = &worker_key[worker_key.len() - 20..];
// Worker-injected process env wins over the typed secret.
assert_eq!(suffix.as_deref(), Some(expected));
std::env::remove_var(ANTHROPIC_API_KEY_ENV);
}
#[test]
#[serial_test::serial]
fn suffix_uses_typed_secret_when_no_worker_env() {
std::env::remove_var(ANTHROPIC_API_KEY_ENV);
let typed_key = "sk-ant-api03-TTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTT-typed-suffix!";
let raw_key = "sk-ant-api03-RRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRR-rawvl-suffix!";
let secrets = HashMap::from([
(
"my-auth".to_string(),
ManagedSecretValue::anthropic_api_key(typed_key),
),
(
"ANTHROPIC_API_KEY".to_string(),
ManagedSecretValue::raw_value(raw_key),
),
]);
let suffix = resolve_anthropic_api_key_suffix(&secrets);
let expected = &typed_key[typed_key.len() - 20..];
// Typed secret wins over RawValue when no worker env is present.
assert_eq!(suffix.as_deref(), Some(expected));
}
#[test]
#[serial_test::serial]
fn suffix_uses_raw_value_when_no_worker_env_or_typed() {
std::env::remove_var(ANTHROPIC_API_KEY_ENV);
let raw_key = "sk-ant-api03-RRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRR-rawvl-suffix!";
let secrets = HashMap::from([(
"ANTHROPIC_API_KEY".to_string(),
ManagedSecretValue::raw_value(raw_key),
)]);
let suffix = resolve_anthropic_api_key_suffix(&secrets);
let expected = &raw_key[raw_key.len() - 20..];
assert_eq!(suffix.as_deref(), Some(expected));
}
#[test]
fn prepare_claude_settings_creates_settings_file() {
let tmp = TempDir::new().unwrap();

194
app/src/ai/agent_sdk/driver_tests.rs
View File

@@ -1,4 +1,9 @@
use std::{ffi::OsString, sync::Arc, time::Duration};
use std::{
collections::HashMap,
ffi::OsString,
sync::Arc,
time::Duration,
};
use futures::channel::oneshot;
use warp_cli::agent::Harness;
@@ -9,7 +14,7 @@ use warp_cli::{
use warp_core::channel::ChannelState;
use super::{
IdleTimeoutSender, LEGACY_OZ_PARENT_LISTENER_MANAGED_EXTERNALLY_ENV,
build_secret_env_vars, IdleTimeoutSender, LEGACY_OZ_PARENT_LISTENER_MANAGED_EXTERNALLY_ENV,
LEGACY_OZ_PARENT_STATE_ROOT_ENV, OZ_MESSAGE_LISTENER_MANAGED_EXTERNALLY_ENV,
OZ_MESSAGE_LISTENER_STATE_ROOT_ENV,
};
@@ -19,6 +24,7 @@ use crate::ai::agent::{
};
use crate::ai::mcp::parsing::normalize_mcp_json;
use crate::ai::{agent_sdk::task_env_vars, ambient_agents::AmbientAgentTaskId};
use warp_managed_secrets::ManagedSecretValue;
#[test]
fn test_normalize_single_cli_server() {
@@ -386,3 +392,187 @@ fn json_format_input_omits_filepath_and_description_for_proto_upload_result() {
assert!(value.get("filepath").is_none());
assert!(value.get("description").is_none());
}
// ── build_secret_env_vars tests ──────────────────────────────────────────────
#[test]
#[serial_test::serial]
fn raw_value_only_writes_under_secret_name() {
std::env::remove_var("MY_SECRET");
let secrets = HashMap::from([(
"MY_SECRET".to_string(),
ManagedSecretValue::raw_value("s3cret"),
)]);
let env_vars = build_secret_env_vars(&secrets);
assert_eq!(
env_vars.get(&OsString::from("MY_SECRET")),
Some(&OsString::from("s3cret"))
);
assert_eq!(env_vars.len(), 1);
}
#[test]
#[serial_test::serial]
fn anthropic_api_key_writes_anthropic_env_var() {
std::env::remove_var("ANTHROPIC_API_KEY");
let secrets = HashMap::from([(
"my-custom-name".to_string(),
ManagedSecretValue::anthropic_api_key("sk-ant-test-key"),
)]);
let env_vars = build_secret_env_vars(&secrets);
assert_eq!(
env_vars.get(&OsString::from("ANTHROPIC_API_KEY")),
Some(&OsString::from("sk-ant-test-key"))
);
}
#[test]
#[serial_test::serial]
fn typed_secret_overrides_raw_value_with_same_env_name() {
std::env::remove_var("ANTHROPIC_API_KEY");
let typed_key = "sk-ant-typed-key-abcdef";
let raw_key = "sk-ant-raw-key-ghijkl";
let secrets = HashMap::from([
(
"my-auth".to_string(),
ManagedSecretValue::anthropic_api_key(typed_key),
),
(
"ANTHROPIC_API_KEY".to_string(),
ManagedSecretValue::raw_value(raw_key),
),
]);
// Run multiple times to defeat HashMap iteration order flakiness.
for _ in 0..20 {
let env_vars = build_secret_env_vars(&secrets);
assert_eq!(
env_vars.get(&OsString::from("ANTHROPIC_API_KEY")),
Some(&OsString::from(typed_key)),
"Typed secret must always override RawValue with the same env name"
);
}
}
#[test]
#[serial_test::serial]
fn bedrock_api_key_writes_all_bedrock_env_vars() {
std::env::remove_var("AWS_BEARER_TOKEN_BEDROCK");
std::env::remove_var("CLAUDE_CODE_USE_BEDROCK");
std::env::remove_var("AWS_REGION");
let secrets = HashMap::from([
(
"bedrock-secret".to_string(),
ManagedSecretValue::anthropic_bedrock_api_key("token-123", "us-west-2"),
),
(
"AWS_REGION".to_string(),
ManagedSecretValue::raw_value("eu-west-1"),
),
]);
let env_vars = build_secret_env_vars(&secrets);
assert_eq!(
env_vars.get(&OsString::from("AWS_BEARER_TOKEN_BEDROCK")),
Some(&OsString::from("token-123"))
);
assert_eq!(
env_vars.get(&OsString::from("CLAUDE_CODE_USE_BEDROCK")),
Some(&OsString::from("1"))
);
assert_eq!(
env_vars.get(&OsString::from("AWS_REGION")),
Some(&OsString::from("us-west-2")),
"Typed Bedrock secret should win over RawValue for AWS_REGION"
);
}
#[test]
#[serial_test::serial]
fn bedrock_access_key_writes_all_aws_env_vars() {
std::env::remove_var("AWS_ACCESS_KEY_ID");
std::env::remove_var("AWS_SECRET_ACCESS_KEY");
std::env::remove_var("AWS_SESSION_TOKEN");
std::env::remove_var("CLAUDE_CODE_USE_BEDROCK");
std::env::remove_var("AWS_REGION");
let secrets = HashMap::from([(
"bedrock-access".to_string(),
ManagedSecretValue::anthropic_bedrock_access_key(
"AKID",
"secret-key",
Some("session-tok".to_string()),
"ap-southeast-1",
),
)]);
let env_vars = build_secret_env_vars(&secrets);
assert_eq!(
env_vars.get(&OsString::from("AWS_ACCESS_KEY_ID")),
Some(&OsString::from("AKID"))
);
assert_eq!(
env_vars.get(&OsString::from("AWS_SECRET_ACCESS_KEY")),
Some(&OsString::from("secret-key"))
);
assert_eq!(
env_vars.get(&OsString::from("AWS_SESSION_TOKEN")),
Some(&OsString::from("session-tok"))
);
assert_eq!(
env_vars.get(&OsString::from("CLAUDE_CODE_USE_BEDROCK")),
Some(&OsString::from("1"))
);
assert_eq!(
env_vars.get(&OsString::from("AWS_REGION")),
Some(&OsString::from("ap-southeast-1"))
);
}
#[test]
#[serial_test::serial]
fn raw_value_skipped_when_process_env_already_set() {
std::env::set_var("WORKER_TOKEN", "injected-value");
let secrets = HashMap::from([(
"WORKER_TOKEN".to_string(),
ManagedSecretValue::raw_value("managed-value"),
)]);
let env_vars = build_secret_env_vars(&secrets);
// The worker-injected env var wins; env_vars should NOT contain it
// because the child inherits the process env directly.
assert!(!env_vars.contains_key(&OsString::from("WORKER_TOKEN")));
std::env::remove_var("WORKER_TOKEN");
}
#[test]
#[serial_test::serial]
fn worker_injected_env_wins_over_typed_secret() {
std::env::set_var("ANTHROPIC_API_KEY", "worker-key");
let secrets = HashMap::from([(
"my-auth".to_string(),
ManagedSecretValue::anthropic_api_key("managed-key"),
)]);
let env_vars = build_secret_env_vars(&secrets);
// The typed secret should be skipped entirely; the child inherits
// ANTHROPIC_API_KEY from the process env.
assert!(!env_vars.contains_key(&OsString::from("ANTHROPIC_API_KEY")));
std::env::remove_var("ANTHROPIC_API_KEY");
}
#[test]
#[serial_test::serial]
fn worker_injected_env_skips_entire_bedrock_secret() {
// Only AWS_REGION is worker-injected; the entire Bedrock secret should
// be atomically skipped — no partial insertion.
std::env::set_var("AWS_REGION", "us-east-1");
std::env::remove_var("AWS_BEARER_TOKEN_BEDROCK");
std::env::remove_var("CLAUDE_CODE_USE_BEDROCK");
let secrets = HashMap::from([(
"bedrock-secret".to_string(),
ManagedSecretValue::anthropic_bedrock_api_key("token-456", "eu-central-1"),
)]);
let env_vars = build_secret_env_vars(&secrets);
assert!(
!env_vars.contains_key(&OsString::from("AWS_BEARER_TOKEN_BEDROCK")),
"Entire Bedrock secret must be skipped when any field is worker-injected"
);
assert!(!env_vars.contains_key(&OsString::from("CLAUDE_CODE_USE_BEDROCK")));
assert!(!env_vars.contains_key(&OsString::from("AWS_REGION")));
std::env::remove_var("AWS_REGION");
}