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
48 changes: 38 additions & 10 deletions crates/core/src/api/llm.rs
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,16 @@ pub use nemo_relay_types::api::llm::{
LLM_REQUEST_INTERCEPT_OUTCOME_SCHEMA, LlmAttributes, LlmRequest, LlmRequestInterceptOutcome,
};

const OBSERVABILITY_CREDENTIAL_HEADERS: [&str; 7] = [
"authorization",
"proxy-authorization",
"cookie",
"x-api-key",
"api-key",
"anthropic-api-key",
"x-goog-api-key",
];

#[derive(Clone)]
struct CapturedLlmScopeStack(ScopeStackHandle);

Expand Down Expand Up @@ -430,14 +440,15 @@ fn emit_llm_start_with_subscribers(
.map_err(|error| FlowError::Internal(error.to_string()))?;
state.llm_sanitize_request_entries(&scope_locals)
};
let observable_request = remove_observability_credential_headers(request.clone());
let mut sanitized_request = NemoRelayContextState::llm_sanitize_request_snapshot_chain(
request.clone(),
observable_request.clone(),
LlmSanitizeRequestContext::for_request_codec(request_codec.clone()),
&entries,
);
let request_changed = sanitized_request
.as_ref()
.is_some_and(|sanitized_request| sanitized_request != request);
.is_some_and(|sanitized_request| sanitized_request != &observable_request);
let mut annotated_request = match (sanitized_request.as_ref(), request_codec.as_deref()) {
(Some(sanitized_request), Some(codec)) if request_changed => {
codec.decode(sanitized_request).ok().map(Arc::new)
Expand Down Expand Up @@ -474,6 +485,15 @@ fn emit_llm_start_with_subscribers(
Ok(())
}

fn remove_observability_credential_headers(mut request: LlmRequest) -> LlmRequest {
request.headers.retain(|name, _| {
!OBSERVABILITY_CREDENTIAL_HEADERS
.iter()
.any(|credential_header| name.eq_ignore_ascii_case(credential_header))
});
request
}

fn emit_pending_request_marks(
handle: &LlmHandle,
marks: Vec<PendingMarkSpec>,
Expand Down Expand Up @@ -601,11 +621,14 @@ fn emit_optimization_marks_with(
/// cannot be read safely.
///
/// # Notes
/// Sanitize-request guardrails affect only the emitted start-event payload, not
/// the caller-owned [`LlmRequest`]. When the owning agent is not fresh, the
/// emitted request annotation is limited to the current user turn. Managed
/// calls with a request codec also apply that projection to the event input,
/// without changing the request used for provider execution.
/// The runtime removes standard credential headers (`authorization`,
/// `proxy-authorization`, `cookie`, `x-api-key`, `api-key`,
/// `anthropic-api-key`, and `x-goog-api-key`) from the event-only request copy
/// before sanitize-request guardrails run. This does not change the
/// caller-owned [`LlmRequest`]. When the owning agent is not fresh, the emitted
/// request annotation is limited to the current user turn. Managed calls with a
/// request codec also apply that projection to the event input, without changing
/// the request used for provider execution.
pub fn llm_call(params: LlmCallParams<'_>) -> Result<LlmHandle> {
let handle_params = CreateLlmHandleParams::builder()
.name(params.name)
Expand Down Expand Up @@ -917,9 +940,11 @@ fn emit_llm_end_without_output(
/// execution intercepts, codecs, or the callback itself.
///
/// # Notes
/// The LLM-start event is emitted before execution intercepts run. When
/// execution fails after that point, the runtime still emits an LLM-end event
/// without an output payload.
/// The LLM-start event is emitted before execution intercepts run. Before
/// sanitize-request guardrails run, the runtime removes standard credential
/// headers from the event-only request copy; the request passed to execution is
/// unchanged. When execution fails after that point, the runtime still emits an
/// LLM-end event without an output payload.
///
/// Response codecs enrich observability output only and do not change the
/// value returned to the caller.
Expand Down Expand Up @@ -1108,6 +1133,9 @@ pub async fn llm_call_execute(params: LlmCallExecuteParams) -> Result<Json> {
///
/// # Notes
/// The LLM-start event is emitted before stream execution intercepts run.
/// Before sanitize-request guardrails run, the runtime removes standard
/// credential headers from the event-only request copy; the request passed to
/// stream execution is unchanged.
///
/// The returned stream emits chunk-level results while the runtime defers the
/// LLM-end event until the collector and finalizer complete.
Expand Down
146 changes: 146 additions & 0 deletions crates/core/tests/unit/llm_api_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,34 @@ fn request() -> LlmRequest {
}
}

fn request_with_credential_headers() -> LlmRequest {
let mut headers = serde_json::Map::new();
for (name, value) in [
("Authorization", "Bearer authorization-secret"),
("PrOxY-AuThOrIzAtIoN", "Basic proxy-secret"),
("COOKIE", "session=cookie-secret"),
("X-Api-Key", "x-api-key-secret"),
("API-KEY", "api-key-secret"),
("Anthropic-Api-Key", "anthropic-api-key-secret"),
("X-GoOg-Api-Key", "x-goog-api-key-secret"),
] {
headers.insert(name.to_string(), json!(value));
}
headers.insert("x-request-id".to_string(), json!("safe-request-id"));
LlmRequest {
headers,
content: json!({"messages": [], "model": "demo"}),
}
}

fn assert_observable_credential_headers_are_removed(request: &LlmRequest) {
assert_eq!(request.headers.len(), 1);
assert_eq!(
request.headers.get("x-request-id"),
Some(&json!("safe-request-id"))
);
}

fn multi_turn_request() -> LlmRequest {
LlmRequest {
headers: serde_json::Map::new(),
Expand Down Expand Up @@ -253,6 +281,124 @@ fn redacted_request() -> LlmRequest {
}
}

#[test]
fn credential_headers_are_removed_before_request_sanitizers_and_event_emission() {
let _guard = lock_global_runtime();
reset_global();
set_thread_scope_stack(create_scope_stack());

let request = request_with_credential_headers();
let sanitizer_requests = Arc::new(Mutex::new(Vec::<LlmRequest>::new()));
let sanitizer_capture = Arc::clone(&sanitizer_requests);
register_llm_sanitize_request_guardrail(
"credential-header-redaction",
1,
Arc::new(move |request, _context| {
sanitizer_capture.lock().unwrap().push(request.clone());
Some(request)
}),
)
.unwrap();

let events = Arc::new(Mutex::new(Vec::<Event>::new()));
let event_capture = Arc::clone(&events);
register_subscriber(
"credential-header-redaction",
Arc::new(move |event| event_capture.lock().unwrap().push(event.clone())),
)
.unwrap();

llm_call(
LlmCallParams::builder()
.name("credential-header-manual")
.request(&request)
.build(),
)
.unwrap();

let provider_requests = Arc::new(Mutex::new(Vec::<LlmRequest>::new()));
let buffered_provider_requests = Arc::clone(&provider_requests);
let runtime = tokio::runtime::Runtime::new().unwrap();
runtime.block_on(async {
llm_call_execute(
LlmCallExecuteParams::builder()
.name("credential-header-buffered")
.request(request.clone())
.func(Arc::new(move |request| {
buffered_provider_requests.lock().unwrap().push(request);
Box::pin(async { Ok(json!({"ok": true})) })
}))
.build(),
)
.await
.unwrap();

let streaming_provider_requests = Arc::clone(&provider_requests);
let mut stream = llm_stream_call_execute(
LlmStreamCallExecuteParams::builder()
.name("credential-header-streaming")
.request(request.clone())
.func(Arc::new(move |request| {
streaming_provider_requests.lock().unwrap().push(request);
Box::pin(async {
Ok(LlmJsonStream::new(tokio_stream::iter(vec![Ok(json!({
"chunk": true
}))])))
})
}))
.collector(Box::new(|_chunk| Ok(())))
.finalizer(Box::new(|| json!({"ok": true})))
.build(),
)
.await
.unwrap();
while let Some(chunk) = stream.next().await {
chunk.unwrap();
}
});

flush_subscribers().unwrap();
for sanitized in sanitizer_requests.lock().unwrap().iter() {
assert_observable_credential_headers_are_removed(sanitized);
}
assert_eq!(sanitizer_requests.lock().unwrap().len(), 3);

let events = events.lock().unwrap();
let start_events = [
"credential-header-manual",
"credential-header-buffered",
"credential-header-streaming",
]
.into_iter()
.map(|name| {
events
.iter()
.find(|event| {
event.name() == name && event.scope_category() == Some(ScopeCategory::Start)
})
.cloned()
.unwrap_or_else(|| panic!("missing LLM start event {name}"))
})
.collect::<Vec<_>>();
drop(events);
assert_eq!(start_events.len(), 3);
for event in start_events {
let input: LlmRequest = serde_json::from_value(event.input().cloned().unwrap()).unwrap();
assert_observable_credential_headers_are_removed(&input);
}

let provider_requests = provider_requests.lock().unwrap();
assert_eq!(provider_requests.len(), 2);
assert!(
provider_requests
.iter()
.all(|provider| provider == &request)
);

assert!(deregister_llm_sanitize_request_guardrail("credential-header-redaction").unwrap());
assert!(deregister_subscriber("credential-header-redaction").unwrap());
}

fn secret_response() -> Json {
json!({
"id": "chatcmpl-test",
Expand Down
7 changes: 7 additions & 0 deletions docs/about-nemo-relay/concepts/middleware.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -191,6 +191,13 @@ guardrails important:
- If you need to change the real execution path, use an intercept
- If you need to change only the emitted payload, use a sanitize guardrail

Before LLM request sanitizers run, Relay removes standard credential headers
from the event-only request copy: `authorization`, `proxy-authorization`,
`cookie`, `x-api-key`, `api-key`, `anthropic-api-key`, and `x-goog-api-key`.
Header-name matching is case-insensitive. This protects sanitizer callbacks,
subscribers, and exporters without changing the request sent to the provider.
Configure a PII sanitizer for custom credential header names.

## Codec-Aware LLM Sanitizers

Every LLM sanitize guardrail receives the payload first and a required
Expand Down
Loading