Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 30 additions & 0 deletions crates/cli/src/agents/claude/launch.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,22 @@ use crate::process::{PreparedAgentLaunch, insert_after_host};
pub(crate) fn prepare(
launch: &mut PreparedAgentLaunch,
gateway_url: &str,
proxy_credential: &crate::provider_auth::TransparentProxyCredential,
dry_run: bool,
) -> Result<(), CliError> {
let proxy_header = format!(
"{}: {}",
crate::provider_auth::TRANSPARENT_PROXY_CREDENTIAL_HEADER,
proxy_credential.expose()
);
let custom_headers = std::env::var("ANTHROPIC_CUSTOM_HEADERS")
.ok()
.filter(|value| !value.trim().is_empty())
.map_or_else(
|| proxy_header.clone(),
|value| replace_custom_header(&value, &proxy_header),
);
launch.set_secret_env("ANTHROPIC_CUSTOM_HEADERS", custom_headers);
Comment thread
coderabbitai[bot] marked this conversation as resolved.
if dry_run {
insert_after_host(
&mut launch.argv,
Expand Down Expand Up @@ -80,6 +94,22 @@ pub(crate) fn prepare(
Ok(())
}

fn replace_custom_header(existing: &str, replacement: &str) -> String {
let replacement_name = replacement
.split_once(':')
.map_or(replacement, |(name, _)| name)
.trim();
existing
.lines()
.filter(|line| {
line.split_once(':')
.is_none_or(|(name, _)| !name.trim().eq_ignore_ascii_case(replacement_name))
})
.chain(std::iter::once(replacement))
.collect::<Vec<_>>()
.join("\n")
}

pub(crate) fn settings_overlay(
argv: &[String],
host_index: usize,
Expand Down
6 changes: 4 additions & 2 deletions crates/cli/src/agents/codex/launch.rs
Original file line number Diff line number Diff line change
Expand Up @@ -198,8 +198,10 @@ fn canonical_json(value: Value) -> Value {

fn gateway_provider_config(gateway_url: &str) -> String {
format!(
"model_providers.nemo-relay-openai={{name=\"NeMo Relay OpenAI\",base_url={},wire_api=\"responses\",requires_openai_auth=true,supports_websockets=false}}",
toml_string(gateway_url)
"model_providers.nemo-relay-openai={{name=\"NeMo Relay OpenAI\",base_url={},wire_api=\"responses\",requires_openai_auth=true,supports_websockets=false,env_http_headers={{{}={}}}}}",
toml_string(gateway_url),
toml_string(crate::provider_auth::TRANSPARENT_PROXY_CREDENTIAL_HEADER),
toml_string(crate::provider_auth::TRANSPARENT_PROXY_CREDENTIAL_ENV),
)
}

Expand Down
7 changes: 7 additions & 0 deletions crates/cli/src/agents/hermes/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,13 @@ pub(crate) fn transparent_config(
"base_url".into(),
Value::String(format!("{}/v1", gateway_url.trim_end_matches('/'))),
);
model.insert(
"api_key".into(),
Value::String(format!(
"${{{}}}",
crate::provider_auth::TRANSPARENT_PROXY_CREDENTIAL_ENV
)),
);
object.insert("model".into(), Value::Object(model));
serde_yaml::to_string(&root).map_err(|error| CliError::Install(error.to_string()))
}
Expand Down
9 changes: 8 additions & 1 deletion crates/cli/src/agents/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -383,11 +383,18 @@ pub(crate) fn prepare_launch(
launch: &mut crate::process::PreparedAgentLaunch,
gateway_url: &str,
resolved: &crate::configuration::ResolvedConfig,
proxy_credential: &crate::provider_auth::TransparentProxyCredential,
dry_run: bool,
) -> Result<(), crate::error::CliError> {
launch.set_secret_env(
crate::provider_auth::TRANSPARENT_PROXY_CREDENTIAL_ENV,
proxy_credential.expose(),
);
match agent {
CodingAgent::Codex => codex::launch::prepare(launch, gateway_url),
CodingAgent::ClaudeCode => claude::launch::prepare(launch, gateway_url, dry_run),
CodingAgent::ClaudeCode => {
claude::launch::prepare(launch, gateway_url, proxy_credential, dry_run)
}
CodingAgent::Hermes => hermes::launch::prepare(
launch,
resolved.agents.hermes.hooks_path.as_deref(),
Expand Down
3 changes: 1 addition & 2 deletions crates/cli/src/agents/shared/alignment.rs
Original file line number Diff line number Diff line change
Expand Up @@ -546,8 +546,7 @@ pub(crate) fn gateway_upstream_url_override(

/// Remove or preserve agent-native auth before generic provider auth injection.
///
/// Codex strips ChatGPT auth tokens only when an OpenAI API key is available to
/// replace them.
/// Codex strips ChatGPT auth tokens only when an OpenAI API key is available to replace them.
pub(crate) fn gateway_forward_headers(
headers: &HeaderMap,
route: GatewayRouteKind,
Expand Down
3 changes: 3 additions & 0 deletions crates/cli/src/error.rs
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,8 @@ pub(crate) enum CliError {
InvalidPayload(String),
#[error("payload too large: {0}")]
PayloadTooLarge(String),
#[error("unauthorized gateway client: {0}")]
Unauthorized(String),
#[error("gateway upstream error: {0}")]
Upstream(#[from] reqwest::Error),
#[error("{0}")]
Expand Down Expand Up @@ -98,6 +100,7 @@ impl IntoResponse for CliError {
let status = match (guardrail_reason.is_some(), &self) {
(true, _) => StatusCode::FORBIDDEN,
(false, Self::PayloadTooLarge(_)) => StatusCode::PAYLOAD_TOO_LARGE,
(false, Self::Unauthorized(_)) => StatusCode::UNAUTHORIZED,
(false, Self::InvalidPayload(_)) => StatusCode::BAD_REQUEST,
(false, Self::Upstream(_)) => StatusCode::BAD_GATEWAY,
(false, Self::ProviderFailure(failure)) => failure
Expand Down
78 changes: 51 additions & 27 deletions crates/cli/src/gateway/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -65,12 +65,11 @@ const MAX_UPSTREAM_ERROR_BODY_BYTES: usize = 64 * 1024;
/// design that splits a raw byte stream between client and runtime.
pub(crate) async fn passthrough(
State(state): State<AppState>,
request: Request<Body>,
mut request: Request<Body>,
) -> Result<Response<Body>, CliError> {
state.touch();
let allow_environment_provider_auth = state.allows_environment_provider_auth(request.headers());
let prepared =
prepare_gateway_request(&state.config, request, allow_environment_provider_auth).await?;
let authorization = state.authorize_provider_request(request.headers_mut())?;
let prepared = prepare_gateway_request(&state.config, request, authorization).await?;
let prep = state
.sessions
.prepare_gateway_call(&prepared.headers, build_llm_gateway_start(&prepared))
Expand Down Expand Up @@ -131,7 +130,7 @@ async fn run_unmanaged_gateway(
&prepared.body_bytes,
&prepared.headers,
None,
ProviderForwarding::new(prepared.provider, prepared.allow_environment_provider_auth),
ProviderForwarding::new(prepared.provider, prepared.authorization),
)
.await?;
let status = response.status();
Expand Down Expand Up @@ -253,8 +252,7 @@ fn build_buffered_func(
let url = prepared.upstream_url.clone();
let body_bytes = prepared.body_bytes.clone();
let headers = prepared.headers.clone();
let forwarding =
ProviderForwarding::new(prepared.provider, prepared.allow_environment_provider_auth);
let forwarding = ProviderForwarding::new(prepared.provider, prepared.authorization);
Arc::new(move |request| {
let http = http.clone();
let method = method.clone();
Expand Down Expand Up @@ -436,8 +434,7 @@ fn build_streaming_func(
let url = prepared.upstream_url.clone();
let body_bytes = prepared.body_bytes.clone();
let headers = prepared.headers.clone();
let forwarding =
ProviderForwarding::new(prepared.provider, prepared.allow_environment_provider_auth);
let forwarding = ProviderForwarding::new(prepared.provider, prepared.authorization);
Arc::new(move |request| {
let http = http.clone();
let method = method.clone();
Expand Down Expand Up @@ -688,7 +685,7 @@ fn encode_sse_frame(event_json: &Value, route: ProviderRoute) -> String {

// Forwards the buffered request to the upstream provider with only the safe request headers. This
// is shared by the buffered and streaming managed funcs so header filtering stays consistent.
// Agent-native credential quirks are normalized by alignment before provider auth injection runs.
// Source authentication is normalized at ingress, before interceptors can select a target.
async fn forward_upstream_request(
http: &reqwest::Client,
method: &Method,
Expand All @@ -698,31 +695,36 @@ async fn forward_upstream_request(
effective_request: Option<&LlmRequest>,
forwarding: ProviderForwarding,
) -> Result<reqwest::Response, reqwest::Error> {
debug_assert_eq!(
forwarding
.authorization
.source_credential
.provider_credential_present(),
crate::provider_auth::has_provider_credential(headers)
);
let effective = effective_dispatch_request(
body_bytes,
headers,
effective_request,
url,
forwarding.route,
);
let sanitized = strip_replaceable_agent_auth_headers(
&effective.headers,
effective.route,
forwarding.allow_environment_provider_auth,
forwarding.source_route,
);
let mut upstream = http
.request(method.clone(), &effective.url)
.body(effective.body_bytes.clone());
for (name, value) in &sanitized {
if should_forward_request_header(name, &sanitized) {
for (name, value) in &effective.headers {
if should_forward_request_header(name, &effective.headers) {
upstream = upstream.header(name, value);
}
}
upstream = inject_provider_auth(
upstream,
effective.route,
&sanitized,
forwarding.allow_environment_provider_auth,
effective.target_route,
&effective.headers,
matches!(
effective.credential_policy,
TargetCredentialPolicy::SourceOrEnvironment
) && forwarding.authorization.allow_environment_provider_auth,
);
upstream.send().await
}
Expand All @@ -732,7 +734,14 @@ struct EffectiveUpstreamRequest {
body_bytes: Bytes,
headers: HeaderMap,
url: String,
route: ProviderRoute,
target_route: ProviderRoute,
credential_policy: TargetCredentialPolicy,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum TargetCredentialPolicy {
SourceOrEnvironment,
ExplicitTarget,
}

#[cfg(test)]
Expand Down Expand Up @@ -765,7 +774,8 @@ fn effective_dispatch_request(
body_bytes: body_bytes.clone(),
headers,
url: url.to_string(),
route,
target_route: route,
credential_policy: TargetCredentialPolicy::SourceOrEnvironment,
};
};

Expand All @@ -782,7 +792,8 @@ fn effective_dispatch_request(
body_bytes: body_bytes.clone(),
headers,
url: url.to_string(),
route,
target_route: route,
credential_policy: TargetCredentialPolicy::SourceOrEnvironment,
};
}
}
Expand All @@ -801,6 +812,17 @@ fn effective_dispatch_request(
.and_then(|value| ProviderRoute::from_dispatch_override(&value));
continue;
}
}
let credential_policy = if override_url.is_some() || dispatch_route_header_seen {
crate::provider_auth::remove_provider_credentials(&mut headers);
TargetCredentialPolicy::ExplicitTarget
} else {
TargetCredentialPolicy::SourceOrEnvironment
};
// Observable source headers exclude credentials. Applying the rewritten map only after source
// credentials are removed lets an explicit target binding add its own authorization without
// inheriting credentials intended for the original provider.
for (name, value) in &request.headers {
let Ok(name) = HeaderName::from_bytes(name.as_bytes()) else {
continue;
};
Expand All @@ -820,7 +842,8 @@ fn effective_dispatch_request(
} else {
override_url.unwrap_or_else(|| url.to_string())
},
route: override_route.unwrap_or(route),
target_route: override_route.unwrap_or(route),
credential_policy,
}
}

Expand Down Expand Up @@ -934,7 +957,7 @@ async fn passthrough_streaming(
&prepared.body_bytes,
&prepared.headers,
None,
ProviderForwarding::new(prepared.provider, prepared.allow_environment_provider_auth),
ProviderForwarding::new(prepared.provider, prepared.authorization),
)
.await?;
let status = response.status();
Expand Down Expand Up @@ -1078,7 +1101,8 @@ pub(crate) async fn models(
.path_and_query()
.map(|p| p.as_str())
.unwrap_or(parts.uri.path());
let allow_environment_provider_auth = state.allows_environment_provider_auth(&parts.headers);
let authorization = state.authorize_provider_request(&mut parts.headers)?;
let allow_environment_provider_auth = authorization.allow_environment_provider_auth;
parts.headers.remove(BOOTSTRAP_CLIENT_TOKEN_HEADER);
let upstream_url = gateway_upstream_url_override(
provider,
Expand Down
16 changes: 12 additions & 4 deletions crates/cli/src/gateway/request.rs
Original file line number Diff line number Diff line change
Expand Up @@ -30,13 +30,13 @@ pub(super) struct PreparedGatewayRequest {
pub(super) body_bytes: Bytes,
pub(super) request_json: Value,
pub(super) streaming: bool,
pub(super) allow_environment_provider_auth: bool,
pub(super) authorization: crate::provider_auth::ProviderRequestAuthorization,
}

pub(super) async fn prepare_gateway_request(
config: &crate::configuration::GatewayConfig,
request: Request<Body>,
allow_environment_provider_auth: bool,
mut authorization: crate::provider_auth::ProviderRequestAuthorization,
) -> Result<PreparedGatewayRequest, CliError> {
let (mut parts, body) = request.into_parts();
parts.headers.remove(BOOTSTRAP_CLIENT_TOKEN_HEADER);
Expand All @@ -56,9 +56,17 @@ pub(super) async fn prepare_gateway_request(
provider,
&parts.headers,
path_and_query,
allow_environment_provider_auth,
authorization.allow_environment_provider_auth,
)
.unwrap_or_else(|| provider.upstream_url(config, path_and_query));
parts.headers = super::routes::strip_replaceable_agent_auth_headers(
&parts.headers,
provider,
authorization.allow_environment_provider_auth,
);
authorization.source_credential = authorization
.source_credential
.after_source_normalization(&parts.headers);
let streaming = request_json
.get("stream")
.and_then(Value::as_bool)
Expand All @@ -72,7 +80,7 @@ pub(super) async fn prepare_gateway_request(
body_bytes,
request_json,
streaming,
allow_environment_provider_auth,
authorization,
})
}

Expand Down
1 change: 1 addition & 0 deletions crates/cli/src/gateway/response.rs
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@ pub(super) fn should_forward_request_header(name: &HeaderName, headers: &HeaderM
&& name != http::header::HOST
&& name != http::header::CONTENT_LENGTH
&& name.as_str() != BOOTSTRAP_CLIENT_TOKEN_HEADER
&& name.as_str() != crate::provider_auth::TRANSPARENT_PROXY_CREDENTIAL_HEADER
// Strip Accept-Encoding so upstreams return identity-encoded bodies; otherwise the
// observability capture (`output.value` on LLM spans, ATIF trajectory bodies) records
// gzip/br/zstd bytes that downstream consumers can't read. Bandwidth cost is paid only
Expand Down
13 changes: 8 additions & 5 deletions crates/cli/src/gateway/routes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,15 +16,18 @@ pub(super) enum ProviderRoute {

#[derive(Clone, Copy)]
pub(super) struct ProviderForwarding {
pub(super) route: ProviderRoute,
pub(super) allow_environment_provider_auth: bool,
pub(super) source_route: ProviderRoute,
pub(super) authorization: crate::provider_auth::ProviderRequestAuthorization,
}

impl ProviderForwarding {
pub(super) fn new(route: ProviderRoute, allow_environment_provider_auth: bool) -> Self {
pub(super) fn new(
source_route: ProviderRoute,
authorization: crate::provider_auth::ProviderRequestAuthorization,
) -> Self {
Self {
route,
allow_environment_provider_auth,
source_route,
authorization,
}
}
}
Expand Down
1 change: 1 addition & 0 deletions crates/cli/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ mod mcp;
mod mcp_environment;
mod plugins;
mod process;
mod provider_auth;
mod server;
mod sessions;

Expand Down
Loading
Loading