use std::sync::{Arc, Mutex};
use async_stream::stream;
use axum::body::{Body, Bytes};
use axum::extract::State;
use axum::http::{HeaderMap, HeaderName, HeaderValue, Method, Request, Response, StatusCode};
use futures_util::StreamExt;
use nemo_relay::api::llm::{
LlmCallExecuteParams, LlmRequest, LlmStreamCallExecuteParams, llm_call_execute,
llm_stream_call_execute,
};
use nemo_relay::api::runtime::{
LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn, TASK_SCOPE_STACK,
};
use nemo_relay::codec::anthropic::{AnthropicMessagesCodec, AnthropicMessagesStreamingCodec};
use nemo_relay::codec::openai_chat::{OpenAIChatCodec, OpenAIChatStreamingCodec};
use nemo_relay::codec::openai_responses::{OpenAIResponsesCodec, OpenAIResponsesStreamingCodec};
use nemo_relay::codec::streaming::StreamingCodec;
use nemo_relay::codec::traits::LlmResponseCodec;
use nemo_relay::error::FlowError;
use serde_json::{Map, Value, json};
use crate::alignment::{self, GatewayRouteKind};
use crate::config::header_string;
use crate::error::CliError;
use crate::server::AppState;
use crate::session::{GatewayCallPrep, LlmGatewayStart, SessionManager};
const MAX_BODY_BYTES: usize = 100 * 1024 * 1024;
pub(crate) async fn passthrough(
State(state): State<AppState>,
request: Request<Body>,
) -> Result<Response<Body>, CliError> {
let prepared = prepare_gateway_request(&state.config, request).await?;
let prep = state
.sessions
.prepare_gateway_call(&prepared.headers, build_llm_gateway_start(&prepared))
.await?;
run_managed_gateway(state, prepared, prep).await
}
struct PreparedGatewayRequest {
method: Method,
headers: HeaderMap,
path: String,
provider: ProviderRoute,
upstream_url: String,
body_bytes: Bytes,
request_json: Value,
streaming: bool,
}
async fn prepare_gateway_request(
config: &crate::config::GatewayConfig,
request: Request<Body>,
) -> Result<PreparedGatewayRequest, CliError> {
let (parts, body) = request.into_parts();
let provider = ProviderRoute::from_path(parts.uri.path()).ok_or_else(|| {
CliError::InvalidPayload(format!("unsupported gateway path {}", parts.uri.path()))
})?;
let body_bytes = axum::body::to_bytes(body, MAX_BODY_BYTES)
.await
.map_err(|error| CliError::InvalidPayload(error.to_string()))?;
let request_json = serde_json::from_slice::<Value>(&body_bytes).unwrap_or(Value::Null);
let path_and_query = parts
.uri
.path_and_query()
.map(|p| p.as_str())
.unwrap_or(parts.uri.path());
let upstream_url = gateway_upstream_url_override(provider, &parts.headers, path_and_query)
.unwrap_or_else(|| provider.upstream_url(config, path_and_query));
let streaming = request_json
.get("stream")
.and_then(Value::as_bool)
.unwrap_or(false);
Ok(PreparedGatewayRequest {
method: parts.method,
headers: parts.headers,
path: parts.uri.path().to_string(),
provider,
upstream_url,
body_bytes,
request_json,
streaming,
})
}
fn build_llm_gateway_start(request: &PreparedGatewayRequest) -> LlmGatewayStart {
LlmGatewayStart {
session_id: gateway_session_id(&request.headers, &request.request_json, request.provider),
provider: request.provider.name().to_string(),
model_name: request
.request_json
.get("model")
.and_then(Value::as_str)
.map(ToOwned::to_owned),
subagent_id: gateway_subagent_id(&request.headers),
conversation_id: gateway_identifier(
&request.headers,
&request.request_json,
"x-nemo-relay-conversation-id",
&[
&["conversation_id"],
&["conversationId"],
&["conversation", "id"],
],
),
generation_id: gateway_identifier(
&request.headers,
&request.request_json,
"x-nemo-relay-generation-id",
&[&["generation_id"], &["generationId"], &["generation", "id"]],
),
request_id: gateway_identifier(
&request.headers,
&request.request_json,
"x-nemo-relay-request-id",
&[
&["request_id"],
&["requestId"],
&["request", "id"],
&["metadata", "request_id"],
],
)
.or_else(|| header_string(&request.headers, "x-request-id")),
request: LlmRequest {
headers: observable_headers(&request.headers),
content: request.request_json.clone(),
},
streaming: request.streaming,
metadata: json!({ "gateway_path": request.path }),
}
}
type UpstreamResponseInfo = Arc<Mutex<Option<(StatusCode, HeaderMap)>>>;
type UpstreamErrorSlot = Arc<Mutex<Option<reqwest::Error>>>;
async fn run_managed_gateway(
state: AppState,
prepared: PreparedGatewayRequest,
prep: GatewayCallPrep,
) -> Result<Response<Body>, CliError> {
let codecs = codecs_for_route(prepared.provider);
if prepared.streaming {
run_managed_streaming(state, prepared, prep, codecs).await
} else {
run_managed_buffered(state, prepared, prep, codecs).await
}
}
struct RouteCodecs {
streaming: Option<Box<dyn StreamingCodec>>,
response: Option<Arc<dyn LlmResponseCodec>>,
}
fn codecs_for_route(route: ProviderRoute) -> RouteCodecs {
match route {
ProviderRoute::AnthropicMessages => RouteCodecs {
streaming: Some(Box::new(AnthropicMessagesStreamingCodec::new())),
response: Some(Arc::new(AnthropicMessagesCodec) as Arc<dyn LlmResponseCodec>),
},
ProviderRoute::OpenAiResponses => RouteCodecs {
streaming: Some(Box::new(OpenAIResponsesStreamingCodec::new())),
response: Some(Arc::new(OpenAIResponsesCodec) as Arc<dyn LlmResponseCodec>),
},
ProviderRoute::OpenAiChatCompletions => RouteCodecs {
streaming: Some(Box::new(OpenAIChatStreamingCodec::new())),
response: Some(Arc::new(OpenAIChatCodec) as Arc<dyn LlmResponseCodec>),
},
ProviderRoute::AnthropicCountTokens | ProviderRoute::OpenAiModels => RouteCodecs {
streaming: None,
response: None,
},
}
}
async fn run_managed_buffered(
state: AppState,
prepared: PreparedGatewayRequest,
prep: GatewayCallPrep,
codecs: RouteCodecs,
) -> Result<Response<Body>, CliError> {
let upstream_info: UpstreamResponseInfo = Arc::new(Mutex::new(None));
let upstream_error: UpstreamErrorSlot = Arc::new(Mutex::new(None));
let response_bytes: Arc<Mutex<Option<Bytes>>> = Arc::new(Mutex::new(None));
let func = build_buffered_func(
state.clone(),
&prepared,
upstream_info.clone(),
upstream_error.clone(),
response_bytes.clone(),
);
let GatewayCallPrep {
scope_stack,
session_id,
provider_name,
request,
parent,
attributes,
metadata,
model_name,
owner_subagent_id,
} = prep;
let provider_for_event = provider_name.clone();
let params = LlmCallExecuteParams::builder()
.name(provider_for_event)
.request(request)
.func(func)
.parent_opt(parent)
.attributes(attributes)
.metadata(metadata)
.model_name_opt(model_name)
.response_codec_opt(codecs.response)
.build();
let result = TASK_SCOPE_STACK
.scope(scope_stack, async move { llm_call_execute(params).await })
.await;
match result {
Ok(response_json) => {
state
.sessions
.record_gateway_response_hints(&session_id, owner_subagent_id, response_json)
.await;
state.sessions.finish_gateway_call(&session_id).await;
let (status, headers) = upstream_info
.lock()
.expect("upstream info lock poisoned")
.take()
.unwrap_or((StatusCode::OK, HeaderMap::new()));
let bytes = response_bytes
.lock()
.expect("response bytes lock poisoned")
.take()
.unwrap_or_default();
build_response(status, headers, Body::from(bytes))
}
Err(error) => {
state.sessions.finish_gateway_call(&session_id).await;
Err(translate_runtime_error(error, &upstream_error))
}
}
}
fn build_buffered_func(
state: AppState,
prepared: &PreparedGatewayRequest,
upstream_info: UpstreamResponseInfo,
upstream_error: UpstreamErrorSlot,
response_bytes: Arc<Mutex<Option<Bytes>>>,
) -> LlmExecutionNextFn {
let http = state.http.clone();
let method = prepared.method.clone();
let url = prepared.upstream_url.clone();
let body_bytes = prepared.body_bytes.clone();
let headers = prepared.headers.clone();
let route = prepared.provider;
Arc::new(move |request| {
let http = http.clone();
let method = method.clone();
let url = url.clone();
let body_bytes = body_bytes.clone();
let headers = headers.clone();
let upstream_info = upstream_info.clone();
let upstream_error = upstream_error.clone();
let response_bytes = response_bytes.clone();
Box::pin(async move {
let response = match forward_upstream_request(
&http,
&method,
&url,
&body_bytes,
&headers,
Some(&request),
route,
)
.await
{
Ok(response) => response,
Err(error) => {
let message = error.to_string();
*upstream_error.lock().expect("upstream error lock poisoned") = Some(error);
return Err(FlowError::Internal(message));
}
};
let status = response.status();
let response_headers = response_headers(response.headers());
let bytes = match response.bytes().await {
Ok(bytes) => bytes,
Err(error) => {
let message = error.to_string();
*upstream_error.lock().expect("upstream error lock poisoned") = Some(error);
return Err(FlowError::Internal(message));
}
};
let json = serde_json::from_slice::<Value>(&bytes)
.unwrap_or_else(|_| json!({ "body_bytes": bytes.len() }));
*upstream_info.lock().expect("upstream info lock poisoned") =
Some((status, response_headers));
*response_bytes.lock().expect("response bytes lock poisoned") = Some(bytes);
Ok(json)
})
})
}
async fn run_managed_streaming(
state: AppState,
prepared: PreparedGatewayRequest,
prep: GatewayCallPrep,
codecs: RouteCodecs,
) -> Result<Response<Body>, CliError> {
let upstream_info: UpstreamResponseInfo = Arc::new(Mutex::new(None));
let upstream_error: UpstreamErrorSlot = Arc::new(Mutex::new(None));
let func = build_streaming_func(
state.clone(),
&prepared,
upstream_info.clone(),
upstream_error.clone(),
);
let provider_route = prepared.provider;
let Some(streaming_codec) = codecs.streaming else {
state.sessions.finish_gateway_call(&prep.session_id).await;
return passthrough_streaming(state, prepared).await;
};
let collector = streaming_codec.collector();
let final_response = Arc::new(Mutex::new(None));
let final_response_for_finalizer = final_response.clone();
let original_finalizer = streaming_codec.finalizer();
let finalizer = Box::new(move || {
let response = original_finalizer();
*final_response_for_finalizer
.lock()
.expect("stream final response lock poisoned") = Some(response.clone());
response
});
let GatewayCallPrep {
scope_stack,
session_id,
provider_name,
request,
parent,
attributes,
metadata,
model_name,
owner_subagent_id,
} = prep;
let params = LlmStreamCallExecuteParams::builder()
.name(provider_name)
.request(request)
.func(func)
.collector(collector)
.finalizer(finalizer)
.parent_opt(parent)
.attributes(attributes)
.metadata(metadata)
.model_name_opt(model_name)
.response_codec_opt(codecs.response)
.build();
let json_stream_result = TASK_SCOPE_STACK
.scope(
scope_stack,
async move { llm_stream_call_execute(params).await },
)
.await;
let json_stream = match json_stream_result {
Ok(json_stream) => json_stream,
Err(error) => {
state.sessions.finish_gateway_call(&session_id).await;
return Err(translate_runtime_error(error, &upstream_error));
}
};
let (status, headers) = upstream_info
.lock()
.expect("upstream info lock poisoned")
.take()
.unwrap_or((StatusCode::OK, HeaderMap::new()));
let body = client_sse_body(
json_stream,
provider_route,
state.sessions.clone(),
session_id.clone(),
owner_subagent_id,
final_response,
);
build_response(status, headers, body)
}
fn build_streaming_func(
state: AppState,
prepared: &PreparedGatewayRequest,
upstream_info: UpstreamResponseInfo,
upstream_error: UpstreamErrorSlot,
) -> LlmStreamExecutionNextFn {
let http = state.http.clone();
let method = prepared.method.clone();
let url = prepared.upstream_url.clone();
let body_bytes = prepared.body_bytes.clone();
let headers = prepared.headers.clone();
let route = prepared.provider;
Arc::new(move |request| {
let http = http.clone();
let method = method.clone();
let url = url.clone();
let body_bytes = body_bytes.clone();
let headers = headers.clone();
let upstream_info = upstream_info.clone();
let upstream_error = upstream_error.clone();
Box::pin(async move {
let response = match forward_upstream_request(
&http,
&method,
&url,
&body_bytes,
&headers,
Some(&request),
route,
)
.await
{
Ok(response) => response,
Err(error) => {
let message = error.to_string();
*upstream_error.lock().expect("upstream error lock poisoned") = Some(error);
return Err(FlowError::Internal(message));
}
};
let status = response.status();
let response_headers = response_headers(response.headers());
*upstream_info.lock().expect("upstream info lock poisoned") =
Some((status, response_headers));
let json_stream = sse_json_stream(response);
Ok(json_stream)
})
})
}
fn sse_json_stream(response: reqwest::Response) -> LlmJsonStream {
use nemo_relay::codec::streaming::SseEventDecoder;
let mut decoder = SseEventDecoder::new();
let mut bytes = response.bytes_stream();
let stream = stream! {
while let Some(chunk) = bytes.next().await {
match chunk {
Ok(buffer) => {
match decoder.push_bytes(&buffer) {
Ok(events) => {
for event in events {
yield Ok(event.data);
}
}
Err(error) => {
yield Err(error);
return;
}
}
}
Err(error) => {
yield Err(FlowError::Internal(error.to_string()));
return;
}
}
}
match decoder.finish() {
Ok(Some(event)) => yield Ok(event.data),
Ok(None) => {}
Err(error) => yield Err(error),
}
};
Box::pin(stream)
}
fn client_sse_body(
json_stream: LlmJsonStream,
route: ProviderRoute,
sessions: SessionManager,
session_id: String,
owner_subagent_id: Option<String>,
final_response: Arc<Mutex<Option<Value>>>,
) -> Body {
let mut json_stream = json_stream;
let mut guard = GatewayCallGuard::new(sessions, session_id, owner_subagent_id, final_response);
let stream = stream! {
while let Some(item) = json_stream.next().await {
match item {
Ok(event_json) => {
let frame = encode_sse_frame(&event_json, route);
yield Ok::<Bytes, CliError>(Bytes::from(frame));
}
Err(error) => {
guard.finish().await;
yield Err(CliError::InvalidPayload(error.to_string()));
return;
}
}
}
guard.finish().await;
if matches!(route, ProviderRoute::OpenAiChatCompletions) {
yield Ok::<Bytes, CliError>(Bytes::from_static(b"data: [DONE]\n\n"));
}
};
Body::from_stream(stream)
}
struct GatewayCallGuard {
sessions: Option<SessionManager>,
session_id: String,
owner_subagent_id: Option<String>,
final_response: Arc<Mutex<Option<Value>>>,
}
impl GatewayCallGuard {
fn new(
sessions: SessionManager,
session_id: String,
owner_subagent_id: Option<String>,
final_response: Arc<Mutex<Option<Value>>>,
) -> Self {
Self {
sessions: Some(sessions),
session_id,
owner_subagent_id,
final_response,
}
}
async fn finish(&mut self) {
if let Some(sessions) = self.sessions.take() {
let response = self
.final_response
.lock()
.expect("stream final response lock poisoned")
.take();
if let Some(response) = response {
sessions
.record_gateway_response_hints(
&self.session_id,
self.owner_subagent_id.clone(),
response,
)
.await;
}
sessions.finish_gateway_call(&self.session_id).await;
}
}
}
impl Drop for GatewayCallGuard {
fn drop(&mut self) {
let Some(sessions) = self.sessions.take() else {
return;
};
let session_id = self.session_id.clone();
let owner_subagent_id = self.owner_subagent_id.clone();
let response = self
.final_response
.lock()
.expect("stream final response lock poisoned")
.take();
if let Ok(handle) = tokio::runtime::Handle::try_current() {
handle.spawn(async move {
if let Some(response) = response {
sessions
.record_gateway_response_hints(&session_id, owner_subagent_id, response)
.await;
}
sessions.finish_gateway_call(&session_id).await;
});
}
}
}
fn encode_sse_frame(event_json: &Value, route: ProviderRoute) -> String {
let serialized = serde_json::to_string(event_json).unwrap_or_else(|_| "null".to_string());
let event_name = match route {
ProviderRoute::AnthropicMessages | ProviderRoute::OpenAiResponses => event_json
.get("type")
.and_then(Value::as_str)
.map(ToOwned::to_owned),
_ => None,
};
match event_name {
Some(name) => format!("event: {name}\ndata: {serialized}\n\n"),
None => format!("data: {serialized}\n\n"),
}
}
async fn forward_upstream_request(
http: &reqwest::Client,
method: &Method,
url: &str,
body_bytes: &Bytes,
headers: &HeaderMap,
effective_request: Option<&LlmRequest>,
route: ProviderRoute,
) -> Result<reqwest::Response, reqwest::Error> {
let (body_bytes, headers) = effective_upstream_request(body_bytes, headers, effective_request);
let sanitized = strip_replaceable_agent_auth_headers(&headers, route);
let mut upstream = http.request(method.clone(), url).body(body_bytes.clone());
for (name, value) in &sanitized {
if should_forward_request_header(name) {
upstream = upstream.header(name, value);
}
}
upstream = inject_provider_auth(upstream, route, &sanitized);
upstream.send().await
}
fn effective_upstream_request(
body_bytes: &Bytes,
headers: &HeaderMap,
effective_request: Option<&LlmRequest>,
) -> (Bytes, HeaderMap) {
let Some(request) = effective_request else {
return (body_bytes.clone(), headers.clone());
};
let body_bytes = if request.content.is_null() {
body_bytes.clone()
} else {
match serde_json::to_vec(&request.content) {
Ok(serialized) => Bytes::from(serialized),
Err(error) => {
eprintln!(
"nemo-relay CLI gateway: failed to serialize rewritten LLM request body; forwarding original request: {error}"
);
return (body_bytes.clone(), headers.clone());
}
}
};
let mut headers = headers.clone();
for (name, value) in &request.headers {
let Ok(name) = HeaderName::from_bytes(name.as_bytes()) else {
continue;
};
let Some(value) = json_header_value(value) else {
continue;
};
headers.insert(name, value);
}
(body_bytes, headers)
}
fn json_header_value(value: &Value) -> Option<HeaderValue> {
let rendered = match value {
Value::String(value) => value.clone(),
value => serde_json::to_string(value).ok()?,
};
HeaderValue::from_str(&rendered).ok()
}
fn inject_provider_auth(
builder: reqwest::RequestBuilder,
route: ProviderRoute,
inbound: &HeaderMap,
) -> reqwest::RequestBuilder {
inject_provider_auth_with_env(builder, route, inbound, |key| std::env::var(key).ok())
}
fn inject_provider_auth_with_env<F>(
builder: reqwest::RequestBuilder,
route: ProviderRoute,
inbound: &HeaderMap,
env_lookup: F,
) -> reqwest::RequestBuilder
where
F: Fn(&str) -> Option<String>,
{
let already_authed = inbound.contains_key(http::header::AUTHORIZATION)
|| inbound.contains_key("x-api-key")
|| inbound.contains_key("api-key")
|| inbound.contains_key("anthropic-api-key");
if already_authed {
return builder;
}
let (env_var, header_name) = match route {
ProviderRoute::OpenAiResponses
| ProviderRoute::OpenAiChatCompletions
| ProviderRoute::OpenAiModels => ("OPENAI_API_KEY", http::header::AUTHORIZATION.as_str()),
ProviderRoute::AnthropicMessages | ProviderRoute::AnthropicCountTokens => {
("ANTHROPIC_API_KEY", "x-api-key")
}
};
let Some(value) = env_lookup(env_var) else {
return builder;
};
let value = value.trim().to_string();
if value.is_empty() {
return builder;
}
let header_value = match route {
ProviderRoute::OpenAiResponses
| ProviderRoute::OpenAiChatCompletions
| ProviderRoute::OpenAiModels => format!("Bearer {value}"),
ProviderRoute::AnthropicMessages | ProviderRoute::AnthropicCountTokens => value,
};
builder.header(header_name, header_value)
}
async fn passthrough_streaming(
state: AppState,
prepared: PreparedGatewayRequest,
) -> Result<Response<Body>, CliError> {
let response = forward_upstream_request(
&state.http,
&prepared.method,
&prepared.upstream_url,
&prepared.body_bytes,
&prepared.headers,
None,
prepared.provider,
)
.await?;
let status = response.status();
let headers = response_headers(response.headers());
let mut bytes = response.bytes_stream();
let body = Body::from_stream(stream! {
while let Some(chunk) = bytes.next().await {
yield chunk;
}
});
build_response(status, headers, body)
}
fn translate_runtime_error(error: FlowError, upstream_error: &UpstreamErrorSlot) -> CliError {
if let Some(upstream) = upstream_error
.lock()
.expect("upstream error lock poisoned")
.take()
{
return CliError::Upstream(upstream);
}
match error {
FlowError::GuardrailRejected(reason) => CliError::GuardrailRejected(reason),
other => CliError::InvalidPayload(other.to_string()),
}
}
pub(crate) async fn models(
State(state): State<AppState>,
request: Request<Body>,
) -> Result<Response<Body>, CliError> {
let (parts, _body) = request.into_parts();
if parts.method != Method::GET {
return build_response(
StatusCode::METHOD_NOT_ALLOWED,
HeaderMap::new(),
Body::empty(),
);
}
let provider = ProviderRoute::OpenAiModels;
let path_and_query = parts
.uri
.path_and_query()
.map(|p| p.as_str())
.unwrap_or(parts.uri.path());
let upstream_url = gateway_upstream_url_override(provider, &parts.headers, path_and_query)
.unwrap_or_else(|| provider.upstream_url(&state.config, path_and_query));
let sanitized = strip_replaceable_agent_auth_headers(&parts.headers, provider);
let mut upstream = state.http.get(upstream_url);
for (name, value) in &sanitized {
if should_forward_request_header(name) {
upstream = upstream.header(name, value);
}
}
upstream = inject_provider_auth(upstream, provider, &sanitized);
let upstream_response = upstream.send().await?;
let status = upstream_response.status();
let headers = response_headers(upstream_response.headers());
let bytes = upstream_response.bytes().await?;
build_response(status, headers, Body::from(bytes))
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ProviderRoute {
OpenAiResponses,
OpenAiChatCompletions,
OpenAiModels,
AnthropicMessages,
AnthropicCountTokens,
}
impl ProviderRoute {
fn from_path(path: &str) -> Option<Self> {
match path {
"/responses" => Some(Self::OpenAiResponses),
"/v1/responses" => Some(Self::OpenAiResponses),
"/chat/completions" => Some(Self::OpenAiChatCompletions),
"/v1/chat/completions" => Some(Self::OpenAiChatCompletions),
"/models" => Some(Self::OpenAiModels),
"/v1/models" => Some(Self::OpenAiModels),
"/v1/messages" => Some(Self::AnthropicMessages),
"/v1/messages/count_tokens" => Some(Self::AnthropicCountTokens),
_ => None,
}
}
const fn name(self) -> &'static str {
match self {
Self::OpenAiResponses => "openai.responses",
Self::OpenAiChatCompletions => "openai.chat_completions",
Self::OpenAiModels => "openai.models",
Self::AnthropicMessages => "anthropic.messages",
Self::AnthropicCountTokens => "anthropic.count_tokens",
}
}
fn upstream_url(self, config: &crate::config::GatewayConfig, path_and_query: &str) -> String {
let base = match self {
Self::OpenAiResponses | Self::OpenAiChatCompletions | Self::OpenAiModels => {
config.openai_base_url.as_str()
}
Self::AnthropicMessages | Self::AnthropicCountTokens => {
config.anthropic_base_url.as_str()
}
};
self.upstream_url_with_base(base, path_and_query)
}
fn upstream_url_with_base(self, base: &str, path_and_query: &str) -> String {
let base = base.trim_end_matches('/');
let path_and_query = match self {
Self::OpenAiResponses | Self::OpenAiChatCompletions | Self::OpenAiModels => {
normalize_openai_path_for_base(base, path_and_query)
}
_ => path_and_query.to_string(),
};
format!("{base}{path_and_query}")
}
const fn alignment_route(self) -> GatewayRouteKind {
match self {
Self::OpenAiResponses => GatewayRouteKind::OpenAiResponses,
Self::OpenAiChatCompletions => GatewayRouteKind::OpenAiChatCompletions,
Self::OpenAiModels => GatewayRouteKind::OpenAiModels,
Self::AnthropicMessages => GatewayRouteKind::AnthropicMessages,
Self::AnthropicCountTokens => GatewayRouteKind::AnthropicCountTokens,
}
}
}
fn normalize_openai_path_for_base(base: &str, path_and_query: &str) -> String {
match (base.ends_with("/v1"), path_and_query.starts_with("/v1/")) {
(true, true) => path_and_query
.strip_prefix("/v1")
.expect("path was checked to start with /v1")
.to_string(),
(false, false) => format!("/v1{path_and_query}"),
_ => path_and_query.to_string(),
}
}
fn gateway_upstream_url_override(
route: ProviderRoute,
headers: &HeaderMap,
path_and_query: &str,
) -> Option<String> {
gateway_upstream_url_override_with_openai_key_state(
route,
headers,
path_and_query,
env_var_is_nonempty("OPENAI_API_KEY"),
)
}
fn gateway_upstream_url_override_with_openai_key_state(
route: ProviderRoute,
headers: &HeaderMap,
path_and_query: &str,
has_openai_replacement_key: bool,
) -> Option<String> {
alignment::gateway_upstream_url_override(
headers,
route.alignment_route(),
path_and_query,
has_openai_replacement_key,
)
}
fn strip_replaceable_agent_auth_headers(headers: &HeaderMap, route: ProviderRoute) -> HeaderMap {
strip_replaceable_agent_auth_headers_with_openai_key_state(
headers,
route,
env_var_is_nonempty("OPENAI_API_KEY"),
)
}
fn strip_replaceable_agent_auth_headers_with_openai_key_state(
headers: &HeaderMap,
route: ProviderRoute,
has_openai_replacement_key: bool,
) -> HeaderMap {
alignment::gateway_forward_headers(headers, route.alignment_route(), has_openai_replacement_key)
}
fn env_var_is_nonempty(name: &str) -> bool {
std::env::var(name)
.ok()
.filter(|value| !value.trim().is_empty())
.is_some()
}
fn gateway_session_id(headers: &HeaderMap, body: &Value, route: ProviderRoute) -> Option<String> {
alignment::gateway_session_id(headers, body, route.alignment_route())
}
fn gateway_subagent_id(headers: &HeaderMap) -> Option<String> {
alignment::gateway_subagent_id(headers)
}
fn gateway_identifier(
headers: &HeaderMap,
body: &Value,
header_name: &'static str,
body_paths: &[&[&str]],
) -> Option<String> {
alignment::gateway_identifier(headers, body, header_name, body_paths)
}
fn observable_headers(headers: &HeaderMap) -> Map<String, Value> {
let mut output = Map::new();
for (name, value) in headers {
if should_record_header(name)
&& let Ok(value) = value.to_str()
{
output.insert(name.as_str().to_string(), json!(value));
}
}
output
}
fn response_headers(headers: &HeaderMap) -> HeaderMap {
let mut output = HeaderMap::new();
for (name, value) in headers {
if !is_hop_by_hop(name) && name != http::header::CONTENT_LENGTH {
output.append(name.clone(), value.clone());
}
}
output
}
fn build_response(
status: StatusCode,
headers: HeaderMap,
body: Body,
) -> Result<Response<Body>, CliError> {
let mut builder = Response::builder().status(status);
for (name, value) in &headers {
builder = builder.header(name, value);
}
Ok(builder.body(body)?)
}
fn should_forward_request_header(name: &HeaderName) -> bool {
!is_hop_by_hop(name)
&& name != http::header::HOST
&& name != http::header::CONTENT_LENGTH
&& name != http::header::ACCEPT_ENCODING
}
fn should_record_header(name: &HeaderName) -> bool {
should_forward_request_header(name)
&& name != http::header::AUTHORIZATION
&& name != http::header::COOKIE
&& name.as_str() != "x-api-key"
&& name.as_str() != "api-key"
&& name.as_str() != "anthropic-api-key"
}
fn is_hop_by_hop(name: &HeaderName) -> bool {
matches!(
name.as_str(),
"connection"
| "keep-alive"
| "proxy-authenticate"
| "proxy-authorization"
| "te"
| "trailer"
| "transfer-encoding"
| "upgrade"
)
}
#[cfg(test)]
#[path = "../tests/coverage/gateway_tests.rs"]
mod tests;