From b478cbe42bd85904e65f99194a7177fffa7d0438 Mon Sep 17 00:00:00 2001 From: Maryam Najafian Date: Thu, 23 Jul 2026 09:31:54 -0700 Subject: [PATCH 1/9] fix: prefer response model in observability exporters (#543) #### Overview Prefer response-derived LLM model names over requested model names when Relay exports observability data. This keeps ATIF, OpenTelemetry, and OpenInference aligned when a router, alias, or versioned endpoint returns a different model than the request asked for. - [x] I confirm this contribution is my own work, or I have the right to submit it under this project's license. - [x] I searched existing issues and open pull requests, and this does not duplicate existing work. #### Details This change fixes model-name precedence for LLM observability exports. - Updates `model_name_for_llm_event()` to prefer, in order, the normalized response model, manual response model, requested or profile model, normalized request model, and manual request fallback. - Updates the OpenTelemetry and OpenInference exporters to preserve one final end-time model-name attribute, preferring response-derived model data and keeping the start or request model only as fallback. - Updates ATIF pair resolution to prefer manual response models when normalized response data is unavailable. - Adds regression coverage for: - normalized response model overriding requested model - requested model remaining the fallback when no response model exists - manual response model overriding requested model across ATIF, OpenTelemetry, and OpenInference Validation run: - `cargo fmt --all` - `cargo clippy --workspace --all-targets -- -D warnings` - `cargo test -p nemo-relay helper_functions_cover_additional_openinference_branches --lib` - `cargo test -p nemo-relay helper_functions_cover_additional_otel_branches --lib` - `cargo test -p nemo-relay test_exporters_agree_on_model_name --lib` - `cargo test -p nemo-relay test_exporters_prefer_response_model_name_over_requested_model --lib` - `cargo test -p nemo-relay test_exporters_prefer_manual_response_model_name_over_requested_model --lib` - `just test-rust` - `just test-python` - `just test-go` - `just test-node` - `PATH="$HOME/.local/nemo-relay-tools/bin:$PATH" uv run pre-commit run --files crates/core/src/observability/atif.rs crates/core/src/observability/mod.rs crates/core/src/observability/openinference.rs crates/core/src/observability/otel.rs crates/core/tests/unit/observability/exporter_parity_tests.rs crates/core/tests/unit/observability/openinference_tests.rs crates/core/tests/unit/observability/otel_tests.rs` #### Where should the reviewer start? Start in `crates/core/src/observability/mod.rs` at `model_name_for_llm_event()`, then review how that precedence is applied in: - `crates/core/src/observability/atif.rs` - `crates/core/src/observability/otel.rs` - `crates/core/src/observability/openinference.rs` The most useful regression coverage is in: - `crates/core/tests/unit/observability/exporter_parity_tests.rs` #### Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to) - Relates to: none Authors: - Maryam Najafian (https://github.com/mnajafian-nv) Approvers: - Will Killian (https://github.com/willkill07) URL: https://github.com/NVIDIA/NeMo-Relay/pull/543 --- crates/core/src/observability/atif.rs | 1 + crates/core/src/observability/mod.rs | 17 ++-- .../core/src/observability/openinference.rs | 9 ++ crates/core/src/observability/otel.rs | 9 ++ .../observability/exporter_parity_tests.rs | 99 +++++++++++++++++++ .../unit/observability/openinference_tests.rs | 26 +++++ .../tests/unit/observability/otel_tests.rs | 25 +++++ 7 files changed, 176 insertions(+), 10 deletions(-) diff --git a/crates/core/src/observability/atif.rs b/crates/core/src/observability/atif.rs index 0a908db60..83c28dd24 100644 --- a/crates/core/src/observability/atif.rs +++ b/crates/core/src/observability/atif.rs @@ -2705,6 +2705,7 @@ fn normalized_response_model_name(event: &Event) -> Option { fn effective_model_for_pair(start: &Event, end: &Event) -> Option { normalized_response_model_name(end) + .or_else(|| manual::model_name_from_manual_llm_output(end.output()).map(ToOwned::to_owned)) .or_else(|| { start .model_name() diff --git a/crates/core/src/observability/mod.rs b/crates/core/src/observability/mod.rs index 2edd20043..9f1ecb7d2 100644 --- a/crates/core/src/observability/mod.rs +++ b/crates/core/src/observability/mod.rs @@ -523,27 +523,24 @@ pub(crate) fn merge_usage( } pub(crate) fn model_name_for_llm_event(event: &crate::api::event::Event) -> Option { - if let Some(model_name) = event.model_name() { - return Some(model_name.to_string()); - } if event.category().map(|category| category.as_str()) != Some("llm") { return None; } + let manual_response_model = + manual::model_name_from_manual_llm_output(event.output()).map(ToOwned::to_owned); + let manual_request_model = + manual::model_name_from_manual_llm_output(event.input()).map(ToOwned::to_owned); event .normalized_llm_response() .and_then(|response| response.as_ref().model.clone()) + .or(manual_response_model) + .or_else(|| event.model_name().map(ToOwned::to_owned)) .or_else(|| { event .normalized_llm_request() .and_then(|request| request.as_ref().model.clone()) }) - .or_else(|| { - event - .output() - .or_else(|| event.input()) - .and_then(|payload| manual::model_name_from_manual_llm_output(Some(payload))) - .map(ToOwned::to_owned) - }) + .or(manual_request_model) } #[cfg(any(feature = "otel", feature = "openinference"))] diff --git a/crates/core/src/observability/openinference.rs b/crates/core/src/observability/openinference.rs index 28f122eef..9b02f10c7 100644 --- a/crates/core/src/observability/openinference.rs +++ b/crates/core/src/observability/openinference.rs @@ -562,6 +562,7 @@ fn build_grpc_metadata(headers: &HashMap) -> Result struct ActiveSpan { span: Span, span_context: SpanContext, + start_model_name: Option, projected_attributes: Vec, } @@ -656,6 +657,7 @@ impl OpenInferenceEventProcessor { self.remove_completed_span_context(event.uuid()); let parent_context = self.parent_context(event); let is_trace_root = !parent_context.span().span_context().is_valid(); + let start_model_name = model_name_for_llm_event(event); let mut span = self .tracer .span_builder(span_name(event)) @@ -665,6 +667,9 @@ impl OpenInferenceEventProcessor { .with_span_id(relay_span_id(event.uuid())) .start_with_context(&self.tracer, &parent_context); let mut attributes = start_attributes(event); + if start_model_name.is_some() { + attributes.retain(|attribute| attribute.key.as_str() != oi::llm::MODEL_NAME.as_str()); + } if is_trace_root { push_session_identity_attributes(&mut attributes, event); } @@ -676,6 +681,7 @@ impl OpenInferenceEventProcessor { ActiveSpan { span, span_context, + start_model_name, projected_attributes, }, ); @@ -688,6 +694,9 @@ impl OpenInferenceEventProcessor { self.record_completed_span_context(event.uuid(), active_span.span_context.clone()); super::set_span_status_from_event_metadata(&mut active_span.span, event); let mut attributes = end_attributes(event); + if let Some(model_name) = model_name_for_llm_event(event).or(active_span.start_model_name) { + attributes.push(KeyValue::new(oi::llm::MODEL_NAME, model_name)); + } if !self.attribute_mappings.is_empty() { let mut projected_attributes = active_span.projected_attributes; projected_attributes.extend(attributes.iter().cloned()); diff --git a/crates/core/src/observability/otel.rs b/crates/core/src/observability/otel.rs index 5fab7868a..d9b19a1d7 100644 --- a/crates/core/src/observability/otel.rs +++ b/crates/core/src/observability/otel.rs @@ -555,6 +555,7 @@ fn build_grpc_metadata(headers: &HashMap) -> Result struct ActiveSpan { span: Span, span_context: SpanContext, + start_model_name: Option, projected_attributes: Vec, } @@ -649,6 +650,7 @@ impl OtelEventProcessor { self.remove_completed_span_context(event.uuid()); let parent_context = self.parent_context(event); let is_trace_root = !parent_context.span().span_context().is_valid(); + let start_model_name = model_name_for_llm_event(event); let mut span = self .tracer .span_builder(span_name(event)) @@ -658,6 +660,9 @@ impl OtelEventProcessor { .with_span_id(relay_span_id(event.uuid())) .start_with_context(&self.tracer, &parent_context); let mut attributes = start_attributes(event); + if start_model_name.is_some() { + attributes.retain(|attribute| attribute.key.as_str() != "nemo_relay.model_name"); + } if is_trace_root { push_session_identity_attributes(&mut attributes, event); } @@ -669,6 +674,7 @@ impl OtelEventProcessor { ActiveSpan { span, span_context, + start_model_name, projected_attributes, }, ); @@ -682,6 +688,9 @@ impl OtelEventProcessor { super::set_span_status_from_event_metadata(&mut active_span.span, event); let mut attributes = end_attributes(event); + if let Some(model_name) = model_name_for_llm_event(event).or(active_span.start_model_name) { + attributes.push(KeyValue::new("nemo_relay.model_name", model_name)); + } if !self.attribute_mappings.is_empty() { let mut projected_attributes = active_span.projected_attributes; projected_attributes.extend(attributes.iter().cloned()); diff --git a/crates/core/tests/unit/observability/exporter_parity_tests.rs b/crates/core/tests/unit/observability/exporter_parity_tests.rs index 832fe5fcc..4f064356a 100644 --- a/crates/core/tests/unit/observability/exporter_parity_tests.rs +++ b/crates/core/tests/unit/observability/exporter_parity_tests.rs @@ -527,6 +527,105 @@ fn test_exporters_agree_on_model_name() { ); } +#[test] +fn test_exporters_prefer_response_model_name_over_requested_model() { + let exports = run_llm_scenario( + chat_request_content("requested-model"), + chat_response_output("response-model"), + ); + + assert_eq!( + exports.agent_step().model_name.as_deref(), + Some("response-model") + ); + assert_eq!( + exports + .otel_attrs("model-call") + .get("nemo_relay.model_name"), + Some(&"response-model".to_string()) + ); + assert_eq!( + exports + .openinference_attrs("model-call") + .get("llm.model_name"), + Some(&"response-model".to_string()) + ); +} + +#[test] +fn test_exporters_fall_back_to_requested_model_when_response_model_is_missing() { + let exports = run_llm_scenario( + chat_request_content("requested-model"), + json!({ + "id": "chatcmpl-no-model", + "object": "chat.completion", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "hello"}, + "finish_reason": "stop" + }] + }), + ); + + assert_eq!( + exports + .otel_attrs("model-call") + .get("nemo_relay.model_name"), + Some(&"requested-model".to_string()) + ); + assert_eq!( + exports + .openinference_attrs("model-call") + .get("llm.model_name"), + Some(&"requested-model".to_string()) + ); +} + +#[test] +fn test_exporters_prefer_manual_response_model_name_over_requested_model() { + let uuid = Uuid::now_v7(); + let start = llm_event_with_model( + ScopeCategory::Start, + uuid, + "model-call", + json!({"prompt": "manual prompt"}), + "requested-model", + ); + let end = llm_event_with_model( + ScopeCategory::End, + uuid, + "model-call", + json!({ + "content": "manual answer", + "model": "response-model" + }), + "requested-model", + ); + assert!( + end.normalized_llm_response().is_none(), + "payload must exercise the manual response-model fallback, not a codec", + ); + + let exports = export_through_all_exporters(&[start, end]); + + assert_eq!( + exports.agent_step().model_name.as_deref(), + Some("response-model") + ); + assert_eq!( + exports + .otel_attrs("model-call") + .get("nemo_relay.model_name"), + Some(&"response-model".to_string()) + ); + assert_eq!( + exports + .openinference_attrs("model-call") + .get("llm.model_name"), + Some(&"response-model".to_string()) + ); +} + // =================================================================== // Tool-call parity // =================================================================== diff --git a/crates/core/tests/unit/observability/openinference_tests.rs b/crates/core/tests/unit/observability/openinference_tests.rs index 537b43774..03882d884 100644 --- a/crates/core/tests/unit/observability/openinference_tests.rs +++ b/crates/core/tests/unit/observability/openinference_tests.rs @@ -3343,6 +3343,32 @@ fn helper_functions_cover_additional_openinference_branches() { raw_model_attributes.get(oi::llm::MODEL_NAME.as_str()), Some(&"raw-model".to_string()) ); + let response_model_end = Event::Scope(ScopeEvent::new( + BaseEvent::builder() + .name("chat") + .data(json!({ + "id": "chatcmpl-response-model", + "model": "response-model", + "choices": [{ + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop" + }] + })) + .build(), + ScopeCategory::End, + Vec::new(), + EventCategory::llm(), + Some( + CategoryProfile::builder() + .model_name("requested-model") + .build(), + ), + )); + let response_model_attributes = attr_map(&common_attributes(&response_model_end)); + assert_eq!( + response_model_attributes.get(oi::llm::MODEL_NAME.as_str()), + Some(&"response-model".to_string()) + ); assert_eq!( llm_attributes.get("openinference.metadata.phase"), Some(&"done".to_string()) diff --git a/crates/core/tests/unit/observability/otel_tests.rs b/crates/core/tests/unit/observability/otel_tests.rs index dde05684c..c162e8746 100644 --- a/crates/core/tests/unit/observability/otel_tests.rs +++ b/crates/core/tests/unit/observability/otel_tests.rs @@ -1778,6 +1778,31 @@ fn helper_functions_cover_additional_otel_branches() { raw_model_attributes.get("nemo_relay.model_name"), Some(&"raw-model".to_string()) ); + let response_model_event = make_scope_event_with_profile( + ScopeCategory::End, + Uuid::now_v7(), + None, + "chat", + ScopeType::Llm, + Some(json!({ + "id": "chatcmpl-response-model", + "model": "response-model", + "choices": [{ + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop" + }] + })), + Some( + CategoryProfile::builder() + .model_name("requested-model") + .build(), + ), + ); + let response_model_attributes = attr_map(&common_attributes(&response_model_event)); + assert_eq!( + response_model_attributes.get("nemo_relay.model_name"), + Some(&"response-model".to_string()) + ); let tool_event = Event::Scope(ScopeEvent::new( BaseEvent::builder() From 0bcbb775c8b2ecb17fd8a1feae5241133d1e3b03 Mon Sep 17 00:00:00 2001 From: Will Killian <2007799+willkill07@users.noreply.github.com> Date: Thu, 23 Jul 2026 13:18:25 -0400 Subject: [PATCH 2/9] chore: prepare 0.7 release notes (#544) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit #### Overview Prepare the repository for the upcoming 0.7 release. - [x] I confirm this contribution is my own work, or I have the right to submit it under this project's license. - [x] I searched existing issues and open pull requests, and this does not duplicate existing work. #### Details - Stop nightly alpha tag creation for `release/0.6`. - Replace 0.6-specific release-note highlights and fixed issues with 0.7 placeholders. - Carry forward known issues and update the migration reference to 0.6 to 0.7. #### Where should the reviewer start? Review `.github/nightly-alpha-branches.yaml` for the nightly-tag change, then `docs/about-nemo-relay/release-notes/index.mdx` for the 0.7 release-note structure. #### Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to) - Relates to: none ## Summary by CodeRabbit * **Documentation** * Updated NVIDIA NeMo Relay release notes to document version **0.7** (including placeholders for **Highlights** and **Fixed Known Issues**). * Updated migration guidance to reference upgrading **0.6 → 0.7**, replacing detailed **0.6** content with a **0.7** placeholder. * **Chores** * Adjusted nightly alpha branch configuration to remove the **release/0.6** branch, leaving **main** only. Authors: - Will Killian (https://github.com/willkill07) Approvers: - Eric Evans II (https://github.com/ericevans-nv) URL: https://github.com/NVIDIA/NeMo-Relay/pull/544 --- .github/nightly-alpha-branches.yaml | 1 - docs/about-nemo-relay/release-notes/index.mdx | 61 +++--------- docs/reference/migration-guides.mdx | 92 +------------------ 3 files changed, 15 insertions(+), 139 deletions(-) diff --git a/.github/nightly-alpha-branches.yaml b/.github/nightly-alpha-branches.yaml index 354590d1d..04baa981f 100644 --- a/.github/nightly-alpha-branches.yaml +++ b/.github/nightly-alpha-branches.yaml @@ -3,4 +3,3 @@ branches: - main - - release/0.6 diff --git a/docs/about-nemo-relay/release-notes/index.mdx b/docs/about-nemo-relay/release-notes/index.mdx index 84ff8e406..8d9550fa2 100644 --- a/docs/about-nemo-relay/release-notes/index.mdx +++ b/docs/about-nemo-relay/release-notes/index.mdx @@ -1,7 +1,7 @@ --- title: "Release Notes for NVIDIA NeMo Relay" sidebar-title: "Release Notes" -description: "Review highlights, compatibility updates, fixed known issues, and current known issues for NVIDIA NeMo Relay 0.6." +description: "Review highlights, compatibility updates, fixed known issues, and current known issues for NVIDIA NeMo Relay 0.7." template-library-version: "1.0.0" position: 6 --- @@ -25,32 +25,13 @@ This is the Release Notes template. Document one release version per page and us This page contains the release notes for [NVIDIA NeMo Relay](/about-nemo-relay/overview). -## Release 0.6 +## Release 0.7 -NVIDIA NeMo Relay 0.6 strengthens local coding-agent observability, event -sanitization, observability exports, and dynamic plugin lifecycle management. +NVIDIA NeMo Relay 0.7 release notes are in preparation. ### Highlights -- Coding agents now connect through `nemo-relay mcp`, which starts or adopts a - shared authenticated gateway before hooks or routed provider traffic arrive. - Codex, Claude Code, and Hermes clients can share the gateway and coordinate - recovery before idle shutdown. -- Relay adds global, scope-local, and plugin-installed sanitizers for mark, - scope-start, and scope-end events. The PII redaction plugin now supports - ordered, composable profiles and the opt-in `trajectory_context` preset. -- ATOF configuration version 2 supports multiple independently configured file - and stream sinks. OpenTelemetry and OpenInference now use typed attributes, - and trace exporters can project selected marks as tool spans. -- Rust, Python, and Node.js embedding hosts can own native and worker dynamic - plugin lifecycles. Experimental Go and C entry points expose the same - source-first lifecycle. -- The coding-agent gateway supports lossless request annotations for Anthropic - Messages, OpenAI Chat Completions, and OpenAI Responses generation routes. -- The CLI adds recursive plugin configuration editing and configurable - human-readable or JSONL operational logging. -- The opt-in Switchyard integration can validate and route buffered or - streaming provider requests through a separately managed Decision API. +- _Highlights will be added for the 0.7 release._ ### Support Matrix and Compatibility Updates @@ -59,32 +40,14 @@ supported platforms and architectures, worker runtimes, coding agents, and integrations. It also records current limitations, including platform-specific worker requirements. -This release requires migration work for persistent coding-agent installations, -ATOF configuration, typed observability attributes, annotated-request plugin -consumers, and managed LLM streams. Refer to the [Migration -Guides](/reference/migration-guides) for upgrade actions from 0.5 to 0.6. - -### Fixed Known Issues in 0.6 - -- Coding-agent gateway generation routes now decode request annotations for - Anthropic Messages, OpenAI Chat Completions, and OpenAI Responses while - preserving unchanged nested fields, explicit nulls, and provider - representations. -- Installed coding-agent sessions now acquire the shared gateway through MCP - before hooks or routed provider traffic, preventing cold-start loss and - coordinating concurrent startup and recovery. -- Transparent Codex routing now recognizes `at-...` access tokens and sends - them to the ChatGPT Codex backend without rewriting unrelated bearer tokens - or provider API keys. -- Relay now asks Codex to prefer the readable legacy multi-agent path during - managed runs and restores the user's prior setting during uninstall. -- ATIF now reports the normalized provider response model when a routed or - translated call runs on a model that differs from the request. -- The NeMo Flow migration skill now skips credential-bearing dotenv files and - symbolic links, requires exact project-root confirmation, and refuses write - mode for filesystem roots and home directories. - -## Known Issues in 0.6 +Migration guidance for upgrading from 0.6 to 0.7 will be added to the +[Migration Guides](/reference/migration-guides) before the release. + +### Fixed Known Issues in 0.7 + +- _Fixed issues will be added for the 0.7 release._ + +## Known Issues in 0.7 - Go and the raw C FFI remain experimental and source-first. Generated API pages focus on Rust, Python, and Node.js. diff --git a/docs/reference/migration-guides.mdx b/docs/reference/migration-guides.mdx index 4431926e1..debc2cda7 100644 --- a/docs/reference/migration-guides.mdx +++ b/docs/reference/migration-guides.mdx @@ -6,100 +6,14 @@ position: 6 {/* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 */} -Use this page to plan an upgrade from NeMo Relay 0.5 to 0.6. It groups the +Use this page to plan an upgrade from NeMo Relay 0.6 to 0.7. It will group the actions from the release notes by the surface you operate. If you skip one or more releases, review the migration guides and release notes for every intervening release in sequence. -## Upgrade to NeMo Relay 0.6 +## Upgrade to NeMo Relay 0.7 -### Coding-Agent Integrations - -Reinstall every persistent Codex, Claude Code, and Hermes Agent integration: - -```bash -nemo-relay install --force -``` - -Or refresh all detected supported hosts: - -```bash -nemo-relay install all --force -``` - -Version 0.5 installations do not use the 0.6 MCP-owned gateway lifecycle, -agent-owned hooks, generation fencing, or user-scoped configuration. Confirm -that the host satisfies the current minimum version before reinstalling: -Claude Code 2.1.121, Codex CLI 0.143.0, or Hermes Agent 0.18.2. Then run -`nemo-relay doctor --plugin ` to verify the refreshed installation. - -### Dynamic Plugins and Workers - -Rebuild Rust native plugins and Rust `grpc-v1` workers that consume annotated -LLM requests against NeMo Relay 0.6. Update Python workers to the 0.6 worker -SDK and declare: - -```toml -[compat] -relay = ">=0.6,<1.0" -``` - -The annotation and LLM request-intercept outcome envelopes gained fields and -variants. A plugin that registers an LLM request intercept cannot claim a -compatibility range that admits Relay 0.5. Node.js, Go, and raw C FFI callbacks -continue to receive JSON, but their consumers must not exhaustively match role -or component discriminator strings. Go and raw C FFI remain experimental and -source-first. - -### Streaming Consumers - -Close a managed LLM stream when you stop consuming it early so Relay can stop -the producer and emit the interrupted end event. Replace direct Rust -`LlmJsonStream` construction with its constructors, handle the result from Go -`LlmStream.Close`, use Python `await stream.aclose()`, Node.js -`await stream.close()`, or call `nemo_relay_stream_close` before freeing a C -stream. - -### Exporters and Observability Queries - -Update ATOF configuration to version 2 and replace the legacy output fields -with the tagged `atof.sinks` list. Direct ATOF exporter construction now takes -one typed file or stream sink. - -Update OpenTelemetry and OpenInference queries, dashboards, and processors to -the typed attribute paths. The former raw `*_json` payload attributes are no -longer emitted. Use `attribute_mappings` only when an older key must continue -to be available. - -ATIF no longer models marks as synthetic system steps. Use ATOF for canonical -mark data or enable the OpenTelemetry or OpenInference `mark_projection = -"tool"` visualization when a trace viewer needs visible mark nodes. Review -any consumers that assume `step.model_name` always identifies the requested -model; it now uses the effective response model where available. - -### Middleware, Sanitizers, and Runtime APIs - -Update PII redaction configuration to the composable `profiles` form when you -need ordered policies. Existing single-policy configuration continues to work, -but it cannot be combined with `profiles`. The optional -`trajectory_context` preset changes observability payloads only; it does not -change provider requests or client-visible responses. - -Sanitizer failures now drop the affected event. If your policy requires -fail-open behavior, handle the failure inside the sanitizer and return only the -safe fields. - -For Rust consumers, update exhaustive matches and direct struct literals for -the expanded public enums and types. Prefer the current builders and typed sink -constructors over direct literals where available. - -### Skills and Automation - -Replace retired 0.5 skill directories with the task-oriented public entry -points, including `nemo-relay-install`, `nemo-relay-get-started`, the -`nemo-relay-instrument-*` skills, and `nemo-relay-plugin-*` skills. The -[NeMo Relay User Skills](https://github.com/NVIDIA/NeMo-Relay/blob/main/skills/README.md) -catalog lists the current paths. +Migration guidance will be added before the 0.7 release. ## Related Release Information From 70551bc1dff14c50706a77d0c7484a502ef8d29a Mon Sep 17 00:00:00 2001 From: Maryam Najafian Date: Thu, 23 Jul 2026 11:06:55 -0700 Subject: [PATCH 3/9] fix: emit OpenInference system prompts as input messages (#542) #### Overview Align OpenInference export behavior with the expected attribute contract by stopping Relay from writing prompt text into `llm.system`. System content is now exported as ordered `llm.input_messages` entries instead. - [x] I confirm this contribution is my own work, or I have the right to submit it under this project's license. - [x] I searched existing issues and open pull requests, and this does not duplicate existing work. #### Details - Remove `llm.system` emission from both replay-payload projection and annotated request projection in the OpenInference exporter. - Preserve replay `systemPrompt` content by exporting it as the leading `llm.input_messages.` entry with role `system`. - Preserve `AnnotatedLlmRequest.instructions` as the leading system input message when present. - Preserve annotated `Message::System` content in `llm.input_messages` instead of the legacy `llm.system` field. - Harden replay input projection so `prompt` still falls back when replay `messages` is empty or every replay message is skipped, and skip incomplete replay message objects unless they provide both a string `role` and displayable `content`. - Update unit coverage to assert that `llm.system` is absent, system and user messages are exported in the expected order, replay fallback remains contiguous, and migrated system-message content is preserved. Validation: - `cargo fmt --all` - `just test-rust` - `cargo clippy --workspace --all-targets -- -D warnings` - `just test-python` - `just test-go` - `just test-node` - `uv run pre-commit run --files crates/core/src/observability/openinference.rs crates/core/tests/unit/observability/openinference_tests.rs` #### Where should the reviewer start? Start in `crates/core/src/observability/openinference.rs`, especially the request projection path around `push_llm_request_attributes()` and `push_replay_input_messages()`. The most relevant regression coverage is in `crates/core/tests/unit/observability/openinference_tests.rs`. #### Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to) - Relates to: none ## Summary by CodeRabbit * **Bug Fixes** * Improved OpenInference tracing by representing the system prompt as the first `llm.input_messages` entry with `role: system`, instead of emitting a separate `llm.system` scalar. * Ensured consistent `llm.input_messages` indexing/ordering for both annotated requests (from instructions) and replay inputs, with correct shifting of subsequent messages. * Added more robust replay handling: only well-formed message entries (with required role/content) are projected, and fallback behavior works for empty or incomplete replay inputs. * **Tests** * Updated and expanded OpenInference flattened-attribute tests to match the new system-message semantics and indexing. Authors: - Maryam Najafian (https://github.com/mnajafian-nv) Approvers: - Will Killian (https://github.com/willkill07) URL: https://github.com/NVIDIA/NeMo-Relay/pull/542 --- .../core/src/observability/openinference.rs | 72 ++++-- .../unit/observability/openinference_tests.rs | 233 +++++++++++++++++- 2 files changed, 275 insertions(+), 30 deletions(-) diff --git a/crates/core/src/observability/openinference.rs b/crates/core/src/observability/openinference.rs index 9b02f10c7..362015853 100644 --- a/crates/core/src/observability/openinference.rs +++ b/crates/core/src/observability/openinference.rs @@ -1009,9 +1009,6 @@ fn push_llm_request_attributes(attributes: &mut Vec, event: &Event) { if let Some(provider) = input.get("provider").and_then(Json::as_str) { attributes.push(KeyValue::new(oi::llm::PROVIDER, provider.to_string())); } - if let Some(system) = input.get("systemPrompt").and_then(display_text_from_json) { - attributes.push(KeyValue::new(oi::llm::SYSTEM, system)); - } push_replay_input_messages(attributes, input); return; } @@ -1047,13 +1044,16 @@ fn push_annotated_request_attributes( attributes: &mut Vec, request: &AnnotatedLlmRequest, ) { - if let Some(system) = request.system_prompt() { - attributes.push(KeyValue::new(oi::llm::SYSTEM, system.to_string())); - } if let Some(params) = request.params.as_ref().and_then(to_json_string) { attributes.push(KeyValue::new(oi::llm::INVOCATION_PARAMETERS, params)); } - push_annotated_input_messages(attributes, &request.messages); + let mut next_index = 0usize; + if let Some(instructions) = request.instructions.as_ref().and_then(message_content_text) { + push_message_role(attributes, "llm.input_messages", next_index, "system"); + push_message_text_value(attributes, "llm.input_messages", next_index, instructions); + next_index += 1; + } + push_annotated_input_messages(attributes, &request.messages, next_index); if let Some(tools) = request.tools.as_deref() { push_annotated_tools(attributes, tools); } @@ -1102,8 +1102,13 @@ fn push_optimization_attributes( crate::observability::push_common_optimization_attributes(attributes, summary); } -fn push_annotated_input_messages(attributes: &mut Vec, messages: &[Message]) { - for (index, message) in messages.iter().enumerate() { +fn push_annotated_input_messages( + attributes: &mut Vec, + messages: &[Message], + start_index: usize, +) { + for (offset, message) in messages.iter().enumerate() { + let index = start_index + offset; let role = match message { Message::System { .. } => "system", Message::Developer { .. } => "developer", @@ -1238,36 +1243,51 @@ fn is_openclaw_replay_payload(content: &serde_json::Map) -> bool { } fn push_replay_input_messages(attributes: &mut Vec, input: &Json) { + let mut next_index = 0usize; + if let Some(system_prompt) = input.get("systemPrompt").and_then(display_text_from_json) { + push_message_role(attributes, "llm.input_messages", next_index, "system"); + attributes.push(KeyValue::new( + format!("llm.input_messages.{next_index}.message.content"), + system_prompt, + )); + next_index += 1; + } if let Some(messages) = input.get("messages").and_then(Json::as_array) { - for (index, message) in messages.iter().enumerate() { - push_replay_input_message(attributes, index, message); + let first_message_index = next_index; + for message in messages { + if push_replay_input_message(attributes, next_index, message) { + next_index += 1; + } + } + if next_index > first_message_index { + return; } - return; } if let Some(prompt) = input.get("prompt").and_then(display_text_from_json) { - push_message_role(attributes, "llm.input_messages", 0, "user"); + push_message_role(attributes, "llm.input_messages", next_index, "user"); attributes.push(KeyValue::new( - "llm.input_messages.0.message.content", + format!("llm.input_messages.{next_index}.message.content"), prompt, )); } } -fn push_replay_input_message(attributes: &mut Vec, index: usize, message: &Json) { +fn push_replay_input_message(attributes: &mut Vec, index: usize, message: &Json) -> bool { let Some(object) = message.as_object() else { - return; + return false; + }; + let Some(role) = object.get("role").and_then(Json::as_str) else { + return false; + }; + let Some(text) = object.get("content").and_then(display_text_from_json) else { + return false; }; - if !object.contains_key("role") && !object.contains_key("content") { - return; - } - let role = object.get("role").and_then(Json::as_str).unwrap_or("user"); push_message_role(attributes, "llm.input_messages", index, role); - if let Some(text) = object.get("content").and_then(display_text_from_json) { - attributes.push(KeyValue::new( - format!("llm.input_messages.{index}.message.content"), - text, - )); - } + attributes.push(KeyValue::new( + format!("llm.input_messages.{index}.message.content"), + text, + )); + true } fn push_replay_response_attributes(attributes: &mut Vec, output: &Json) { diff --git a/crates/core/tests/unit/observability/openinference_tests.rs b/crates/core/tests/unit/observability/openinference_tests.rs index 03882d884..22a3ca14d 100644 --- a/crates/core/tests/unit/observability/openinference_tests.rs +++ b/crates/core/tests/unit/observability/openinference_tests.rs @@ -405,6 +405,38 @@ fn sample_openinference_annotated_request() -> AnnotatedLlmRequest { } } +fn sample_openinference_annotated_request_with_instructions() -> AnnotatedLlmRequest { + AnnotatedLlmRequest { + instructions: Some(MessageContent::Text("Use concise answers.".to_string())), + api_specific: None, + messages: vec![Message::User { + content: MessageContent::Text("Search docs.".to_string()), + name: None, + }], + model: Some("gpt-4o".to_string()), + params: Some(GenerationParams { + temperature: Some(0.2), + max_tokens: Some(128), + top_p: None, + stop: None, + }), + tools: Some(vec![ToolDefinition::Function { + function: FunctionDefinition { + name: "search_docs".to_string(), + description: Some("Search the docs corpus.".to_string()), + parameters: Some(json!({ + "type": "object", + "properties": {"query": {"type": "string"}} + })), + strict: None, + extra: serde_json::Map::new(), + }, + extra: serde_json::Map::new(), + }]), + ..empty_annotated_request() + } +} + fn sample_openinference_annotated_response() -> AnnotatedLlmResponse { AnnotatedLlmResponse { message: Some(MessageContent::Text("I will search docs.".to_string())), @@ -1705,11 +1737,17 @@ fn openclaw_replay_payloads_emit_flattened_openinference_llm_attributes() { assert_eq!(spans.len(), 1); let attributes = attr_map(&spans[0].attributes); assert_attr(&attributes, "llm.provider", "nvidia-inference"); - assert_attr(&attributes, "llm.system", "Use reliable sources."); - assert_attr(&attributes, "llm.input_messages.0.message.role", "user"); + assert!(!attributes.contains_key("llm.system")); + assert_attr(&attributes, "llm.input_messages.0.message.role", "system"); assert_attr( &attributes, "llm.input_messages.0.message.content", + "Use reliable sources.", + ); + assert_attr(&attributes, "llm.input_messages.1.message.role", "user"); + assert_attr( + &attributes, + "llm.input_messages.1.message.content", "Find the answer.", ); assert_attr( @@ -1743,6 +1781,132 @@ fn openclaw_replay_payloads_emit_flattened_openinference_llm_attributes() { assert_no_attr_contains(&attributes, "secret-token"); } +#[test] +fn openclaw_replay_payloads_fall_back_to_prompt_when_replay_messages_are_empty() { + let (provider, exporter) = make_provider(); + let mut processor = + OpenInferenceEventProcessor::new(provider.clone(), "test-scope".to_string()); + let uuid = Uuid::now_v7(); + + processor.process(&make_start_event( + uuid, + None, + "openclaw-model-call", + ScopeType::Llm, + Some(json!({ + "headers": {"authorization": "Bearer secret-token"}, + "content": { + "provider": "nvidia-inference", + "model": "claude-sonnet-4", + "systemPrompt": "Use reliable sources.", + "prompt": "Find the answer.", + "messages": [], + "placeholderRequest": false, + "source": "openclaw.llm_output" + } + })), + )); + processor.process(&make_end_event( + uuid, + None, + "openclaw-model-call", + ScopeType::Llm, + Some(json!({ + "role": "assistant", + "content": "I will search.", + "openclaw": { + "duration_ms": 42 + } + })), + )); + + processor.force_flush().unwrap(); + + let spans = exporter.get_finished_spans().unwrap(); + assert_eq!(spans.len(), 1); + let attributes = attr_map(&spans[0].attributes); + assert!(!attributes.contains_key("llm.system")); + assert_attr(&attributes, "llm.input_messages.0.message.role", "system"); + assert_attr( + &attributes, + "llm.input_messages.0.message.content", + "Use reliable sources.", + ); + assert_attr(&attributes, "llm.input_messages.1.message.role", "user"); + assert_attr( + &attributes, + "llm.input_messages.1.message.content", + "Find the answer.", + ); + assert!(!attributes.contains_key("llm.input_messages.2.message.role")); + assert!(!attributes.contains_key("llm.input_messages.2.message.content")); +} + +#[test] +fn openclaw_replay_payloads_fall_back_to_prompt_when_replay_messages_are_incomplete() { + let (provider, exporter) = make_provider(); + let mut processor = + OpenInferenceEventProcessor::new(provider.clone(), "test-scope".to_string()); + let uuid = Uuid::now_v7(); + + processor.process(&make_start_event( + uuid, + None, + "openclaw-model-call", + ScopeType::Llm, + Some(json!({ + "headers": {"authorization": "Bearer secret-token"}, + "content": { + "provider": "nvidia-inference", + "model": "claude-sonnet-4", + "systemPrompt": "Use reliable sources.", + "prompt": "Find the answer.", + "messages": [ + {"content": "content without role"}, + {"role": "user"}, + {"role": "user", "content": ""} + ], + "placeholderRequest": false, + "source": "openclaw.llm_output" + } + })), + )); + processor.process(&make_end_event( + uuid, + None, + "openclaw-model-call", + ScopeType::Llm, + Some(json!({ + "role": "assistant", + "content": "I will search.", + "openclaw": { + "duration_ms": 42 + } + })), + )); + + processor.force_flush().unwrap(); + + let spans = exporter.get_finished_spans().unwrap(); + assert_eq!(spans.len(), 1); + let attributes = attr_map(&spans[0].attributes); + assert!(!attributes.contains_key("llm.system")); + assert_attr(&attributes, "llm.input_messages.0.message.role", "system"); + assert_attr( + &attributes, + "llm.input_messages.0.message.content", + "Use reliable sources.", + ); + assert_attr(&attributes, "llm.input_messages.1.message.role", "user"); + assert_attr( + &attributes, + "llm.input_messages.1.message.content", + "Find the answer.", + ); + assert!(!attributes.contains_key("llm.input_messages.2.message.role")); + assert!(!attributes.contains_key("llm.input_messages.2.message.content")); +} + #[test] fn openclaw_replay_tool_call_alias_fields_emit_openinference_attributes() { let (provider, exporter) = make_provider(); @@ -4219,8 +4383,13 @@ fn annotated_llm_payloads_emit_flattened_openinference_message_and_tool_attribut let spans = exporter.get_finished_spans().unwrap(); assert_eq!(spans.len(), 1); let attributes = attr_map(&spans[0].attributes); - assert_attr(&attributes, "llm.system", "Use concise answers."); + assert!(!attributes.contains_key("llm.system")); assert_attr(&attributes, "llm.input_messages.0.message.role", "system"); + assert_attr( + &attributes, + "llm.input_messages.0.message.content", + "Use concise answers.", + ); assert_attr(&attributes, "llm.input_messages.1.message.role", "user"); assert_attr( &attributes, @@ -4301,7 +4470,7 @@ fn annotated_input_projection_covers_extended_roles_and_native_text() { }, ]; let mut attributes = Vec::new(); - push_annotated_input_messages(&mut attributes, &messages); + push_annotated_input_messages(&mut attributes, &messages, 0); let attributes = attr_map(&attributes); assert_attr( &attributes, @@ -4354,6 +4523,62 @@ fn annotated_input_projection_covers_extended_roles_and_native_text() { ); } +#[test] +fn annotated_llm_instructions_emit_leading_system_input_message() { + let (provider, exporter) = make_provider(); + let mut processor = + OpenInferenceEventProcessor::new(provider.clone(), "test-scope".to_string()); + let uuid = Uuid::now_v7(); + + processor.process(&make_scope_event_with_profile( + ScopeCategory::Start, + uuid, + None, + "annotated-chat-with-instructions", + ScopeType::Llm, + None, + Some( + CategoryProfile::builder() + .annotated_request(Arc::new( + sample_openinference_annotated_request_with_instructions(), + )) + .build(), + ), + )); + processor.process(&make_scope_event_with_profile( + ScopeCategory::End, + uuid, + None, + "annotated-chat-with-instructions", + ScopeType::Llm, + None, + Some( + CategoryProfile::builder() + .annotated_response(Arc::new(sample_openinference_annotated_response())) + .build(), + ), + )); + + processor.force_flush().unwrap(); + + let spans = exporter.get_finished_spans().unwrap(); + assert_eq!(spans.len(), 1); + let attributes = attr_map(&spans[0].attributes); + assert!(!attributes.contains_key("llm.system")); + assert_attr(&attributes, "llm.input_messages.0.message.role", "system"); + assert_attr( + &attributes, + "llm.input_messages.0.message.content", + "Use concise answers.", + ); + assert_attr(&attributes, "llm.input_messages.1.message.role", "user"); + assert_attr( + &attributes, + "llm.input_messages.1.message.content", + "Search docs.", + ); +} + #[test] fn hermes_exact_api_payloads_emit_openinference_text_usage_and_metadata() { let (provider, exporter) = make_provider(); From f69802c7feb290ec080b1dc23b3a9ffc6263f238 Mon Sep 17 00:00:00 2001 From: Hans Arnholm Date: Thu, 23 Jul 2026 13:00:30 -0700 Subject: [PATCH 4/9] feat(observability): route ATIF files with scope metadata (#505) #### Overview Allow ATIF filename templates to derive path-safe route fragments from top-level scope metadata, with concise literal fallbacks when metadata is optional. - [x] I confirm this contribution is my own work, or I have the right to submit it under this project's license. - [x] I searched existing issues and open pull requests, and this does not duplicate existing work. #### Details - Extend the existing `filename_template` setting with nested `{metadata.}` placeholders and optional `{metadata.:-fallback}` values. - Render one filename for local, S3, and HTTP storage so every existing destination uses the same metadata-aware routing behavior. - Accept only path-safe relative metadata fragments and skip only the affected trajectory when a required value is missing, non-string, or unsafe. - Preserve `{session_id}` as the required per-trajectory identity and log structured `atif_destination_render_failed` warnings for render failures. - Add focused routing, fallback, path-safety, and recovery coverage, plus user documentation and observability skill guidance. Validation: - `cargo fmt --all` - `just test-rust` - `cargo clippy --workspace --all-targets -- -D warnings` - `just test-python` (535 passed) - `just test-go` - `just test-node` (277 passed) - `just docs` - `uv run pre-commit run --all-files` - Built `nemo-relay` for `x86_64-unknown-linux-musl`, confirmed it is a static x86-64 ELF, and executed it in a Linux/amd64 Alpine container. This change is additive and has no breaking configuration changes. #### Where should the reviewer start? Start with `render_atif_filename` in `crates/core/src/observability/plugin_component.rs`, then review `atif_filename_template_routes_by_metadata_and_skips_invalid_paths` for the end-to-end behavior and path-safety decision. #### Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to) - Relates to: none ## Summary by CodeRabbit - **New Features** - Enhanced ATIF `filename_template` to support `{metadata.:-fallback}` placeholders for routing trajectories into nested, metadata-derived directories. - **Bug Fixes** - Stricter template validation and runtime rendering: malformed templates or unsafe/missing metadata now skip only the affected trajectory and log `atif_destination_render_failed`. - **Documentation** - Updated ATIF docs and references with the metadata placeholder syntax, `:-` fallback rules, path-safety requirements, and behavior when rendering fails. - **Tests** - Added/updated unit tests for unclosed placeholder rejection, safe-vs-unsafe metadata interpolation, traversal-like value handling, and updated destination preparation behavior. Authors: - Hans Arnholm (https://github.com/cypres) Approvers: - Will Killian (https://github.com/willkill07) URL: https://github.com/NVIDIA/NeMo-Relay/pull/505 --- .../src/observability/plugin_component.rs | 149 +++++++++++++++-- .../observability/plugin_component_tests.rs | 158 +++++++++++++++++- docs/configure-plugins/observability/atif.mdx | 39 ++++- .../references/atif.md | 10 ++ 4 files changed, 337 insertions(+), 19 deletions(-) diff --git a/crates/core/src/observability/plugin_component.rs b/crates/core/src/observability/plugin_component.rs index eebe80e36..e52f78cd3 100644 --- a/crates/core/src/observability/plugin_component.rs +++ b/crates/core/src/observability/plugin_component.rs @@ -256,8 +256,10 @@ pub struct AtifSectionConfig { /// [`storage`]: Self::storage #[serde(default, skip_serializing_if = "Option::is_none")] pub output_directory: Option, - /// Filename template. `{session_id}` is replaced with the top-level trajectory scope UUID. - /// When [`storage`] is non-empty, the rendered filename is appended to each backend's key prefix. + /// Filename template. `{session_id}` is replaced with the top-level trajectory scope UUID, and + /// `{metadata.:-fallback}` placeholders use path-safe strings from the top-level scope + /// metadata or the optional literal fallback. When [`storage`] is non-empty, the rendered + /// filename is appended to each backend's key prefix. /// /// [`storage`]: Self::storage #[serde(default = "default_atif_filename_template")] @@ -846,11 +848,8 @@ fn register_atif_dispatcher( section: AtifSectionConfig, ctx: &mut PluginRegistrationContext, ) -> PluginResult<()> { - if !section.filename_template.contains("{session_id}") { - return Err(PluginError::InvalidConfig( - "ATIF filename_template must contain '{session_id}'".to_string(), - )); - } + validate_atif_filename_template(§ion.filename_template) + .map_err(PluginError::InvalidConfig)?; let mut storage_vec = Vec::with_capacity(section.storage.len()); for (index, entry) in section.storage.iter().enumerate() { @@ -1171,9 +1170,22 @@ impl AtifDispatcher { // subscriber is attached after that start event has already been // emitted. let session_id = event.uuid().to_string(); + let (filename, local_path) = match self.prepare_destination(&session_id, event.metadata()) { + Ok(destination) => destination, + Err(error) => { + log::warn!( + target: "nemo_relay.observability", + event = "atif_destination_render_failed", + plugin_kind = OBSERVABILITY_PLUGIN_KIND, + exporter = "atif", + session_id = session_id.as_str(); + "ATIF destination rendering failed: {error}" + ); + return None; + } + }; let exporter = AtifExporter::new(session_id.clone(), self.agent_info()); (exporter.subscriber())(event); - let (filename, local_path) = self.prepare_destination(&session_id); let correlation = AtifCorrelation::from_event(event); self.scope_owners.insert(event.uuid(), event.uuid()); self.agents.insert( @@ -1360,13 +1372,14 @@ impl AtifDispatcher { } } - fn prepare_destination(&self, session_id: &str) -> (String, Option) { - let filename = self - .config - .filename_template - .replace("{session_id}", session_id); + fn prepare_destination( + &self, + session_id: &str, + metadata: Option<&Json>, + ) -> Result<(String, Option), String> { + let filename = render_atif_filename(&self.config.filename_template, session_id, metadata)?; if !self.config.storage.is_empty() { - return (filename, None); + return Ok((filename, None)); } let directory = self .config @@ -1374,7 +1387,7 @@ impl AtifDispatcher { .clone() .unwrap_or_else(default_output_directory); let path = directory.join(&filename); - (filename, Some(path)) + Ok((filename, Some(path))) } fn sink_targets(&self) -> Vec { @@ -1393,6 +1406,108 @@ impl AtifDispatcher { } } +fn is_valid_atif_metadata_selector(selector: &str) -> bool { + !selector.is_empty() + && selector.split('.').all(|segment| { + !segment.is_empty() + && segment + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_')) + }) +} + +fn parse_atif_metadata_expression(expression: &str) -> Result<(&str, Option<&str>), String> { + let (selector, fallback) = expression + .split_once(":-") + .map_or((expression, None), |(key, value)| (key, Some(value))); + if !is_valid_atif_metadata_selector(selector) { + return Err(format!( + "ATIF filename_template metadata placeholder '{{metadata.{selector}}}' must contain a dot-separated path of ASCII letters, digits, '-' or '_'" + )); + } + Ok((selector, fallback)) +} + +fn validate_atif_filename_template(template: &str) -> Result<(), String> { + const PREFIX: &str = "{metadata."; + + if !template.contains("{session_id}") { + return Err("ATIF filename_template must contain '{session_id}'".to_string()); + } + + let mut cursor = 0; + while let Some(relative_start) = template[cursor..].find(PREFIX) { + let selector_start = cursor + relative_start + PREFIX.len(); + let end = template[selector_start..] + .find('}') + .map(|relative_end| selector_start + relative_end) + .ok_or_else(|| { + "ATIF filename_template contains an unclosed metadata placeholder".to_string() + })?; + let (_, fallback) = parse_atif_metadata_expression(&template[selector_start..end])?; + if let Some(fallback) = fallback + && !is_safe_atif_metadata_path(fallback) + { + return Err(format!( + "ATIF filename_template fallback '{fallback}' must be a path-safe relative fragment" + )); + } + cursor = end + 1; + } + Ok(()) +} + +fn render_atif_filename( + template: &str, + session_id: &str, + metadata: Option<&Json>, +) -> Result { + const PREFIX: &str = "{metadata."; + + let mut rendered = template.replace("{session_id}", session_id); + let mut cursor = 0; + while let Some(relative_start) = rendered[cursor..].find(PREFIX) { + let start = cursor + relative_start; + let selector_start = start + PREFIX.len(); + let end = rendered[selector_start..] + .find('}') + .map(|relative_end| selector_start + relative_end) + .ok_or_else(|| { + "ATIF filename_template contains an unclosed metadata placeholder".to_string() + })?; + let expression = rendered[selector_start..end].to_string(); + let (selector, fallback) = parse_atif_metadata_expression(&expression)?; + let value = selector + .split('.') + .fold(metadata, |value, segment| value?.get(segment)) + .and_then(Json::as_str) + .or(fallback) + .ok_or_else(|| { + format!( + "filename_template placeholder '{{metadata.{selector}}}' must resolve to a string" + ) + })?; + if !is_safe_atif_metadata_path(value) { + return Err(format!( + "metadata path '{selector}' must be a path-safe relative fragment" + )); + } + rendered.replace_range(start..=end, value); + cursor = start + value.len(); + } + Ok(rendered) +} + +fn is_safe_atif_metadata_path(value: &str) -> bool { + !value.is_empty() + && value.split('/').all(|segment| { + !matches!(segment, "" | "." | "..") + && segment.bytes().all(|byte| { + byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_' | b'.' | b'~') + }) + }) +} + fn atif_dispatcher_subscriber( manager: Arc>, subscriber_prefix: String, @@ -2376,14 +2491,14 @@ fn validate_atif_values( policy: &ConfigPolicy, section: &AtifSectionConfig, ) { - if !section.filename_template.contains("{session_id}") { + if let Err(message) = validate_atif_filename_template(§ion.filename_template) { push_policy_diag( diagnostics, policy.unsupported_value, "observability.unsupported_value", Some("atif".to_string()), Some("filename_template".to_string()), - "ATIF filename_template must contain '{session_id}'".to_string(), + message, ); } for (index, storage) in section.storage.iter().enumerate() { diff --git a/crates/core/tests/unit/observability/plugin_component_tests.rs b/crates/core/tests/unit/observability/plugin_component_tests.rs index 128d473aa..6b17f869d 100644 --- a/crates/core/tests/unit/observability/plugin_component_tests.rs +++ b/crates/core/tests/unit/observability/plugin_component_tests.rs @@ -704,6 +704,35 @@ fn unknown_fields_and_bad_values_follow_policy() { assert!(ignore_report.diagnostics.is_empty()); } +#[test] +fn atif_filename_template_syntax_is_rejected_before_activation() { + let _guard = crate::observability::test_mutex().lock().unwrap(); + reset_runtime(); + + let valid_report = validate_plugin_config(&plugin_config(json!({ + "atif": { + "filename_template": "{metadata.workflow_id:-unassigned}/trajectory-{session_id}.json" + } + }))); + assert!(!valid_report.has_errors()); + + let malformed = "trajectory-{session_id}.json/{metadata.tenant"; + let invalid_report = validate_plugin_config(&plugin_config(json!({ + "atif": {"filename_template": malformed} + }))); + assert!(invalid_report.diagnostics.iter().any(|diag| { + diag.field.as_deref() == Some("filename_template") + && diag.message.contains("unclosed metadata placeholder") + })); + + let error = futures::executor::block_on(initialize_plugins_exact(plugin_config(json!({ + "policy": {"unsupported_value": "ignore"}, + "atif": {"enabled": true, "filename_template": malformed} + })))) + .unwrap_err(); + assert!(error.to_string().contains("unclosed metadata placeholder")); +} + #[test] fn invalid_shapes_and_strict_policy_are_reported() { let _guard = crate::observability::test_mutex().lock().unwrap(); @@ -1293,6 +1322,57 @@ fn atif_defaults_create_one_file_per_top_level_agent() { assert!(!second_serialized.contains("first-agent")); } +#[test] +fn atif_filename_template_routes_by_metadata_and_skips_invalid_paths() { + let _guard = crate::observability::test_mutex().lock().unwrap(); + reset_runtime(); + let dir = temp_dir("observability-atif-metadata-template"); + + let config = plugin_config(json!({ + "atif": { + "enabled": true, + "output_directory": dir, + "filename_template": "{metadata.routing.artifact_path}/trajectory-{session_id}.json" + } + })); + futures::executor::block_on(initialize_plugins_exact(config)).unwrap(); + + let invalid = crate::api::scope::push_scope( + PushScopeParams::builder() + .name("invalid-metadata-path-agent") + .scope_type(ScopeType::Agent) + .metadata(json!({"routing": {"artifact_path": "../escape"}})) + .build(), + ) + .unwrap(); + pop(&invalid); + + let valid = crate::api::scope::push_scope( + PushScopeParams::builder() + .name("valid-metadata-path-agent") + .scope_type(ScopeType::Agent) + .metadata(json!({"routing": {"artifact_path": "tenant-a/session-123"}})) + .build(), + ) + .unwrap(); + pop(&valid); + + clear_plugin_configuration().unwrap(); + let invalid_filename = format!("trajectory-{}.json", invalid.uuid); + assert!( + !dir.join(&invalid_filename).exists() + && !dir.join("../escape").join(&invalid_filename).exists(), + "unsafe metadata path should not produce a trajectory file" + ); + assert!( + dir.join(format!( + "tenant-a/session-123/trajectory-{}.json", + valid.uuid + )) + .exists() + ); +} + #[test] fn atif_routes_global_descendant_events_by_parent_uuid() { let _guard = crate::observability::test_mutex().lock().unwrap(); @@ -1763,7 +1843,7 @@ fn write_atif_reports_missing_local_path_and_unregistered_remote_sink() { #[test] fn atif_dispatcher_default_output_path_uses_current_directory() { let dispatcher = AtifDispatcher::new(AtifSectionConfig::default()); - let (filename, local_path) = dispatcher.prepare_destination("session-1"); + let (filename, local_path) = dispatcher.prepare_destination("session-1", None).unwrap(); assert_eq!(filename, "nemo-relay-atif-session-1.json"); assert_eq!( local_path.unwrap(), @@ -1773,6 +1853,82 @@ fn atif_dispatcher_default_output_path_uses_current_directory() { ); } +#[test] +fn atif_metadata_template_values_must_be_safe_path_fragments() { + assert!( + validate_atif_filename_template( + "{metadata.workflow_id:-unassigned}/trajectory-{session_id}.json" + ) + .is_ok() + ); + assert_eq!( + render_atif_filename( + "{metadata.workflow_id:-unassigned}/trajectory-{session_id}.json", + "scope-id", + None + ) + .unwrap(), + "unassigned/trajectory-scope-id.json" + ); + + assert!(is_safe_atif_metadata_path( + "tenant-a/team_1.session-123~retry" + )); + + for value in [ + "", + "/absolute", + "trailing/", + "double//slash", + ".", + "../escape", + "tenant/../escape", + r"tenant\session", + "tenant name", + "tenant:session", + ] { + assert!( + !is_safe_atif_metadata_path(value), + "metadata path should be rejected: {value:?}" + ); + } + + let dispatcher = AtifDispatcher::new(AtifSectionConfig { + filename_template: "{metadata.artifact_path}/trajectory-{session_id}.json".to_string(), + ..AtifSectionConfig::default() + }); + assert!(dispatcher.prepare_destination("session-1", None).is_err()); + let non_string = json!({"artifact_path": 123}); + assert!( + dispatcher + .prepare_destination("session-1", Some(&non_string)) + .is_err() + ); + + for template in [ + "{metadata.}/trajectory-{session_id}.json", + "{metadata.tenant..id}/trajectory-{session_id}.json", + "{metadata.tenant/trajectory-{session_id}.json", + "{metadata.{tenant}}/trajectory-{session_id}.json", + "{metadata.missing:-../escape}/trajectory-{session_id}.json", + ] { + assert!( + validate_atif_filename_template(template).is_err(), + "template should fail configuration validation: {template:?}" + ); + let dispatcher = AtifDispatcher::new(AtifSectionConfig { + filename_template: template.to_string(), + ..AtifSectionConfig::default() + }); + assert!( + dispatcher + .prepare_destination("session-1", Some(&json!({"tenant": "tenant-a"}))) + .is_err(), + "template should be rejected: {template:?}" + ); + } +} + #[test] fn atif_payload_merges_correlation_with_existing_trajectory_extra() { let agent_uuid = Uuid::now_v7(); diff --git a/docs/configure-plugins/observability/atif.mdx b/docs/configure-plugins/observability/atif.mdx index 3acb8b227..c358404aa 100644 --- a/docs/configure-plugins/observability/atif.mdx +++ b/docs/configure-plugins/observability/atif.mdx @@ -68,9 +68,46 @@ The following table describes the top-level ATIF settings: | `tool_definitions` | Omitted | Optional ATIF tool metadata. | | `extra` | Omitted | Optional ATIF agent metadata. | | `output_directory` | Current working directory | Directory containing trajectory files. Ignored when `storage` is non-empty. | -| `filename_template` | `nemo-relay-atif-{session_id}.json` | Must contain `{session_id}`. | +| `filename_template` | `nemo-relay-atif-{session_id}.json` | Must contain `{session_id}`. Can contain `{metadata.}` placeholders for metadata-based routing. | | `storage` | Omitted | Optional list of remote storage destinations. When non-empty, trajectories are uploaded to every configured backend instead of being written locally. Refer to [Remote Storage](#remote-storage). | +### Metadata-Based Paths + +Use `{metadata.}` placeholders in `filename_template` to route +trajectories with top-level scope metadata. Dots select nested fields, and +templates can use multiple placeholders. For example: + +```toml +[components.config.atif] +enabled = true +output_directory = "logs" +filename_template = "{metadata.atif_prefix:-unassigned}/trajectory-{session_id}.json" +``` + +With scope metadata `{"atif_prefix":"tenant-a/session-123"}`, Relay writes +`logs/tenant-a/session-123/trajectory-.json`. The template must +still contain `{session_id}`. + +Use `:-` to provide a literal fallback when metadata can be absent, for example +`{metadata.atif_prefix:-unassigned}`. Relay also uses the fallback when the +metadata value is not a string. + +Each metadata placeholder must resolve to a string containing a non-empty, +relative path fragment. Slash-separated segments can contain ASCII letters, +digits, `-`, `_`, `.`, and `~`. Relay rejects empty segments, `.` and `..` +segments, absolute paths, backslashes, spaces, and other characters. If a +placeholder has no fallback and is missing or non-string, or if its resolved +value is unsafe, Relay skips that trajectory and logs +`atif_destination_render_failed`. The rendered filename applies to local, S3, +and HTTP storage in the same way as a static filename. + +The CLI gateway parses `x-nemo-relay-session-metadata` as JSON and merges it +into the top-level scope metadata: + +```http +x-nemo-relay-session-metadata: {"atif_prefix":"tenant-a/session-123"} +``` + ## Remote Storage Use `storage` when local trace files are not durable across sessions, such as in diff --git a/skills/nemo-relay-plugin-observability/references/atif.md b/skills/nemo-relay-plugin-observability/references/atif.md index febc22725..6e8f0226a 100644 --- a/skills/nemo-relay-plugin-observability/references/atif.md +++ b/skills/nemo-relay-plugin-observability/references/atif.md @@ -60,6 +60,16 @@ live OTLP spans. - Response codecs can improve LLM end annotations, but they do not change the caller-visible LLM response. +## Plugin-Managed File Routing + +- `filename_template` must contain `{session_id}` and can use + `{metadata.}` to select nested string values from top-level scope + metadata. Use `{metadata.:-fallback}` when the value is optional. +- Metadata values must be relative path fragments using ASCII letters, digits, + `-`, `_`, `.`, `~`, and safe `/`-separated segments. +- Missing, non-string, or unsafe metadata values skip only the affected + trajectory and produce an `atif_destination_render_failed` warning. + ## Checklist - [ ] Session and agent metadata chosen From 2b8a88445b9f23c310817d6ba99a36f4988f25e7 Mon Sep 17 00:00:00 2001 From: Maryam Najafian Date: Thu, 23 Jul 2026 15:02:38 -0700 Subject: [PATCH 5/9] fix: avoid duplicate ATIF user steps on continuations (#545) #### Overview Fixes an ATIF exporter bug where a continued LLM request after tool work could emit a second duplicate `user` step for the same logical turn. The exporter now keeps one `user` step per fresh user turn, preserves same-turn continuation requests on the matching `agent` step, and avoids carrying an unpaired continuation request onto a later unrelated step. - [x] I confirm this contribution is my own work, or I have the right to submit it under this project's license. - [x] I searched existing issues and open pull requests, and this does not duplicate existing work. #### Details - Update the ATIF LLM-start mapping so only fresh user turns emit a `user` step. - Treat same-turn continuations as part of the matching `agent` step instead of repeating the user message. - Preserve the full continuation request in `Step.extra.llm_request` so request fidelity is not lost. - Add turn-state detection for raw chat requests, OpenAI Responses-style inputs, and annotated requests. - Make raw message turn-state detection content-aware so user messages that carry only tool-use or tool-result continuation content stay in the same logical turn. - Keep OpenAI Responses typed non-user continuation items in the same logical turn instead of re-emitting a duplicate `user` step. - Only treat LLM ends with event data as pairable, which prevents a stashed continuation request from leaking onto a later unrelated `agent` step. - Add regression coverage for OpenAI Responses continuations with `function_call_output` and `shell_call`, Anthropic tool-result continuation, and an empty-end continuation followed by a fresh user turn. - Update the existing full agent-loop test to reflect the corrected ATIF step sequence. **Validation:** - `cargo fmt --all` - `cargo test -p nemo-relay observability::atif::tests::` - `cargo test -p nemo-relay-cli --lib -- --nocapture` - `uv run pre-commit run --files crates/core/src/observability/atif.rs crates/core/tests/unit/atif_tests.rs` #### Where should the reviewer start? Start in `crates/core/src/observability/atif.rs`, especially the turn-state detection around `llm_start_user_step_message()`, `chat_messages_turn_state()`, and the pairability lookup in `EventLookupMaps::from_events_with_correlation_events()`. The most important regression coverage is in `crates/core/tests/unit/atif_tests.rs`. #### Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to) - Relates to: none Authors: - Maryam Najafian (https://github.com/mnajafian-nv) Approvers: - Eric Evans II (https://github.com/ericevans-nv) - Will Killian (https://github.com/willkill07) URL: https://github.com/NVIDIA/NeMo-Relay/pull/545 --- crates/core/src/observability/atif.rs | 236 ++++++++++++- crates/core/tests/unit/atif_tests.rs | 468 +++++++++++++++++++++++++- 2 files changed, 684 insertions(+), 20 deletions(-) diff --git a/crates/core/src/observability/atif.rs b/crates/core/src/observability/atif.rs index 83c28dd24..5571dfd81 100644 --- a/crates/core/src/observability/atif.rs +++ b/crates/core/src/observability/atif.rs @@ -18,7 +18,7 @@ //! //! | NeMo Relay Event | ATIF Step | Notes | //! |-----------------|-------------------------|--------------------------------------| -//! | LLM Start | `user` step | Messages extracted from LlmRequest | +//! | LLM Start | `user` step | Fresh user turns only; same-turn continuations stay on the agent step | //! | LLM End | `agent` step | Response content, tool_calls promoted| //! | Tool Start | *(skipped)* | tool_calls come from LLM End instead | //! | Tool End | agent observation | Correlated by `source_call_id` | @@ -38,7 +38,7 @@ use uuid::Uuid; use crate::api::event::{Event, EventNormalizationExt}; use crate::api::runtime::EventSubscriberFn; use crate::api::subscriber::flush_subscribers; -use crate::codec::request::{AnnotatedLlmRequest, Message, MessageContent}; +use crate::codec::request::{AnnotatedLlmRequest, ContentPart, Message, MessageContent}; use crate::codec::response::AnnotatedLlmResponse; use crate::error::Result; use crate::json::Json; @@ -287,6 +287,12 @@ pub struct AtifStepExtra { pub tool_invocations: Option>, } +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum RequestTurnState { + FreshUser, + Continuation, +} + /// A complete ATIF trajectory. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct AtifTrajectory { @@ -930,9 +936,10 @@ fn extract_reasoning_content(output: &Json) -> Option { /// Extract the latest user-facing message from an LLM request payload. /// /// LLM start inputs typically contain `{ "messages": [...], "model": "...", -/// "max_tokens": ..., "tools": [...], "stream": ... }`. For the ATIF user step -/// we emit a schema-compatible message value (string or content-part array) -/// and preserve the full LLM request in `Step.extra.llm_request`. +/// "max_tokens": ..., "tools": [...], "stream": ... }`. For ATIF we emit a +/// schema-compatible message value (string or content-part array) and preserve +/// the full LLM request in `Step.extra.llm_request`, either on the user step or +/// on the matching agent step for same-turn continuations. /// /// Returns the latest user message content if present, a prompt if present, or /// a stringified representation of the input as a last resort. @@ -964,6 +971,135 @@ fn extract_user_messages(input: &Json) -> Json { atif_content_value(input) } +fn llm_start_user_step_message(event: &Event, input: &Json, has_paired_end: bool) -> Option { + let turn_state = event + .annotated_request() + .and_then(|request| request_turn_state_from_annotation(request.as_ref())) + .or_else(|| request_turn_state_from_raw(input)); + + if has_paired_end && matches!(turn_state, Some(RequestTurnState::Continuation)) { + return None; + } + + event + .annotated_request() + .and_then(|request| atif_message_from_annotated_request(request.as_ref())) + .or_else(|| Some(extract_user_messages(input))) +} + +fn request_turn_state_from_raw(input: &Json) -> Option { + let object = input.as_object()?; + if let Some(messages) = object.get("messages").and_then(Json::as_array) + && let Some(state) = chat_messages_turn_state(messages) + { + return Some(state); + } + if let Some(input_items) = object.get("input") + && let Some(state) = openai_responses_turn_state(input_items) + { + return Some(state); + } + object.get("prompt").map(|_| RequestTurnState::FreshUser) +} + +fn request_turn_state_from_annotation(request: &AnnotatedLlmRequest) -> Option { + request + .messages + .iter() + .rev() + .find_map(annotated_message_turn_state) +} + +fn chat_messages_turn_state(messages: &[Json]) -> Option { + messages + .iter() + .rev() + .filter_map(Json::as_object) + .find_map(raw_chat_message_turn_state) +} + +fn raw_chat_message_turn_state( + message: &serde_json::Map, +) -> Option { + match message.get("role").and_then(Json::as_str) { + Some("system" | "developer") => None, + Some("user") | None => Some(match message.get("content") { + Some(content) if raw_content_starts_new_turn(content) => RequestTurnState::FreshUser, + Some(_) => RequestTurnState::Continuation, + None => RequestTurnState::FreshUser, + }), + Some(_) => Some(RequestTurnState::Continuation), + } +} + +fn raw_content_starts_new_turn(content: &Json) -> bool { + match content { + Json::String(_) => true, + Json::Array(parts) => { + if parts.iter().any(raw_content_part_is_user_input) { + return true; + } + !parts.iter().any(raw_content_part_is_tool_continuation) + } + Json::Object(_) => { + raw_content_part_is_user_input(content) + || !raw_content_part_is_tool_continuation(content) + } + _ => true, + } +} + +fn raw_content_part_is_user_input(part: &Json) -> bool { + matches!( + part.get("type").and_then(Json::as_str), + Some( + "text" + | "input_text" + | "image_url" + | "image" + | "input_image" + | "audio" + | "input_audio" + | "file" + | "document" + ) + ) +} + +fn raw_content_part_is_tool_continuation(part: &Json) -> bool { + matches!( + part.get("type").and_then(Json::as_str), + Some("tool_result" | "tool_use") + ) +} + +fn openai_responses_turn_state(input: &Json) -> Option { + if input.is_string() { + return Some(RequestTurnState::FreshUser); + } + + input + .as_array()? + .iter() + .rev() + .filter_map(Json::as_object) + .find_map(openai_responses_item_turn_state) +} + +fn openai_responses_item_turn_state( + item: &serde_json::Map, +) -> Option { + match item.get("type").and_then(Json::as_str) { + Some("message") => match item.get("role").and_then(Json::as_str) { + Some("system" | "developer") => None, + Some("user") | None => Some(RequestTurnState::FreshUser), + Some(_) => Some(RequestTurnState::Continuation), + }, + Some(_) => Some(RequestTurnState::Continuation), + _ => None, + } +} + fn openai_responses_input_message(input: &Json) -> Option { if input.is_string() { return Some(atif_content_value(input)); @@ -1084,6 +1220,64 @@ fn atif_message_from_annotated_request(request: &AnnotatedLlmRequest) -> Option< } } +fn annotated_message_turn_state(message: &Message) -> Option { + match message { + Message::System { .. } | Message::Developer { .. } => None, + Message::User { content, .. } => Some(if annotated_content_starts_new_turn(content) { + RequestTurnState::FreshUser + } else { + RequestTurnState::Continuation + }), + Message::Assistant { .. } + | Message::Tool { .. } + | Message::Function { .. } + | Message::ToolCallItem { .. } + | Message::ToolResultItem { .. } => Some(RequestTurnState::Continuation), + Message::ProviderNative { value, .. } => provider_native_turn_state(value), + } +} + +fn annotated_content_starts_new_turn(content: &MessageContent) -> bool { + match content { + MessageContent::Text(_) => true, + MessageContent::Parts(parts) => { + let has_user_input = parts.iter().any(|part| { + matches!( + part, + ContentPart::Text { .. } + | ContentPart::ImageUrl { .. } + | ContentPart::Image { .. } + | ContentPart::Audio { .. } + | ContentPart::File { .. } + ) + }); + if has_user_input { + return true; + } + let has_tool_continuation = parts.iter().any(|part| { + matches!( + part, + ContentPart::ToolUse { .. } | ContentPart::ToolResult { .. } + ) + }); + !has_tool_continuation + } + } +} + +fn provider_native_turn_state(value: &Json) -> Option { + let object = value.as_object()?; + if object.get("type").is_some() { + return openai_responses_item_turn_state(object); + } + match object.get("role").and_then(Json::as_str) { + Some("system" | "developer") => None, + Some("user") => Some(RequestTurnState::FreshUser), + Some(_) => Some(RequestTurnState::Continuation), + None => None, + } +} + fn atif_message_from_annotated_response(response: &AnnotatedLlmResponse) -> Option { match &response.message { Some(MessageContent::Text(text)) => Some(Json::String(text.clone())), @@ -1380,6 +1574,7 @@ fn json_string_at(value: &Json, path: &[&str]) -> Option { struct EventLookupMaps { name_map: std::collections::HashMap, start_ts_map: std::collections::HashMap>, + llm_end_uuids: HashSet, llm_start_model_names: HashMap, tool_call_ids: std::collections::HashMap, suppressed_llm_events: HashSet, @@ -1407,6 +1602,7 @@ impl EventLookupMaps { ) -> Self { let mut name_map = std::collections::HashMap::new(); let mut start_ts_map = std::collections::HashMap::new(); + let mut llm_end_uuids = HashSet::new(); let mut llm_start_model_names = HashMap::new(); for event in events { if is_start_event(event) { @@ -1418,11 +1614,18 @@ impl EventLookupMaps { llm_start_model_names.insert(event.uuid(), model_name); } } + if event.scope_category() == Some(crate::api::event::ScopeCategory::End) + && event.category().map(|category| category.as_str()) == Some("llm") + && event.data().is_some() + { + llm_end_uuids.insert(event.uuid()); + } } let llm_dedupe = build_llm_dedupe(llm_dedupe_events); Self { name_map, start_ts_map, + llm_end_uuids, llm_start_model_names, tool_call_ids: build_tool_call_correlations(tool_correlation_events), suppressed_llm_events: llm_dedupe.suppressed_events, @@ -1928,6 +2131,7 @@ struct PendingAgentStep { step_idx: Option, ancestry: Option, invocation: Option, + llm_request: Option, llm_response: Option, tool_ancestry: Vec, tool_invocations: Vec, @@ -1947,7 +2151,7 @@ impl PendingAgentStep { let extra = AtifStepExtra { ancestry, invocation: self.invocation.take(), - llm_request: None, + llm_request: self.llm_request.take(), llm_response: self.llm_response.take(), event_payload: None, tool_ancestry: std::mem::take(&mut self.tool_ancestry), @@ -1977,6 +2181,10 @@ impl PendingAgentStep { self.tool_call_order = tool_call_order; } + fn stash_llm_request(&mut self, llm_request: Json) { + self.llm_request = Some(llm_request); + } + fn push_tool_metadata(&mut self, ancestry: AtifAncestry, invocation: AtifInvocationInfo) { self.tool_ancestry.push(ancestry); self.tool_invocations.push(invocation); @@ -2299,6 +2507,11 @@ impl StepConversionState { }; let content = unwrap_llm_request(input); self.current_reasoning_effort = extract_reasoning_effort(&content); + let has_paired_end = lookups.llm_end_uuids.contains(&event.uuid()); + let Some(message) = llm_start_user_step_message(event, &content, has_paired_end) else { + self.current_agent.stash_llm_request(content); + return; + }; let extra = AtifStepExtra { ancestry: build_ancestry(event, &lookups.name_map), invocation: None, @@ -2311,10 +2524,7 @@ impl StepConversionState { self.steps.push(AtifStep { step_id: 0, source: "user".to_string(), - message: event - .annotated_request() - .and_then(|request| atif_message_from_annotated_request(request)) - .unwrap_or_else(|| extract_user_messages(&content)), + message, timestamp: Some(event.timestamp().to_rfc3339()), model_name: None, reasoning_effort: None, @@ -3382,8 +3592,10 @@ fn events_to_steps_for_agent( /// Mapping logic: /// 1. Sort events by timestamp. /// 2. For each LLM pair: -/// - Start event → user step (message = extracted `messages` array from -/// unwrapped LlmRequest content, stripping `max_tokens`/`model`/etc.) +/// - Start event → user step when the request begins a fresh user turn +/// - Start events that continue the same turn after tool work stash +/// `llm_request` on the matching agent step instead of repeating the user +/// message /// - End event → agent step (message = extracted content, metrics from /// token_usage, tool_calls promoted to AtifToolCall entries with parsed /// JSON arguments) diff --git a/crates/core/tests/unit/atif_tests.rs b/crates/core/tests/unit/atif_tests.rs index 10f203bc3..050134f3a 100644 --- a/crates/core/tests/unit/atif_tests.rs +++ b/crates/core/tests/unit/atif_tests.rs @@ -4890,7 +4890,8 @@ fn test_exporter_user_message_extraction() { #[test] fn test_exporter_full_agent_loop() { // Simulate a complete agent loop: LLM→tool_calls→observations→LLM→final answer - // This should produce 5 steps: user, agent+tool_calls, merged obs, user, agent + // The second request continues the same user turn after tool work, so it + // should not create a duplicate user step. let exporter = AtifExporter::new("session-1".to_string(), make_agent_info()); let llm1_uuid = Uuid::now_v7(); let llm2_uuid = Uuid::now_v7(); @@ -4984,8 +4985,8 @@ fn test_exporter_full_agent_loop() { let trajectory = exporter.export().unwrap(); assert_atif_v17_shape(&trajectory); - // Expected: user, agent+tool_calls+observations, user, agent - assert_eq!(trajectory.steps.len(), 4); + // Expected: user, agent+tool_calls+observations, agent + assert_eq!(trajectory.steps.len(), 3); assert_eq!(trajectory.steps[0].source, "user"); assert_eq!(trajectory.steps[0].step_id, 1); @@ -4997,17 +4998,23 @@ fn test_exporter_full_agent_loop() { assert_eq!(tcs[0].function_name, "get_weather"); assert_eq!(tcs[1].function_name, "get_population"); - assert_eq!(trajectory.steps[2].source, "user"); - assert_eq!(trajectory.steps[2].step_id, 3); let obs = trajectory.steps[1].observation.as_ref().unwrap(); assert_eq!(obs.results.len(), 2); - assert_eq!(trajectory.steps[3].source, "agent"); - assert_eq!(trajectory.steps[3].step_id, 4); + assert_eq!(trajectory.steps[2].source, "agent"); + assert_eq!(trajectory.steps[2].step_id, 3); assert_eq!( - trajectory.steps[3].message, + trajectory.steps[2].message, json!("The weather in SF is 62°F and foggy. Population is 873,965.") ); + let final_extra: AtifStepExtra = + serde_json::from_value(trajectory.steps[2].extra.clone().unwrap()).unwrap(); + let llm_request = final_extra.llm_request.unwrap(); + assert_eq!( + llm_request["messages"][0]["content"], + json!("What is the weather and population of SF?") + ); + assert_eq!(llm_request["messages"][3]["tool_call_id"], json!("c2")); // Final metrics should aggregate both LLM calls let fm = trajectory.final_metrics.as_ref().unwrap(); @@ -5015,6 +5022,451 @@ fn test_exporter_full_agent_loop() { assert_eq!(fm.total_completion_tokens, Some(80)); } +#[test] +fn test_exporter_skips_duplicate_user_step_for_openai_responses_tool_continuation() { + let exporter = AtifExporter::new("session-1".to_string(), make_agent_info()); + let llm1_uuid = Uuid::now_v7(); + let llm2_uuid = Uuid::now_v7(); + + let llm1_start = event_builder(llm1_uuid, EventType::Start) + .name("openai.responses") + .scope_type(ScopeType::Llm) + .input(json!({ + "model": "switchyard", + "input": [{ + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Fix pip"}] + }] + })) + .model_name("switchyard") + .build(); + let llm1_end = event_builder(llm1_uuid, EventType::End) + .name("openai.responses") + .scope_type(ScopeType::Llm) + .output(json!({ + "id": "resp_1", + "model": "switchyard", + "output": [{ + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "I will inspect the environment."}] + }] + })) + .model_name("switchyard") + .build(); + let llm2_start = event_builder(llm2_uuid, EventType::Start) + .name("openai.responses") + .scope_type(ScopeType::Llm) + .input(json!({ + "model": "switchyard", + "input": [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Fix pip"}] + }, + { + "type": "function_call", + "call_id": "call_1", + "name": "shell", + "arguments": "{}" + }, + { + "type": "function_call_output", + "call_id": "call_1", + "output": "pip is missing" + } + ] + })) + .model_name("switchyard") + .build(); + let llm2_end = event_builder(llm2_uuid, EventType::End) + .name("openai.responses") + .scope_type(ScopeType::Llm) + .output(json!({ + "id": "resp_2", + "model": "switchyard", + "output": [{ + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "I found that pip is missing."}] + }] + })) + .model_name("switchyard") + .build(); + + { + let mut state = exporter.state.lock().unwrap(); + state + .events + .extend([llm1_start, llm1_end, llm2_start, llm2_end]); + } + + let trajectory = exporter.export().unwrap(); + assert_atif_v17_shape(&trajectory); + assert_eq!(trajectory.steps.len(), 3); + let user_steps = trajectory + .steps + .iter() + .filter(|step| step.source == "user") + .collect::>(); + assert_eq!(user_steps.len(), 1); + assert_eq!(user_steps[0].message, json!("Fix pip")); + + let final_extra: AtifStepExtra = + serde_json::from_value(trajectory.steps[2].extra.clone().unwrap()).unwrap(); + let llm_request = final_extra.llm_request.unwrap(); + assert_eq!( + llm_request["input"][2]["type"], + json!("function_call_output") + ); + assert_eq!(llm_request["input"][2]["output"], json!("pip is missing")); +} + +#[test] +fn test_exporter_skips_duplicate_user_step_for_openai_responses_shell_call_continuation() { + let exporter = AtifExporter::new("session-1".to_string(), make_agent_info()); + let llm1_uuid = Uuid::now_v7(); + let llm2_uuid = Uuid::now_v7(); + let codec = OpenAIResponsesCodec; + + let llm1_start_input = json!({ + "model": "switchyard", + "input": [{ + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Fix pip"}] + }] + }); + let llm1_start = event_builder(llm1_uuid, EventType::Start) + .name("openai.responses") + .scope_type(ScopeType::Llm) + .input(llm1_start_input) + .model_name("switchyard") + .build(); + let llm1_end = event_builder(llm1_uuid, EventType::End) + .name("openai.responses") + .scope_type(ScopeType::Llm) + .output(json!({ + "id": "resp_1", + "model": "switchyard", + "output": [{ + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "I will inspect the environment."}] + }] + })) + .model_name("switchyard") + .build(); + + let llm2_start_input = json!({ + "model": "switchyard", + "input": [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Fix pip"}] + }, + { + "type": "shell_call", + "id": "sh_1", + "call_id": "sh_call", + "action": {"commands": ["which pip"]} + } + ] + }); + let llm2_start = event_builder(llm2_uuid, EventType::Start) + .name("openai.responses") + .scope_type(ScopeType::Llm) + .input(llm2_start_input.clone()) + .annotated_request( + codec + .decode(&LlmRequest { + headers: serde_json::Map::new(), + content: llm2_start_input, + }) + .unwrap(), + ) + .model_name("switchyard") + .build(); + let llm2_end = event_builder(llm2_uuid, EventType::End) + .name("openai.responses") + .scope_type(ScopeType::Llm) + .output(json!({ + "id": "resp_2", + "model": "switchyard", + "output": [{ + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "I found the missing pip binary."}] + }] + })) + .model_name("switchyard") + .build(); + + { + let mut state = exporter.state.lock().unwrap(); + state + .events + .extend([llm1_start, llm1_end, llm2_start, llm2_end]); + } + + let trajectory = exporter.export().unwrap(); + assert_atif_v17_shape(&trajectory); + assert_eq!(trajectory.steps.len(), 3); + let user_steps = trajectory + .steps + .iter() + .filter(|step| step.source == "user") + .collect::>(); + assert_eq!(user_steps.len(), 1); + assert_eq!(user_steps[0].message, json!("Fix pip")); + + let final_extra: AtifStepExtra = + serde_json::from_value(trajectory.steps[2].extra.clone().unwrap()).unwrap(); + let llm_request = final_extra.llm_request.unwrap(); + assert_eq!(llm_request["input"][1]["type"], json!("shell_call")); + assert_eq!(llm_request["input"][1]["call_id"], json!("sh_call")); +} + +#[test] +fn test_exporter_skips_duplicate_user_step_for_anthropic_tool_result_continuation() { + let exporter = AtifExporter::new("session-1".to_string(), make_agent_info()); + let llm1_uuid = Uuid::now_v7(); + let llm2_uuid = Uuid::now_v7(); + + let llm1_start = event_builder(llm1_uuid, EventType::Start) + .name("anthropic.messages") + .scope_type(ScopeType::Llm) + .input(json!({ + "model": "claude-sonnet-4", + "system": "Be concise.", + "messages": [{"role": "user", "content": "Find the file."}] + })) + .model_name("claude-sonnet-4") + .build(); + let llm1_end = event_builder(llm1_uuid, EventType::End) + .name("anthropic.messages") + .scope_type(ScopeType::Llm) + .output(json!({ + "id": "msg_01", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4", + "content": [ + {"type": "text", "text": "I will search for it."}, + { + "type": "tool_use", + "id": "toolu_01", + "name": "search", + "input": {"query": "README.md"} + } + ], + "stop_reason": "tool_use" + })) + .model_name("claude-sonnet-4") + .build(); + let llm2_start = event_builder(llm2_uuid, EventType::Start) + .name("anthropic.messages") + .scope_type(ScopeType::Llm) + .input(json!({ + "model": "claude-sonnet-4", + "system": "Be concise.", + "messages": [ + {"role": "user", "content": "Find the file."}, + {"role": "assistant", "content": [ + { + "type": "tool_use", + "id": "toolu_01", + "name": "search", + "input": {"query": "README.md"} + } + ]}, + {"role": "user", "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_01", + "content": "README.md", + "is_error": false + } + ]} + ] + })) + .model_name("claude-sonnet-4") + .build(); + let llm2_end = event_builder(llm2_uuid, EventType::End) + .name("anthropic.messages") + .scope_type(ScopeType::Llm) + .output(json!({ + "id": "msg_02", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4", + "content": [{"type": "text", "text": "README.md is the relevant file."}], + "stop_reason": "end_turn" + })) + .model_name("claude-sonnet-4") + .build(); + + { + let mut state = exporter.state.lock().unwrap(); + state + .events + .extend([llm1_start, llm1_end, llm2_start, llm2_end]); + } + + let trajectory = exporter.export().unwrap(); + assert_atif_v17_shape(&trajectory); + assert_eq!(trajectory.steps.len(), 3); + let user_steps = trajectory + .steps + .iter() + .filter(|step| step.source == "user") + .collect::>(); + assert_eq!(user_steps.len(), 1); + assert_eq!(user_steps[0].message, json!("Find the file.")); + + let final_extra: AtifStepExtra = + serde_json::from_value(trajectory.steps[2].extra.clone().unwrap()).unwrap(); + let llm_request = final_extra.llm_request.unwrap(); + assert_eq!( + llm_request["messages"][2]["content"][0]["type"], + json!("tool_result") + ); + assert_eq!( + llm_request["messages"][2]["content"][0]["tool_use_id"], + json!("toolu_01") + ); +} + +#[test] +fn test_exporter_does_not_stash_unpaired_continuation_request_on_next_agent_step() { + let exporter = AtifExporter::new("session-1".to_string(), make_agent_info()); + let llm1_uuid = Uuid::now_v7(); + let llm2_uuid = Uuid::now_v7(); + let llm3_uuid = Uuid::now_v7(); + + let llm1_start = event_builder(llm1_uuid, EventType::Start) + .name("openai.responses") + .scope_type(ScopeType::Llm) + .input(json!({ + "model": "switchyard", + "input": [{ + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Fix pip"}] + }] + })) + .model_name("switchyard") + .build(); + let llm1_end = event_builder(llm1_uuid, EventType::End) + .name("openai.responses") + .scope_type(ScopeType::Llm) + .output(json!({ + "id": "resp_1", + "model": "switchyard", + "output": [{ + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "I will inspect the environment."}] + }] + })) + .model_name("switchyard") + .build(); + let llm2_start = event_builder(llm2_uuid, EventType::Start) + .name("openai.responses") + .scope_type(ScopeType::Llm) + .input(json!({ + "model": "switchyard", + "input": [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Fix pip"}] + }, + { + "type": "function_call", + "call_id": "call_1", + "name": "shell", + "arguments": "{}" + }, + { + "type": "function_call_output", + "call_id": "call_1", + "output": "pip is missing" + } + ] + })) + .model_name("switchyard") + .build(); + let llm2_end = event_builder(llm2_uuid, EventType::End) + .name("openai.responses") + .scope_type(ScopeType::Llm) + .model_name("switchyard") + .build(); + let llm3_start = event_builder(llm3_uuid, EventType::Start) + .name("openai.responses") + .scope_type(ScopeType::Llm) + .input(json!({ + "model": "switchyard", + "input": [{ + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Try ensurepip instead"}] + }] + })) + .model_name("switchyard") + .build(); + let llm3_end = event_builder(llm3_uuid, EventType::End) + .name("openai.responses") + .scope_type(ScopeType::Llm) + .output(json!({ + "id": "resp_3", + "model": "switchyard", + "output": [{ + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Use ensurepip to restore pip."}] + }] + })) + .model_name("switchyard") + .build(); + + { + let mut state = exporter.state.lock().unwrap(); + state.events.extend([ + llm1_start, llm1_end, llm2_start, llm2_end, llm3_start, llm3_end, + ]); + } + + let trajectory = exporter.export().unwrap(); + assert_atif_v17_shape(&trajectory); + assert_eq!(trajectory.steps.len(), 5); + + let user_steps = trajectory + .steps + .iter() + .filter(|step| step.source == "user") + .collect::>(); + assert_eq!(user_steps.len(), 3); + assert_eq!(user_steps[0].message, json!("Fix pip")); + assert_eq!(user_steps[1].message, json!("Fix pip")); + assert_eq!(user_steps[2].message, json!("Try ensurepip instead")); + + let final_extra: AtifStepExtra = + serde_json::from_value(trajectory.steps[4].extra.clone().unwrap()).unwrap(); + assert!(final_extra.llm_request.is_none()); + + let llm3_user_extra: AtifStepExtra = + serde_json::from_value(trajectory.steps[3].extra.clone().unwrap()).unwrap(); + let llm3_request = llm3_user_extra.llm_request.unwrap(); + assert_eq!( + llm3_request["input"][0]["content"][0]["text"], + json!("Try ensurepip instead") + ); +} + #[test] fn test_reasoning_content_extracted() { // When an LLM End event carries output["reasoning"], the agent step From 6e13cfd63943a712c61b76513d671a67a7389974 Mon Sep 17 00:00:00 2001 From: Maryam Najafian Date: Thu, 23 Jul 2026 16:50:31 -0700 Subject: [PATCH 6/9] fix: align ATIF model fallback with OpenTelemetry and OpenInference (#547) #### Overview Keep ATIF model attribution aligned with OpenTelemetry and OpenInference when an LLM end event does not provide response-side model attribution. - [x] I confirm this contribution is my own work, or I have the right to submit it under this project's license. - [x] I searched existing issues and open pull requests, and this does not duplicate existing work. #### Details - Update the ATIF agent-step end path to preserve the paired start or request model when the end event has no response-side model attribution. - Keep existing response-first behavior unchanged when a normalized or manual response model is present. - Extend the exporter parity regression to assert that ATIF now matches the existing OpenTelemetry and OpenInference fallback behavior for this case. Validation: - `cargo fmt --all` - `cargo clippy --workspace --all-targets -- -D warnings` - `cargo test -p nemo-relay test_exporters_fall_back_to_requested_model_when_response_model_is_missing --lib` - `cargo test -p nemo-relay test_exporters_prefer_response_model_name_over_requested_model --lib` - `cargo test -p nemo-relay test_exporters_prefer_manual_response_model_name_over_requested_model --lib` - `just test-rust` - `just test-python` - `just test-go` - `just test-node` - `PATH="$HOME/.local/nemo-relay-tools/bin:$PATH" uv run pre-commit run --files crates/core/src/observability/atif.rs crates/core/tests/unit/observability/exporter_parity_tests.rs` #### Where should the reviewer start? Start in `crates/core/src/observability/atif.rs` in `handle_llm_end()`, where the emitted ATIF agent step now falls back to the paired start model only when the end event cannot provide response-side model attribution. The key regression coverage is in `crates/core/tests/unit/observability/exporter_parity_tests.rs`, especially `test_exporters_fall_back_to_requested_model_when_response_model_is_missing`. #### Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to) - Relates to: RELAY-564 ## Summary by CodeRabbit * **Bug Fixes** * Improved observability exports so agent steps retain the requested model name when response metadata is unavailable. * Ensured model names remain consistent across ATIF, OpenTelemetry, and OpenInference exports. * **Tests** * Added coverage verifying the requested model name is preserved in ATIF fallback scenarios. Authors: - Maryam Najafian (https://github.com/mnajafian-nv) Approvers: - Will Killian (https://github.com/willkill07) URL: https://github.com/NVIDIA/NeMo-Relay/pull/547 --- crates/core/src/observability/atif.rs | 4 +++- crates/core/tests/unit/observability/exporter_parity_tests.rs | 4 ++++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/crates/core/src/observability/atif.rs b/crates/core/src/observability/atif.rs index 5571dfd81..9d5402e52 100644 --- a/crates/core/src/observability/atif.rs +++ b/crates/core/src/observability/atif.rs @@ -2553,6 +2553,8 @@ impl StepConversionState { .llm_start_model_names .get(&event.uuid()) .map(String::as_str); + let step_model_name = + effective_response_model_name(event).or_else(|| paired_start_model.map(str::to_owned)); let ancestry = build_ancestry(event, &lookups.name_map); let invocation = build_invocation_info( start_ts, @@ -2582,7 +2584,7 @@ impl StepConversionState { .and_then(|response| atif_message_from_annotated_response(response)) .unwrap_or_else(|| extract_llm_response_message(output)), timestamp: Some(event.timestamp().to_rfc3339()), - model_name: effective_response_model_name(event), + model_name: step_model_name, reasoning_effort, reasoning_content, tool_calls, diff --git a/crates/core/tests/unit/observability/exporter_parity_tests.rs b/crates/core/tests/unit/observability/exporter_parity_tests.rs index 4f064356a..4982038e0 100644 --- a/crates/core/tests/unit/observability/exporter_parity_tests.rs +++ b/crates/core/tests/unit/observability/exporter_parity_tests.rs @@ -567,6 +567,10 @@ fn test_exporters_fall_back_to_requested_model_when_response_model_is_missing() }), ); + assert_eq!( + exports.agent_step().model_name.as_deref(), + Some("requested-model") + ); assert_eq!( exports .otel_attrs("model-call") From 9d16c074e90d289bc5eb8a039369bd389077103a Mon Sep 17 00:00:00 2001 From: Eric Evans II <194135482+ericevans-nv@users.noreply.github.com> Date: Fri, 24 Jul 2026 19:01:31 -0500 Subject: [PATCH 7/9] feat: add size-based operational log rotation (#555) #### Overview Add optional size-based rotation and retention to Relay operational file-log sinks. File sinks remain append-only unless rotation is explicitly configured. - [x] I confirm this contribution is my own work, or I have the right to submit it under this project's license. - [x] I searched existing issues and open pull requests, and this does not duplicate existing work. #### Details - Add validated `max_file_size_bytes` and `retained_files` settings for file sinks. - Require both rotation settings together and reject zero or unsafe retention values. - Rotate complete log records before the next write would exceed the configured size. - Retain the configured number of managed backups while preserving the active file. - Preserve historical backups outside the current retention window when the configured limit is reduced, avoiding automatic deletion of existing data. - Keep file writes behind the existing asynchronous sink and report rotation failures through its error path. - Detect conflicts between active sink paths and generated backup paths. - Preserve existing append-only behavior when rotation is omitted. - Document the configuration and add focused core and CLI coverage. #### Where should the reviewer start? Start with the rotating writer in `crates/core/src/logging/rotation.rs` and its integration with the asynchronous file sink in `crates/core/src/logging/sink.rs`. Focused behavior coverage is in `crates/core/tests/coverage/logging_tests.rs`. #### Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to) - Closes #553 ## Summary by CodeRabbit - **New Features** - Added size-based rotation for operational file log sinks using `max_file_size_bytes` and `retained_files` (with a `retained_files` max of 9). - When enabled, log writing rotates by size while keeping records intact and appending to the active file. - **Bug Fixes** - Enforced that rotation settings must be configured together; invalid values now fail configuration. - Prevented collisions between active log paths and any potential rotated/retained file targets. - **Documentation** - Updated the operational logging TOML examples and guidance for rotation configuration and limits. - **Tests** - Expanded parsing and integration coverage for rotation, retention ordering, boundary rotation, large-record behavior, and collision scenarios. Authors: - Eric Evans II (https://github.com/ericevans-nv) Approvers: - Will Killian (https://github.com/willkill07) - Maryam Najafian (https://github.com/mnajafian-nv) URL: https://github.com/NVIDIA/NeMo-Relay/pull/555 --- crates/cli/src/configuration/logging.rs | 21 +- .../cli/tests/coverage/shared/config_tests.rs | 50 ++++ crates/core/src/logging/config.rs | 72 ++++++ crates/core/src/logging/mod.rs | 6 +- crates/core/src/logging/rotation.rs | 148 ++++++++++++ crates/core/src/logging/sink.rs | 105 ++++++-- crates/core/tests/coverage/logging_tests.rs | 225 +++++++++++++++++- docs/reference/operational-logging.mdx | 9 +- 8 files changed, 602 insertions(+), 34 deletions(-) create mode 100644 crates/core/src/logging/rotation.rs diff --git a/crates/cli/src/configuration/logging.rs b/crates/cli/src/configuration/logging.rs index 85cf7b8f8..69b1909a7 100644 --- a/crates/cli/src/configuration/logging.rs +++ b/crates/cli/src/configuration/logging.rs @@ -7,8 +7,9 @@ use std::path::PathBuf; use nemo_relay::error::FlowError; use nemo_relay::logging::{ - DEFAULT_FILE_FLUSH_INTERVAL_MILLIS, DEFAULT_FILE_SINK_QUEUE_ENTRIES, FileLogSinkConfig, - LogFormat, LogLevel, LogSinkConfig, LoggingConfig, MAX_FILE_SINK_QUEUE_ENTRIES, + DEFAULT_FILE_FLUSH_INTERVAL_MILLIS, DEFAULT_FILE_SINK_QUEUE_ENTRIES, FileLogRotationConfig, + FileLogSinkConfig, LogFormat, LogLevel, LogSinkConfig, LoggingConfig, + MAX_FILE_SINK_QUEUE_ENTRIES, }; use serde::Deserialize; @@ -36,6 +37,8 @@ struct RawFileLogSinkConfig { /// Optional advanced: pending async queue entries per file sink (default /// [`DEFAULT_FILE_SINK_QUEUE_ENTRIES`]). queue_capacity: Option, + max_file_size_bytes: Option, + retained_files: Option, } pub(super) fn apply_file_logging_config( @@ -100,11 +103,25 @@ fn parse_file_log_sink( Some(capacity) => capacity, None => DEFAULT_FILE_SINK_QUEUE_ENTRIES, }; + let rotation = match (config.max_file_size_bytes, config.retained_files) { + (None, None) => None, + (Some(max_file_size_bytes), Some(retained_files)) => Some( + FileLogRotationConfig::new(max_file_size_bytes, retained_files) + .map_err(logging_parse_error)?, + ), + _ => { + return Err(CliError::Config( + "logging sink max_file_size_bytes and retained_files must be configured together" + .into(), + )); + } + }; Ok(LogSinkConfig::File(FileLogSinkConfig { path, level, format, queue_capacity, + rotation, })) } diff --git a/crates/cli/tests/coverage/shared/config_tests.rs b/crates/cli/tests/coverage/shared/config_tests.rs index f74db743a..103e3ff54 100644 --- a/crates/cli/tests/coverage/shared/config_tests.rs +++ b/crates/cli/tests/coverage/shared/config_tests.rs @@ -3271,6 +3271,56 @@ format = "human" } } +#[test] +fn logging_rotation_cli_config_preserves_pair_and_rejects_incomplete_pair() { + let temp = tempfile::tempdir().unwrap(); + let config_path = isolated_config_path(&temp); + let log_path = temp.path().join("relay.log.jsonl"); + std::fs::write( + &config_path, + format!( + r#" +[[logging.sinks]] +path = {} +max_file_size_bytes = 1024 +retained_files = 2 +"#, + toml_basic_string(log_path.to_string_lossy().as_ref()) + ), + ) + .unwrap(); + + let resolved = resolve_server_config(&GatewayOverrides { + config: Some(config_path), + ..GatewayOverrides::default() + }) + .unwrap(); + let LogSinkConfig::File(sink) = &resolved.logging.sinks[0]; + let rotation = sink.rotation.expect("complete rotation configuration"); + assert_eq!(rotation.max_file_size_bytes(), 1024); + assert_eq!(rotation.retained_files(), 2); + + let incomplete_path = isolated_config_path(&temp); + std::fs::write( + &incomplete_path, + r#" +[[logging.sinks]] +path = "relay.log.jsonl" +max_file_size_bytes = 1024 +"#, + ) + .unwrap(); + let error = resolve_server_config(&GatewayOverrides { + config: Some(incomplete_path), + ..GatewayOverrides::default() + }) + .unwrap_err() + .to_string(); + assert!(error.contains( + "logging sink max_file_size_bytes and retained_files must be configured together" + )); +} + #[test] fn logging_rejects_invalid_level_format_missing_path_and_zero_queue() { let temp = tempfile::tempdir().unwrap(); diff --git a/crates/core/src/logging/config.rs b/crates/core/src/logging/config.rs index 868881fb9..6dcd41fb9 100644 --- a/crates/core/src/logging/config.rs +++ b/crates/core/src/logging/config.rs @@ -28,6 +28,13 @@ pub const DEFAULT_FILE_FLUSH_INTERVAL_MILLIS: u64 = 1000; /// configuration above this bound is rejected with a config error. It cannot be raised. pub const MAX_FILE_SINK_QUEUE_ENTRIES: usize = 8_192; +/// Fixed hard maximum number of retained backup files per rotating file sink. +/// +/// Size-based rotation renames existing backup files on each rotation, so an unbounded value can +/// make one log write perform excessive filesystem work. This limit counts backup files and does +/// not include the active log file. +pub const MAX_FILE_SINK_RETAINED_FILES: usize = 9; + /// Operational logging configuration for [`LoggingRuntime::configure`](super::LoggingRuntime::configure). /// /// `level` is the process-wide **minimum severity**: call sites may emit any level, but records @@ -219,6 +226,8 @@ pub struct FileLogSinkConfig { /// Maximum pending asynchronous queue entries for this file sink. Must be greater than 0 and /// at most [`MAX_FILE_SINK_QUEUE_ENTRIES`]. pub queue_capacity: usize, + /// Optional size-based rotation and retention settings. + pub rotation: Option, } impl Default for FileLogSinkConfig { @@ -228,10 +237,56 @@ impl Default for FileLogSinkConfig { level: LogLevel::Info, format: LogFormat::Jsonl, queue_capacity: DEFAULT_FILE_SINK_QUEUE_ENTRIES, + rotation: None, } } } +/// Size-based rotation settings for a file log sink. +/// +/// `retained_files` counts previous log files and excludes the active file. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct FileLogRotationConfig { + max_file_size_bytes: u64, + retained_files: usize, +} + +impl FileLogRotationConfig { + /// Creates validated size-based rotation settings. + pub fn new(max_file_size_bytes: u64, retained_files: usize) -> Result { + if max_file_size_bytes == 0 { + return Err(FlowError::InvalidArgument( + "logging sink max_file_size_bytes must be greater than 0".into(), + )); + } + if retained_files == 0 { + return Err(FlowError::InvalidArgument( + "logging sink retained_files must be greater than 0".into(), + )); + } + if retained_files > MAX_FILE_SINK_RETAINED_FILES { + return Err(FlowError::InvalidArgument(format!( + "logging sink retained_files {retained_files} exceeds maximum \ + {MAX_FILE_SINK_RETAINED_FILES} backup files per sink" + ))); + } + Ok(Self { + max_file_size_bytes, + retained_files, + }) + } + + /// Maximum active file size before the next record triggers rotation. + pub fn max_file_size_bytes(self) -> u64 { + self.max_file_size_bytes + } + + /// Number of previous log files retained in addition to the active file. + pub fn retained_files(self) -> usize { + self.retained_files + } +} + #[derive(Debug, Deserialize)] struct LoggingDocument { logging: Option, @@ -277,6 +332,8 @@ struct RawFileLogSinkConfig { level: Option, format: Option, queue_capacity: Option, + max_file_size_bytes: Option, + retained_files: Option, } impl RawFileLogSinkConfig { @@ -318,11 +375,26 @@ impl RawFileLogSinkConfig { None => DEFAULT_FILE_SINK_QUEUE_ENTRIES, }; + let rotation = match (self.max_file_size_bytes, self.retained_files) { + (None, None) => None, + (Some(max_file_size_bytes), Some(retained_files)) => Some(FileLogRotationConfig::new( + max_file_size_bytes, + retained_files, + )?), + _ => { + return Err(FlowError::InvalidArgument( + "logging sink max_file_size_bytes and retained_files must be configured \ + together" + .into(), + )); + } + }; Ok(LogSinkConfig::File(FileLogSinkConfig { path, level, format, queue_capacity, + rotation, })) } } diff --git a/crates/core/src/logging/mod.rs b/crates/core/src/logging/mod.rs index 5c146888e..9b3c7528c 100644 --- a/crates/core/src/logging/mod.rs +++ b/crates/core/src/logging/mod.rs @@ -8,6 +8,7 @@ mod config; mod format; +mod rotation; mod sink; use std::io::{self, Write}; @@ -21,8 +22,9 @@ use uuid::Uuid; use crate::error::{FlowError, Result}; pub use config::{ - DEFAULT_FILE_FLUSH_INTERVAL_MILLIS, DEFAULT_FILE_SINK_QUEUE_ENTRIES, FileLogSinkConfig, - LogFormat, LogLevel, LogSinkConfig, LoggingConfig, MAX_FILE_SINK_QUEUE_ENTRIES, + DEFAULT_FILE_FLUSH_INTERVAL_MILLIS, DEFAULT_FILE_SINK_QUEUE_ENTRIES, FileLogRotationConfig, + FileLogSinkConfig, LogFormat, LogLevel, LogSinkConfig, LoggingConfig, + MAX_FILE_SINK_QUEUE_ENTRIES, MAX_FILE_SINK_RETAINED_FILES, }; pub(crate) use sink::build_logger; use sink::log_level_filter; diff --git a/crates/core/src/logging/rotation.rs b/crates/core/src/logging/rotation.rs new file mode 100644 index 000000000..3314377c1 --- /dev/null +++ b/crates/core/src/logging/rotation.rs @@ -0,0 +1,148 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Size-based file rotation for operational log sinks. + +use std::fs::{self, File, OpenOptions}; +use std::io::{self, BufWriter, Write}; +use std::path::{Path, PathBuf}; + +pub(crate) struct SizeRotatingFileWriter { + base_path: PathBuf, + file: Option>, + current_size: u64, + max_file_size_bytes: u64, + retained_files: usize, +} + +impl SizeRotatingFileWriter { + pub(crate) fn new( + base_path: PathBuf, + max_file_size_bytes: u64, + retained_files: usize, + ) -> io::Result { + create_parent_directory(&base_path)?; + let file = open_active_file(&base_path, false)?; + let current_size = file.get_ref().metadata()?.len(); + + Ok(Self { + base_path, + file: Some(file), + current_size, + max_file_size_bytes, + retained_files, + }) + } + + fn rotate_if_needed(&mut self, incoming_bytes: usize) -> io::Result<()> { + if self.current_size == 0 + || self.current_size.saturating_add(incoming_bytes as u64) <= self.max_file_size_bytes + { + return Ok(()); + } + + self.rotate() + } + + fn rotate(&mut self) -> io::Result<()> { + let mut file = self + .file + .take() + .ok_or_else(|| io::Error::other("rotating log file is not open"))?; + if let Err(error) = file.flush() { + self.file = Some(file); + return Err(error); + } + drop(file); + + if let Err(error) = rotate_files(&self.base_path, self.retained_files) { + return match self.reopen_after_failed_rotation() { + Ok(()) => Err(error), + Err(reopen_error) => Err(io::Error::new( + reopen_error.kind(), + format!( + "log rotation failed: {error}; failed to reopen active log file: \ + {reopen_error}" + ), + )), + }; + } + + self.file = Some(open_active_file(&self.base_path, true)?); + self.current_size = 0; + Ok(()) + } + + fn reopen_after_failed_rotation(&mut self) -> io::Result<()> { + let file = open_active_file(&self.base_path, false)?; + self.current_size = file.get_ref().metadata()?.len(); + self.file = Some(file); + Ok(()) + } +} + +impl Write for SizeRotatingFileWriter { + fn write(&mut self, buffer: &[u8]) -> io::Result { + self.rotate_if_needed(buffer.len())?; + self.file + .as_mut() + .ok_or_else(|| io::Error::other("rotating log file is not open"))? + .write_all(buffer)?; + self.current_size = self.current_size.saturating_add(buffer.len() as u64); + Ok(buffer.len()) + } + + fn flush(&mut self) -> io::Result<()> { + self.file + .as_mut() + .ok_or_else(|| io::Error::other("rotating log file is not open"))? + .flush() + } +} + +fn create_parent_directory(path: &Path) -> io::Result<()> { + if let Some(parent) = path.parent() + && !parent.as_os_str().is_empty() + { + fs::create_dir_all(parent)?; + } + Ok(()) +} + +fn open_active_file(path: &Path, truncate: bool) -> io::Result> { + let file = OpenOptions::new() + .create(true) + .write(true) + .append(!truncate) + .truncate(truncate) + .open(path)?; + Ok(BufWriter::new(file)) +} + +fn rotate_files(base_path: &Path, retained_files: usize) -> io::Result<()> { + for index in (1..=retained_files).rev() { + let source = if index == 1 { + base_path.to_path_buf() + } else { + rotated_log_path(base_path, index - 1) + }; + if !source.exists() { + continue; + } + + let destination = rotated_log_path(base_path, index); + fs::rename(source, destination)?; + } + Ok(()) +} + +pub(crate) fn rotated_log_path(base_path: &Path, index: usize) -> PathBuf { + let stem = base_path.file_stem().unwrap_or(base_path.as_os_str()); + let mut file_name = stem.to_os_string(); + file_name.push(format!(".{index}")); + if let Some(extension) = base_path.extension() { + file_name.push("."); + file_name.push(extension); + } + base_path.with_file_name(file_name) +} diff --git a/crates/core/src/logging/sink.rs b/crates/core/src/logging/sink.rs index bb0456804..51ac42d10 100644 --- a/crates/core/src/logging/sink.rs +++ b/crates/core/src/logging/sink.rs @@ -10,12 +10,13 @@ use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use spdlog::sink::{AsyncPoolSink, FileSink, OverflowPolicy, StdStreamSink}; +use spdlog::sink::{AsyncPoolSink, FileSink, OverflowPolicy, StdStreamSink, WriteSink}; use spdlog::terminal_style::StyleMode; use spdlog::{Level, LevelFilter, Logger, ThreadPool}; use super::config::{LogLevel, LogSinkConfig, LoggingConfig, MAX_FILE_SINK_QUEUE_ENTRIES}; use super::format::RelayFormatter; +use super::rotation::{SizeRotatingFileWriter, rotated_log_path}; use crate::error::{FlowError, Result}; pub(crate) fn build_logger( @@ -24,7 +25,8 @@ pub(crate) fn build_logger( ) -> Result<(Arc, Vec>)> { let mut sinks: Vec> = Vec::new(); let mut thread_pools = Vec::new(); - let mut resolved_paths: Vec = Vec::new(); + let mut active_paths: Vec = Vec::new(); + let mut reserved_paths: Vec = Vec::new(); let stderr_sink = StdStreamSink::builder() .stderr() @@ -44,36 +46,70 @@ pub(crate) fn build_logger( for sink in &config.sinks { let LogSinkConfig::File(file_sink) = sink; let resolved_path = resolve_log_path(&file_sink.path)?; - if resolved_paths - .iter() - .any(|existing| existing == &resolved_path) - { + if active_paths.contains(&resolved_path) { return Err(FlowError::InvalidArgument(format!( "duplicate logging sink path {}", resolved_path.display() ))); } - resolved_paths.push(resolved_path.clone()); + let candidate_paths = reserved_sink_paths(&resolved_path, file_sink.rotation); + if let Some(collision) = candidate_paths + .iter() + .find(|candidate| reserved_paths.contains(candidate)) + { + return Err(FlowError::InvalidArgument(format!( + "logging sink path conflicts with another active or rotated file: {}", + collision.display() + ))); + } + active_paths.push(resolved_path.clone()); + reserved_paths.extend(candidate_paths); - // FileSink performs the real open/append. AsyncPoolSink is spdlog's stock bounded queue + - // worker pool in front of that file so hot paths enqueue instead of blocking on disk I/O. - // Overflow drops incoming records so a stuck disk cannot stall the process. - let file = FileSink::builder() - .path(&resolved_path) - .truncate(false) - .formatter(RelayFormatter { - format: file_sink.format, - root_relay_id: root_relay_id.clone(), - }) - .level_filter(spdlog_level_filter(file_sink.level)) - .error_handler(stderr_error_handler(&resolved_path.display().to_string())) - .build_arc() - .map_err(|error| { - FlowError::InvalidArgument(format!( - "failed to open logging sink {}: {error}", - resolved_path.display() - )) - })?; + let file: Arc = match file_sink.rotation { + None => FileSink::builder() + .path(&resolved_path) + .truncate(false) + .formatter(RelayFormatter { + format: file_sink.format, + root_relay_id: root_relay_id.clone(), + }) + .level_filter(spdlog_level_filter(file_sink.level)) + .error_handler(stderr_error_handler(&resolved_path.display().to_string())) + .build_arc() + .map_err(|error| { + FlowError::InvalidArgument(format!( + "failed to open logging sink {}: {error}", + resolved_path.display() + )) + })?, + Some(rotation) => WriteSink::builder() + .target( + SizeRotatingFileWriter::new( + resolved_path.clone(), + rotation.max_file_size_bytes(), + rotation.retained_files(), + ) + .map_err(|error| { + FlowError::InvalidArgument(format!( + "failed to open rotating logging sink {}: {error}", + resolved_path.display() + )) + })?, + ) + .formatter(RelayFormatter { + format: file_sink.format, + root_relay_id: root_relay_id.clone(), + }) + .level_filter(spdlog_level_filter(file_sink.level)) + .error_handler(stderr_error_handler(&resolved_path.display().to_string())) + .build_arc() + .map_err(|error| { + FlowError::InvalidArgument(format!( + "failed to open rotating logging sink {}: {error}", + resolved_path.display() + )) + })?, + }; if file_sink.queue_capacity > MAX_FILE_SINK_QUEUE_ENTRIES { return Err(FlowError::InvalidArgument(format!( @@ -132,6 +168,23 @@ pub(crate) fn build_logger( Ok((logger, thread_pools)) } +fn reserved_sink_paths( + resolved_path: &Path, + rotation: Option, +) -> Vec { + let mut paths = vec![resolved_path.to_path_buf()]; + if let Some(rotation) = rotation { + paths.reserve(rotation.retained_files()); + for index in 1..=rotation.retained_files() { + paths.push(logging_path_identity(&rotated_log_path( + resolved_path, + index, + ))); + } + } + paths +} + fn resolve_log_path(path: &Path) -> Result { // Relative paths resolve against process CWD. Absolute paths are unchanged. No `~` or env // expansion. diff --git a/crates/core/tests/coverage/logging_tests.rs b/crates/core/tests/coverage/logging_tests.rs index 2fc6fe558..bf2df1fed 100644 --- a/crates/core/tests/coverage/logging_tests.rs +++ b/crates/core/tests/coverage/logging_tests.rs @@ -2,13 +2,14 @@ // SPDX-License-Identifier: Apache-2.0 use crate::logging::{ - FileLogSinkConfig, LogFormat, LogLevel, LogSinkConfig, LoggingConfig, LoggingRuntime, - MAX_FILE_SINK_QUEUE_ENTRIES, build_logger, format_event_for_test, init_logging, + FileLogRotationConfig, FileLogSinkConfig, LogFormat, LogLevel, LogSinkConfig, LoggingConfig, + LoggingRuntime, MAX_FILE_SINK_QUEUE_ENTRIES, MAX_FILE_SINK_RETAINED_FILES, build_logger, + format_event_for_test, init_logging, }; use serde_json::Value; use spdlog::Level; use std::ffi::{OsStr, OsString}; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use std::sync::{Arc, Barrier, Mutex, MutexGuard}; static LOGGING_TEST_LOCK: Mutex<()> = Mutex::new(()); @@ -239,6 +240,7 @@ queue_capacity = 7 level: LogLevel::Warn, format: LogFormat::Human, queue_capacity: 7, + rotation: None, })] ); } @@ -295,6 +297,97 @@ queue_capacity = 7 ); } +#[test] +fn logging_rotation_configuration_parses_complete_pair_and_rejects_invalid_values() { + let config = LoggingConfig::from_toml_document( + r#" +[logging] + +[[logging.sinks]] +path = "relay.log.jsonl" +max_file_size_bytes = 1024 +retained_files = 2 +"#, + ) + .unwrap(); + + let LogSinkConfig::File(sink) = &config.sinks[0]; + assert_eq!( + sink.rotation, + Some(FileLogRotationConfig::new(1024, 2).unwrap()) + ); + + let invalid_documents = [ + ( + "missing retained_files", + r#" +[logging] +[[logging.sinks]] +path = "relay.log.jsonl" +max_file_size_bytes = 1024 +"# + .to_owned(), + "must be configured together", + ), + ( + "missing max_file_size_bytes", + r#" +[logging] +[[logging.sinks]] +path = "relay.log.jsonl" +retained_files = 2 +"# + .to_owned(), + "must be configured together", + ), + ( + "zero max_file_size_bytes", + r#" +[logging] +[[logging.sinks]] +path = "relay.log.jsonl" +max_file_size_bytes = 0 +retained_files = 2 +"# + .to_owned(), + "max_file_size_bytes must be greater than 0", + ), + ( + "zero retained_files", + r#" +[logging] +[[logging.sinks]] +path = "relay.log.jsonl" +max_file_size_bytes = 1024 +retained_files = 0 +"# + .to_owned(), + "retained_files must be greater than 0", + ), + ( + "retained_files over maximum", + format!( + r#" +[logging] +[[logging.sinks]] +path = "relay.log.jsonl" +max_file_size_bytes = 1024 +retained_files = {} +"#, + MAX_FILE_SINK_RETAINED_FILES + 1 + ), + "exceeds maximum", + ), + ]; + + for (name, document, expected) in invalid_documents { + let error = LoggingConfig::from_toml_document(&document) + .unwrap_err() + .to_string(); + assert!(error.contains(expected), "{name}: {error}"); + } +} + #[test] #[cfg(unix)] fn logging_environment_rejects_non_unicode_values() { @@ -452,6 +545,132 @@ fn wait_for_log_line(path: &std::path::Path, ready: impl Fn(&str) -> bool) -> St std::fs::read_to_string(path).unwrap_or_default() } +fn read_single_jsonl_record(path: &Path) -> Value { + let contents = std::fs::read_to_string(path).unwrap(); + let mut lines = contents.lines(); + let record = serde_json::from_str(lines.next().expect("one JSONL record")).unwrap(); + assert!( + lines.next().is_none(), + "expected exactly one JSONL record in {}, got {contents:?}", + path.display() + ); + record +} + +#[test] +fn logging_rotation_retains_newest_backups_and_complete_records() { + let _lock = lock_logging_tests(); + let temp = tempfile::tempdir().unwrap(); + let path = temp.path().join("relay.log.jsonl"); + let config = LoggingConfig { + sinks: vec![LogSinkConfig::File(FileLogSinkConfig { + path: path.clone(), + rotation: Some(FileLogRotationConfig::new(1, 2).unwrap()), + ..FileLogSinkConfig::default() + })], + ..default_config() + }; + + let runtime = init_logging(&config).unwrap(); + log::info!(target: "nemo_relay.rotation_test", event = "rotation_one"; "rotation one"); + log::info!(target: "nemo_relay.rotation_test", event = "rotation_two"; "rotation two"); + log::info!(target: "nemo_relay.rotation_test", event = "rotation_three"; "rotation three"); + runtime.shutdown(); + + let active = read_single_jsonl_record(&path); + let newest_backup = read_single_jsonl_record(&temp.path().join("relay.log.1.jsonl")); + let oldest_backup = read_single_jsonl_record(&temp.path().join("relay.log.2.jsonl")); + + assert_eq!(active["event"], "logging_shutdown_started"); + assert_eq!(newest_backup["event"], "rotation_three"); + assert_eq!(oldest_backup["event"], "rotation_two"); + assert!(!temp.path().join("relay.log.3.jsonl").exists()); +} + +#[test] +fn logging_rotation_rotates_existing_file_at_boundary() { + let _lock = lock_logging_tests(); + let temp = tempfile::tempdir().unwrap(); + let path = temp.path().join("relay.log.jsonl"); + let existing_record = "x".repeat(64); + std::fs::write(&path, &existing_record).unwrap(); + let config = LoggingConfig { + sinks: vec![LogSinkConfig::File(FileLogSinkConfig { + path: path.clone(), + rotation: Some(FileLogRotationConfig::new(64, 2).unwrap()), + ..FileLogSinkConfig::default() + })], + ..default_config() + }; + + let runtime = init_logging(&config).unwrap(); + runtime.shutdown(); + + assert_eq!( + std::fs::read_to_string(temp.path().join("relay.log.2.jsonl")).unwrap(), + existing_record + ); + assert_eq!( + read_single_jsonl_record(&path)["event"], + "logging_shutdown_started" + ); +} + +#[test] +fn logging_rotation_preserves_historical_backups_outside_retention_window() { + let _lock = lock_logging_tests(); + let temp = tempfile::tempdir().unwrap(); + let path = temp.path().join("relay.log.jsonl"); + let historical_path = temp.path().join("relay.log.3.jsonl"); + let historical_contents = "historical backup outside current retention window\n"; + std::fs::write(&historical_path, historical_contents).unwrap(); + let config = LoggingConfig { + sinks: vec![LogSinkConfig::File(FileLogSinkConfig { + path, + rotation: Some(FileLogRotationConfig::new(1, 2).unwrap()), + ..FileLogSinkConfig::default() + })], + ..default_config() + }; + + let runtime = init_logging(&config).unwrap(); + runtime.shutdown(); + + assert_eq!( + std::fs::read_to_string(historical_path).unwrap(), + historical_contents + ); +} + +#[test] +fn logging_rotation_rejects_generated_backup_path_collision() { + let _lock = lock_logging_tests(); + let temp = tempfile::tempdir().unwrap(); + let path = temp.path().join("relay.log.jsonl"); + let config = LoggingConfig { + sinks: vec![ + LogSinkConfig::File(FileLogSinkConfig { + path: path.clone(), + rotation: Some(FileLogRotationConfig::new(1024, 1).unwrap()), + ..FileLogSinkConfig::default() + }), + LogSinkConfig::File(FileLogSinkConfig { + path: temp.path().join("relay.log.1.jsonl"), + ..FileLogSinkConfig::default() + }), + ], + ..default_config() + }; + + let error = build_logger(&config, "root".into()) + .err() + .expect("generated backup path collision should fail") + .to_string(); + + assert!(error.contains("conflicts with another active or rotated file")); + assert!(error.contains("relay.log.1.jsonl")); +} + #[test] fn sink_level_filter_drops_events_below_sink_minimum() { let _lock = lock_logging_tests(); diff --git a/docs/reference/operational-logging.mdx b/docs/reference/operational-logging.mdx index 2202fccde..29f2d092e 100644 --- a/docs/reference/operational-logging.mdx +++ b/docs/reference/operational-logging.mdx @@ -100,11 +100,18 @@ path = ".nemo-relay/logs/relay.log.jsonl" format = "jsonl" level = "debug" queue_capacity = 1024 +max_file_size_bytes = 10485760 +retained_files = 5 ``` File sink paths are resolved relative to the process working directory. File sinks use asynchronous queues, and `queue_capacity` cannot exceed 8,192 entries -per sink. +per sink. Size-based rotation is optional; when enabled, +`max_file_size_bytes` and `retained_files` must be configured together. +`retained_files` counts backup files in addition to the active file and cannot +exceed 9. A record larger than `max_file_size_bytes` remains intact rather +than being split. File sinks remain append-only when rotation settings are +omitted. ## Rust Library API From 42c7655ff8cdd4300dc80c157e58d09a6096bcb5 Mon Sep 17 00:00:00 2001 From: Will Killian <2007799+willkill07@users.noreply.github.com> Date: Sun, 26 Jul 2026 18:57:24 -0400 Subject: [PATCH 8/9] feat(sanitizers)!: expose per-call codec capabilities (#546) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit #### Overview Exposes the active per-call LLM codec to observability sanitizers across in-process and plugin boundaries, then uses that resolved capability to make PII redaction provider-correct and fail-closed for mixed managed LLM traffic. > [!WARNING] > **BREAKING CHANGE:** LLM request and response sanitizers now use the required two-argument contract `(payload, directional_context) -> optional payload`. No registration-time compatibility shim remains. Rust, Python, Node.js, native plugins, raw C FFI consumers, worker plugins, and downstream bindings must update their callbacks. All Rust worker sanitizer callbacks—mark, scope, tool, and LLM—are now async. Native dynamic plugins remain ABI v2 and workers remain `grpc-v1`, but the v2 context/host tables and grpc-v1 invocation/service definitions change in place. - [x] I confirm this contribution is my own work, or I have the right to submit it under this project's license. - [x] I searched existing issues and open pull requests, and this does not duplicate existing work. #### Details - Replaces split legacy/contextual LLM sanitizer registration with one canonical request and response callback representation while preserving priority, array-order tie breaking, scope ownership, registration rollback, snapshots, and short-circuit omission. - Introduces distinct `LlmSanitizeRequestContext` and `LlmSanitizeResponseContext` values with structured codec identities: none, built-in, runtime, and opaque. Every active codec, including opaque codecs, can be resolved for the lifetime of the callback. - Gives request sanitizers `decode(request)` and `encode(annotated, original)`; gives response sanitizers `decode(response)`. - Extends native ABI v2 and the raw C FFI with borrowed directional codec handles and host-owned transformation operations. Safe wrappers reject retained handles after the callback returns, contain codec panics, preserve ownership for aliased results, and fail closed for malformed output or unrepresentable runtime codec IDs. - Extends `grpc-v1` with directional contexts and authenticated, invocation-scoped codec capabilities. Host RPCs reject forged, expired, unauthorized, and wrong-direction capability IDs, and SDK authors see async codec proxies rather than raw IDs. - Makes Rust worker sanitizer registration async-only while Python workers accept immediate values or awaitables under the same required two-argument callback contract. - Keeps sanitizer failures fail-closed for observability: failures or `None`/`null` results omit the payload and annotation without changing the client-visible request or response. - Invalidates caller-supplied normalized annotations whenever sanitization changes a buffered payload, rebuilding only through the active directional codec. Decode failure or no codec omits the stale annotation. - Updates PII redaction to use the resolved active codec per call, use configured `codec` only when no active codec exists, support provider-agnostic policies without a configured codec, and omit codec-dependent observability data when normalization is unavailable. - Hardens normalized target overlays: object-field removals are applied incrementally, unsafe array removals and inconsistent projections fail closed, and OpenAI Responses keeps the top-level `output_text` alias synchronized with sanitized structured output. - Makes Node custom-codec reference cleanup event-loop-affine so early-dropped streams cannot delete N-API references from a Tokio thread. - Documents parameter order, directional contexts, codec identity and resolution, capability lifetime, omission behavior, native/worker contracts, and mixed-provider PII usage with Fern tabs. Adds the 0.7 breaking-change release note and a complete 0.6-to-0.7 migration procedure. Validation: - Repository-wide pre-commit suite passed before the final review fixes; the final changed-file pre-commit suite also passed, including FFI header sync, Cargo fmt/Clippy/check, Node formatting, and documentation links. - Node: 290 tests passed. - Python: 556 tests passed. - Go: all package tests passed. - Raw FFI: 82 unit tests and 77 integration tests passed. - Targeted buffered and streaming OpenAI Responses PII regressions passed and verified serialized events contain no original secret. - The public Rust native-plugin example passes its standalone `cargo check`. - The affected Rust core, PII, native-plugin, worker protocol/SDK, and binding capability suites passed. One unrelated CLI test in the workspace was contaminated by an ancestor user plugin configuration; the exact test passed from an isolated checkout of this commit. - Documentation build completed with 0 errors. #### Where should the reviewer start? Start with `crates/core/src/api/runtime/callbacks.rs` and `crates/core/src/api/llm.rs` for the canonical callback contract, codec propagation, and annotation invalidation. Then review `crates/plugin/src/lib.rs`, `crates/core/src/plugin/dynamic/native.rs`, `crates/core/src/plugin/dynamic/worker.rs`, and `crates/worker-proto/proto/nemo/relay/worker/v1/plugin_worker.proto` for cross-boundary capability lifetime and authorization. Finish with `crates/pii-redaction/src/builtin.rs` and `crates/pii-redaction/src/overlay.rs` for per-call codec selection and fail-closed redaction, then review `docs/reference/migration-guides.mdx` for the public upgrade contract. #### Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to) - Closes #526 - Closes RELAY-555 ## Summary by CodeRabbit - **New Features** - Codec-aware LLM sanitizers: callbacks now receive per-call request/response codec context for directional normalization/redaction. - Sanitizers can omit observability payloads/annotations by returning no value (`null`/`None`). - Codec-aware support expanded across Node.js, Python, Go, Rust, native plugins, and worker integrations. - PII redaction now selects the active codec per call for mixed-provider traffic. - **Breaking Changes** - Callback signatures changed to `(payload, context)` and must return an optional payload (`null`/`None` to omit). - Native plugins now require Native ABI v2 (rebuilt with the matching SDK/protocol). - **Documentation** - Added migration guidance and updated middleware/plugin docs for codec-aware sanitizer behavior. Authors: - Will Killian (https://github.com/willkill07) Approvers: - Alex Fournier (https://github.com/afourniernv) URL: https://github.com/NVIDIA/NeMo-Relay/pull/546 --- crates/cli/src/gateway/mod.rs | 6 +- crates/core/src/api/llm.rs | 179 ++- crates/core/src/api/registry.rs | 5 +- crates/core/src/api/runtime.rs | 10 +- crates/core/src/api/runtime/callbacks.rs | 180 ++- crates/core/src/api/runtime/state.rs | 17 +- crates/core/src/codec/anthropic.rs | 9 + crates/core/src/codec/openai_chat.rs | 9 + crates/core/src/codec/openai_responses.rs | 9 + crates/core/src/codec/traits.rs | 18 + crates/core/src/plugin/dynamic/native.rs | 431 ++++++- crates/core/src/plugin/dynamic/worker.rs | 851 +++++++++----- crates/core/src/stream.rs | 46 +- .../tests/fixtures/native_plugin/src/lib.rs | 4 +- .../tests/fixtures/worker_plugin/src/main.rs | 102 +- .../tests/integration/api_surface_tests.rs | 22 +- .../tests/integration/middleware_tests.rs | 16 +- .../tests/integration/native_plugin_tests.rs | 17 +- .../core/tests/integration/pipeline_tests.rs | 12 +- crates/core/tests/unit/context_tests.rs | 8 +- .../core/tests/unit/dynamic_worker_tests.rs | 465 +++++++- crates/core/tests/unit/llm_api_tests.rs | 716 +++++++++++- crates/core/tests/unit/native_plugin_tests.rs | 517 ++++++++- crates/core/tests/unit/plugin_tests.rs | 16 +- crates/ffi/nemo_relay.h | 142 ++- crates/ffi/src/api/llm.rs | 140 ++- crates/ffi/src/api/llm_registry.rs | 22 +- crates/ffi/src/api/mod.rs | 13 +- crates/ffi/src/api/plugin.rs | 32 +- crates/ffi/src/api/scope_registry.rs | 24 +- crates/ffi/src/callable.rs | 205 +++- crates/ffi/src/types/mod.rs | 4 + crates/ffi/tests/integration/api_tests.rs | 19 +- .../tests/integration/callable_extra_tests.rs | 88 +- crates/ffi/tests/unit/api_tests.rs | 134 ++- crates/ffi/tests/unit/callable_tests.rs | 129 ++- crates/node/package.json | 4 +- crates/node/plugin.d.ts | 38 +- crates/node/src/api/mod.rs | 477 ++++++-- crates/node/src/callable.rs | 330 +++++- crates/node/src/stream.rs | 4 + crates/node/tests/callback_error_tests.mjs | 12 +- crates/node/tests/event_sanitizers_tests.mjs | 25 +- crates/node/tests/llm_tests.mjs | 291 ++++- crates/pii-redaction/README.md | 2 - crates/pii-redaction/src/builtin.rs | 357 +++++- crates/pii-redaction/src/component.rs | 11 +- crates/pii-redaction/src/overlay.rs | 6 +- .../tests/unit/component_tests.rs | 1013 +++++++++++++++-- crates/plugin/src/lib.rs | 344 +++++- crates/plugin/tests/typed_callbacks.rs | 387 ++++++- crates/python/src/py_api/mod.rs | 46 +- crates/python/src/py_callable.rs | 134 ++- crates/python/src/py_plugin.rs | 14 +- crates/python/src/py_types/codecs.rs | 47 + crates/python/src/py_types/core.rs | 87 +- crates/python/src/py_types/mod.rs | 31 +- .../python/tests/coverage/coverage_tests.rs | 55 +- .../tests/coverage/py_api_coverage_tests.rs | 4 +- .../coverage/py_callable_coverage_tests.rs | 13 +- .../coverage/py_plugin_coverage_tests.rs | 42 +- .../nemo/relay/worker/v1/plugin_worker.proto | 58 + crates/worker/src/lib.rs | 690 ++++++++--- crates/worker/tests/worker_sdk_tests.rs | 329 +++++- docs/about-nemo-relay/concepts/middleware.mdx | 118 ++ docs/about-nemo-relay/release-notes/index.mdx | 35 +- docs/build-plugins/dynamic-plugins/about.mdx | 2 +- .../grpc-worker/grpc-worker-protocol.mdx | 62 +- .../dynamic-plugins/native-dynamic/about.mdx | 33 +- docs/build-plugins/language-binding/about.mdx | 8 + .../configure-plugins/pii-redaction/about.mdx | 8 +- .../pii-redaction/configuration.mdx | 57 +- .../llm-request-intercept-outcomes.mdx | 9 +- docs/reference/migration-guides.mdx | 289 ++++- examples/rust-native-plugin/src/lib.rs | 12 +- go/nemo_relay/adaptive_plugin_test.go | 14 +- go/nemo_relay/callbacks.go | 316 ++++- go/nemo_relay/callbacks_test.go | 71 ++ go/nemo_relay/deregister_test.go | 4 +- go/nemo_relay/error_test.go | 4 +- go/nemo_relay/guardrails/guardrails_test.go | 19 +- go/nemo_relay/llm_test.go | 233 +++- go/nemo_relay/nemo_relay.go | 40 +- go/nemo_relay/plugin.go | 21 +- go/nemo_relay/scope_local_test.go | 15 +- python/nemo_relay/__init__.py | 23 +- python/nemo_relay/__init__.pyi | 35 +- python/nemo_relay/_native.pyi | 33 +- python/nemo_relay/guardrails.py | 22 +- .../plugin/src/nemo_relay_plugin/__init__.py | 15 + python/plugin/src/nemo_relay_plugin/_api.py | 301 +++-- python/tests/plugin/test_worker_sdk.py | 299 ++++- python/tests/test_builtin_codecs.py | 3 +- python/tests/test_llm.py | 181 ++- python/tests/test_scope_local.py | 31 +- 95 files changed, 10286 insertions(+), 1495 deletions(-) diff --git a/crates/cli/src/gateway/mod.rs b/crates/cli/src/gateway/mod.rs index a42444553..b4ebdbe9e 100644 --- a/crates/cli/src/gateway/mod.rs +++ b/crates/cli/src/gateway/mod.rs @@ -26,7 +26,7 @@ use nemo_relay::api::llm::{ llm_stream_call_execute, }; use nemo_relay::api::runtime::{ - LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn, TASK_SCOPE_STACK, + LlmCodecIdentity, LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn, TASK_SCOPE_STACK, }; use nemo_relay::codec::request::AnnotatedLlmRequest; use nemo_relay::codec::resolve::{ @@ -160,6 +160,10 @@ struct GatewayRequestCodec { } impl LlmCodec for GatewayRequestCodec { + fn codec_identity(&self) -> LlmCodecIdentity { + self.inner.codec_identity() + } + fn decode(&self, request: &LlmRequest) -> nemo_relay::error::Result { self.inner.decode(request) } diff --git a/crates/core/src/api/llm.rs b/crates/core/src/api/llm.rs index 8484d18e3..29e4655a9 100644 --- a/crates/core/src/api/llm.rs +++ b/crates/core/src/api/llm.rs @@ -15,11 +15,13 @@ use crate::api::event::{ use crate::api::optimization::{ LlmOptimizationRecorder, finalize_optimization_summary, scope_llm_optimization_recorder, }; +#[cfg(test)] +use crate::api::runtime::LlmCodecIdentity; use crate::api::runtime::NemoRelayContextState; use crate::api::runtime::global_context; use crate::api::runtime::{ EventSubscriberFn, LlmCollectorFn, LlmExecutionNextFn, LlmFinalizerFn, LlmJsonStream, - LlmStreamExecutionNextFn, + LlmSanitizeRequestContext, LlmSanitizeResponseContext, LlmStreamExecutionNextFn, }; use crate::api::runtime::{ScopeStackHandle, current_scope_stack}; use crate::api::scope::event; @@ -391,7 +393,7 @@ fn emit_llm_start( handle: &LlmHandle, request: &LlmRequest, annotated_request: Option>, - request_codec: Option<&dyn LlmCodec>, + request_codec: Option>, ) -> Result<()> { ensure_runtime_owner()?; let subscribers = { @@ -412,7 +414,7 @@ fn emit_llm_start_with_subscribers( handle: &LlmHandle, request: &LlmRequest, annotated_request: Option>, - request_codec: Option<&dyn LlmCodec>, + request_codec: Option>, subscribers: &[EventSubscriberFn], ) -> Result<()> { ensure_runtime_owner()?; @@ -428,36 +430,43 @@ fn emit_llm_start_with_subscribers( .map_err(|error| FlowError::Internal(error.to_string()))?; state.llm_sanitize_request_entries(&scope_locals) }; - let mut sanitized_request = - NemoRelayContextState::llm_sanitize_request_snapshot_chain(request.clone(), &entries); - let mut annotated_request = match request_codec { - Some(codec) - if sanitized_request.headers != request.headers - || sanitized_request.content != request.content => - { - codec.decode(&sanitized_request).ok().map(Arc::new) + let mut sanitized_request = NemoRelayContextState::llm_sanitize_request_snapshot_chain( + 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); + 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) } - _ => annotated_request, + (Some(_), _) if !request_changed => annotated_request, + (None, _) => None, + (Some(_), _) => None, }; let scope_stack = handle.captured_scope_stack(); let agent_is_fresh = { let mut scope_guard = scope_stack.write().expect("scope stack lock poisoned"); scope_guard.take_agent_freshness(handle.parent_uuid) }; - if !agent_is_fresh { + if !agent_is_fresh && let Some(sanitized_request) = sanitized_request.as_mut() { project_llm_request_to_current_user_turn( - &mut sanitized_request, + sanitized_request, &mut annotated_request, - request_codec, + request_codec.as_deref(), ); } - let input = serde_json::to_value(&sanitized_request).unwrap_or(Json::Null); + let input = sanitized_request + .as_ref() + .and_then(|sanitized_request| serde_json::to_value(sanitized_request).ok()); let event = { let context = global_context(); let state = context .read() .map_err(|error| FlowError::Internal(error.to_string()))?; - state.build_llm_start_event(handle, Some(input), annotated_request) + state.build_llm_start_event(handle, input, annotated_request) }; if let Some(event) = sanitize_event_with_scope_stack(event, scope_stack) { NemoRelayContextState::emit_event(&event, subscribers); @@ -626,9 +635,8 @@ struct LlmCallEndBehavior { /// # Parameters /// - `handle`: LLM handle to close. /// - `response`: Raw provider response associated with the end event. -/// - `data`: Optional application payload retained for compatibility. The -/// emitted end event data is the sanitized `response` unless it sanitizes to -/// JSON null, in which case this payload is used. +/// - `data`: Optional application payload retained for compatibility. When the +/// raw `response` is JSON null, this payload is sanitized in its place. /// - `metadata`: Optional JSON metadata recorded on the end event. /// - `annotated_response`: Optional normalized response annotation produced by /// a response codec. When omitted and `response_codec` is supplied, the @@ -693,20 +701,36 @@ fn llm_call_end_with_behavior( let entries = state.llm_sanitize_response_entries(&scope_locals); (entries, subscribers) }; - let sanitized_response = - NemoRelayContextState::llm_sanitize_response_snapshot_chain(response, &entries); - let data = if sanitized_response.is_null() { - data + let response_was_null_without_fallback = response.is_null() && data.is_none(); + let response = if response.is_null() { + data.unwrap_or(response) } else { - Some(sanitized_response) + response }; - let (mut annotated_response, decode_error) = resolve_llm_end_annotation( - annotated_response, - response_codec, - data.as_ref(), - &behavior, - &handle.name, + let sanitized_response = NemoRelayContextState::llm_sanitize_response_snapshot_chain( + response.clone(), + LlmSanitizeResponseContext::for_response_codec(response_codec.clone()), + &entries, ); + let response_changed = sanitized_response + .as_ref() + .is_some_and(|sanitized_response| sanitized_response != &response); + let data = match sanitized_response { + Some(response) if response_was_null_without_fallback && response.is_null() => None, + response => response, + }; + let annotation_omitted = data.as_ref().is_none_or(Json::is_null); + let (mut annotated_response, decode_error) = if annotation_omitted { + (None, None) + } else { + resolve_llm_end_annotation( + (!response_changed).then_some(annotated_response).flatten(), + response_codec, + data.as_ref(), + &behavior, + &handle.name, + ) + }; handle.optimization_recorder.close_for_finalization(None); emit_optimization_marks(handle, &subscribers); let pricing = crate::codec::response::active_pricing_resolver(); @@ -716,7 +740,8 @@ fn llm_call_end_with_behavior( handle.model_name.as_deref(), &pricing, ); - if annotated_response.is_none() + if !annotation_omitted + && annotated_response.is_none() && let Some(summary) = summary { annotated_response = Some(AnnotatedLlmResponse { @@ -753,6 +778,22 @@ fn llm_call_end_with_behavior( } } +#[cfg(test)] +fn sanitize_context_for_request_codec(codec: Option<&dyn LlmCodec>) -> LlmSanitizeRequestContext { + LlmSanitizeRequestContext::with_identity( + codec.map_or(LlmCodecIdentity::None, LlmCodec::codec_identity), + ) +} + +#[cfg(test)] +pub(crate) fn sanitize_context_for_response_codec( + codec: Option<&dyn LlmResponseCodec>, +) -> LlmSanitizeResponseContext { + LlmSanitizeResponseContext::with_identity( + codec.map_or(LlmCodecIdentity::None, LlmResponseCodec::codec_identity), + ) +} + fn resolve_llm_end_annotation( annotated_response: Option>, response_codec: Option>, @@ -780,39 +821,63 @@ fn resolve_llm_end_annotation( fn emit_llm_end_without_output( handle: &LlmHandle, metadata: Option, + response_codec: Option>, lifecycle_subscribers: Option<&[EventSubscriberFn]>, ) -> Result<()> { ensure_runtime_owner()?; - let subscribers = { + let (entries, subscribers) = { let scope_stack = handle.captured_scope_stack(); let scope_guard = scope_stack.read().expect("scope stack lock poisoned"); + let scope_locals = scope_guard.collect_scope_local_registries(|registries| { + ®istries.llm_sanitize_response_guardrails + }); let scope_subscribers = scope_guard.collect_scope_local_subscribers(); - match lifecycle_subscribers { + let subscribers = match lifecycle_subscribers { Some(subscribers) => subscribers.to_vec(), None => snapshot_event_subscribers(scope_subscribers)?, - } + }; + let context = global_context(); + let state = context + .read() + .map_err(|error| FlowError::Internal(error.to_string()))?; + let entries = state.llm_sanitize_response_entries(&scope_locals); + (entries, subscribers) }; + let had_fallback_data = handle.data.is_some(); + let data = handle.data.clone().and_then(|data| { + NemoRelayContextState::llm_sanitize_response_snapshot_chain( + data, + LlmSanitizeResponseContext::for_response_codec(response_codec), + &entries, + ) + }); + let annotation_omitted = + (had_fallback_data && data.is_none()) || data.as_ref().is_some_and(Json::is_null); handle.optimization_recorder.close_for_finalization(None); emit_optimization_marks(handle, &subscribers); let pricing = crate::codec::response::active_pricing_resolver(); - let annotated_response = finalize_optimization_summary( - &handle.optimization_recorder, - None, - handle.model_name.as_deref(), - &pricing, - ) - .map(|summary| { - Arc::new(AnnotatedLlmResponse { - optimization_summary: Some(summary), - ..AnnotatedLlmResponse::default() + let annotated_response = (!annotation_omitted) + .then(|| { + finalize_optimization_summary( + &handle.optimization_recorder, + None, + handle.model_name.as_deref(), + &pricing, + ) }) - }); + .flatten() + .map(|summary| { + Arc::new(AnnotatedLlmResponse { + optimization_summary: Some(summary), + ..AnnotatedLlmResponse::default() + }) + }); let event = { let context = global_context(); let state = context .read() .map_err(|error| FlowError::Internal(error.to_string()))?; - state.end_llm_handle(handle, handle.data.clone(), metadata, annotated_response) + state.end_llm_handle(handle, data, metadata, annotated_response) }; if let Some(event) = sanitize_event_with_scope_stack(event, handle.captured_scope_stack()) { NemoRelayContextState::emit_event(&event, &subscribers); @@ -950,7 +1015,7 @@ pub async fn llm_call_execute(params: LlmCallExecuteParams) -> Result { &handle, &intercepted_request, annotated_request.clone(), - request_codec.as_deref(), + request_codec.clone(), &lifecycle_subscribers, )?; emit_pending_request_marks(&handle, pending_marks, &lifecycle_subscribers)?; @@ -999,8 +1064,12 @@ pub async fn llm_call_execute(params: LlmCallExecuteParams) -> Result { Err(error) => { let end_metadata = metadata_with_otel_status(metadata, "ERROR", Some(error.to_string())); - let _ = - emit_llm_end_without_output(&handle, end_metadata, Some(&lifecycle_subscribers)); + let _ = emit_llm_end_without_output( + &handle, + end_metadata, + response_codec, + Some(&lifecycle_subscribers), + ); Err(error) } } @@ -1136,7 +1205,7 @@ pub async fn llm_stream_call_execute(params: LlmStreamCallExecuteParams) -> Resu &handle, &intercepted_request, annotated_request, - request_codec.as_deref(), + request_codec.clone(), &lifecycle_subscribers, )?; emit_pending_request_marks(&handle, pending_marks, &lifecycle_subscribers)?; @@ -1180,8 +1249,12 @@ pub async fn llm_stream_call_execute(params: LlmStreamCallExecuteParams) -> Resu Err(error) => { let end_metadata = metadata_with_otel_status(metadata, "ERROR", Some(error.to_string())); - let _ = - emit_llm_end_without_output(&handle, end_metadata, Some(&lifecycle_subscribers)); + let _ = emit_llm_end_without_output( + &handle, + end_metadata, + response_codec, + Some(&lifecycle_subscribers), + ); Err(error) } } diff --git a/crates/core/src/api/registry.rs b/crates/core/src/api/registry.rs index 72199095c..984ac43d8 100644 --- a/crates/core/src/api/registry.rs +++ b/crates/core/src/api/registry.rs @@ -111,7 +111,7 @@ macro_rules! global_guardrail_registry_api { state .$field .register( - Guardrail::new(name, priority, guardrail), + Guardrail::new(name, priority, guardrail.into()), ) .map_err(FlowError::AlreadyExists) } @@ -298,7 +298,7 @@ macro_rules! scope_guardrail_registry_api { registries .$field .register( - Guardrail::new(name, priority, guardrail), + Guardrail::new(name, priority, guardrail.into()), ) .map_err(FlowError::AlreadyExists) } @@ -475,6 +475,7 @@ global_guardrail_registry_api!( mark_sanitize_guardrails, EventSanitizeFn ); + global_guardrail_registry_api!( /// Register a global scope-start event sanitizer. register_scope_sanitize_start_guardrail, diff --git a/crates/core/src/api/runtime.rs b/crates/core/src/api/runtime.rs index 8b02d998b..0261657ba 100644 --- a/crates/core/src/api/runtime.rs +++ b/crates/core/src/api/runtime.rs @@ -10,10 +10,12 @@ pub mod state; pub mod subscriber_dispatcher; pub use callbacks::{ - EventSanitizeFn, EventSubscriberFn, LlmCollectorFn, LlmConditionalFn, LlmExecutionFn, - LlmExecutionNextFn, LlmFinalizerFn, LlmJsonStream, LlmRequestInterceptFn, LlmSanitizeRequestFn, - LlmSanitizeResponseFn, LlmStreamExecutionFn, LlmStreamExecutionNextFn, LlmStreamInner, - ToolConditionalFn, ToolExecutionFn, ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, + BuiltinLlmCodec, EventSanitizeFn, EventSubscriberFn, LlmCodecIdentity, LlmCollectorFn, + LlmConditionalFn, LlmExecutionFn, LlmExecutionNextFn, LlmFinalizerFn, LlmJsonStream, + LlmRequestInterceptFn, LlmSanitizeRequestContext, LlmSanitizeRequestFn, + LlmSanitizeResponseContext, LlmSanitizeResponseFn, LlmStreamExecutionFn, + LlmStreamExecutionNextFn, LlmStreamInner, ToolConditionalFn, ToolExecutionFn, + ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, }; pub use global::global_context; pub use scope_stack::{ diff --git a/crates/core/src/api/runtime/callbacks.rs b/crates/core/src/api/runtime/callbacks.rs index 21fe853f1..e40558da1 100644 --- a/crates/core/src/api/runtime/callbacks.rs +++ b/crates/core/src/api/runtime/callbacks.rs @@ -19,6 +19,7 @@ use crate::api::event::{Event, EventSanitizeFields}; use crate::api::llm::{LlmRequest, LlmRequestInterceptOutcome}; use crate::api::tool::ToolExecutionInterceptOutcome; use crate::codec::request::AnnotatedLlmRequest; +use crate::codec::traits::{LlmCodec, LlmResponseCodec}; use crate::error::Result; use crate::json::Json; @@ -132,6 +133,165 @@ pub(crate) type ToolExecutionOutcomeNextFn = Arc< + Sync, >; +/// Relay's built-in LLM codec identities. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BuiltinLlmCodec { + /// OpenAI Chat Completions request and response payloads. + OpenAiChat, + /// OpenAI Responses request and response payloads. + OpenAiResponses, + /// Anthropic Messages request and response payloads. + AnthropicMessages, +} + +impl BuiltinLlmCodec { + /// Stable identifier used in configuration and language bindings. + #[must_use] + pub const fn id(self) -> &'static str { + match self { + Self::OpenAiChat => "openai_chat", + Self::OpenAiResponses => "openai_responses", + Self::AnthropicMessages => "anthropic_messages", + } + } +} + +/// Per-call LLM codec identity supplied to sanitize guardrails. +#[derive(Debug, Clone, PartialEq, Eq, Default)] +pub enum LlmCodecIdentity { + /// No codec was active for this payload direction. + #[default] + None, + /// A Relay built-in codec was active. + BuiltIn(BuiltinLlmCodec), + /// A runtime-registered codec was active, identified by its stable ID. + Runtime(String), + /// A codec was active but does not expose a registered identity. + Opaque, +} + +/// Per-call codec context for LLM request sanitize guardrails. +/// +/// The context distinguishes no codec, Relay built-ins, runtime-registered +/// codecs, and active codecs with no stable identity. +#[derive(Clone, Default)] +pub struct LlmSanitizeRequestContext { + /// Identity of the codec active for this payload direction. + codec: LlmCodecIdentity, + request_codec: Option>, +} + +impl std::fmt::Debug for LlmSanitizeRequestContext { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("LlmSanitizeRequestContext") + .field("codec", &self.codec) + .finish_non_exhaustive() + } +} + +impl LlmSanitizeRequestContext { + /// Construct a context that carries only a codec identity. + /// + /// Identity-only contexts do not carry a codec handle, so + /// [`Self::resolve_codec`] returns `None` even when the identity describes + /// an active codec. + #[must_use] + pub fn with_identity(codec: LlmCodecIdentity) -> Self { + Self { + codec, + ..Self::default() + } + } + + /// Construct request-sanitizer context from the active request codec. + #[must_use] + pub fn for_request_codec(codec: Option>) -> Self { + let identity = codec + .as_deref() + .map_or(LlmCodecIdentity::None, LlmCodec::codec_identity); + Self { + codec: identity, + request_codec: codec, + } + } + + /// Return the identity of the codec active for this payload direction. + #[must_use] + pub fn codec(&self) -> &LlmCodecIdentity { + &self.codec + } + + /// Resolve the active request codec. + /// + /// Returns `None` for contexts constructed with [`Self::with_identity`]. + #[must_use] + pub fn resolve_codec(&self) -> Option> { + self.request_codec.clone() + } +} + +/// Per-call codec context for LLM response sanitize guardrails. +/// +/// The context distinguishes no codec, Relay built-ins, runtime-registered +/// codecs, and active codecs with no stable identity. +#[derive(Clone, Default)] +pub struct LlmSanitizeResponseContext { + /// Identity of the codec active for this payload direction. + codec: LlmCodecIdentity, + response_codec: Option>, +} + +impl std::fmt::Debug for LlmSanitizeResponseContext { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("LlmSanitizeResponseContext") + .field("codec", &self.codec) + .finish_non_exhaustive() + } +} + +impl LlmSanitizeResponseContext { + /// Construct a context that carries only a codec identity. + /// + /// Identity-only contexts do not carry a codec handle, so + /// [`Self::resolve_codec`] returns `None` even when the identity describes + /// an active codec. + #[must_use] + pub fn with_identity(codec: LlmCodecIdentity) -> Self { + Self { + codec, + ..Self::default() + } + } + + /// Construct response-sanitizer context from the active response codec. + #[must_use] + pub fn for_response_codec(codec: Option>) -> Self { + let identity = codec + .as_deref() + .map_or(LlmCodecIdentity::None, LlmResponseCodec::codec_identity); + Self { + codec: identity, + response_codec: codec, + } + } + + /// Return the identity of the codec active for this payload direction. + #[must_use] + pub fn codec(&self) -> &LlmCodecIdentity { + &self.codec + } + + /// Resolve the active response codec. + /// + /// Returns `None` for contexts constructed with [`Self::with_identity`]. + #[must_use] + pub fn resolve_codec(&self) -> Option> { + self.response_codec.clone() + } +} + /// Sanitize an LLM request before the runtime records it. /// /// LLM request sanitizers affect the serialized request payload emitted on @@ -140,10 +300,16 @@ pub(crate) type ToolExecutionOutcomeNextFn = Arc< /// /// # Parameters /// - First argument: LLM request payload to sanitize for observability. +/// - Second argument: Per-call request codec identity and capability. /// /// # Returns -/// Sanitized [`LlmRequest`] for the emitted event. -pub type LlmSanitizeRequestFn = Arc LlmRequest + Send + Sync>; +/// `Some` contains the sanitized request for the emitted event. `None` omits +/// both the raw request payload and its annotation from that event. +/// +/// The context is always supplied and distinguishes no codec, built-in codecs, +/// runtime-registered codecs, and opaque active codecs. +pub type LlmSanitizeRequestFn = + Arc Option + Send + Sync>; /// Sanitize an LLM response before the runtime records it. /// /// These callbacks rewrite the JSON response payload captured on LLM-end @@ -151,10 +317,16 @@ pub type LlmSanitizeRequestFn = Arc LlmRequest + Send + Sy /// /// # Parameters /// - First argument: JSON response payload to sanitize for observability. +/// - Second argument: Per-call response codec identity and capability. /// /// # Returns -/// Sanitized JSON response payload for the emitted event. -pub type LlmSanitizeResponseFn = Arc Json + Send + Sync>; +/// `Some` contains the sanitized response for the emitted event. `None` omits +/// both the raw response payload and its annotation from that event. +/// +/// The context is always supplied and distinguishes no codec, built-in codecs, +/// runtime-registered codecs, and opaque active codecs. +pub type LlmSanitizeResponseFn = + Arc Option + Send + Sync>; /// Decide whether an LLM call is allowed to continue. /// /// The callback receives the current [`LlmRequest`] and can allow execution, diff --git a/crates/core/src/api/runtime/state.rs b/crates/core/src/api/runtime/state.rs index 60a01aa8e..eb3dbb950 100644 --- a/crates/core/src/api/runtime/state.rs +++ b/crates/core/src/api/runtime/state.rs @@ -22,7 +22,8 @@ use crate::api::llm::{LlmHandle, LlmRequest}; use crate::api::registry::{ExecutionIntercept, Guardrail, Intercept}; use crate::api::runtime::callbacks::{ EventSanitizeFn, EventSubscriberFn, LlmConditionalFn, LlmExecutionFn, LlmExecutionNextFn, - LlmRequestInterceptFn, LlmSanitizeRequestFn, LlmSanitizeResponseFn, LlmStreamExecutionFn, + LlmRequestInterceptFn, LlmSanitizeRequestContext, LlmSanitizeRequestFn, + LlmSanitizeResponseContext, LlmSanitizeResponseFn, LlmStreamExecutionFn, LlmStreamExecutionNextFn, LlmStreamExecutionRegistryRefs, ToolConditionalFn, ToolExecutionFn, ToolExecutionNextFn, ToolExecutionOutcomeNextFn, ToolInterceptFn, ToolSanitizeFn, }; @@ -965,11 +966,12 @@ impl NemoRelayContextState { /// The sanitized [`LlmRequest`] after every provided guardrail has run. pub(crate) fn llm_sanitize_request_snapshot_chain( request: LlmRequest, + context: LlmSanitizeRequestContext, entries: &[Guardrail], - ) -> LlmRequest { - let mut value = request; + ) -> Option { + let mut value = Some(request); for entry in entries { - value = (entry.payload)(value); + value = value.and_then(|value| (entry.payload)(value, context.clone())); } value } @@ -1003,11 +1005,12 @@ impl NemoRelayContextState { /// The sanitized response payload after every provided guardrail has run. pub(crate) fn llm_sanitize_response_snapshot_chain( response: Json, + context: LlmSanitizeResponseContext, entries: &[Guardrail], - ) -> Json { - let mut value = response; + ) -> Option { + let mut value = Some(response); for entry in entries { - value = (entry.payload)(value); + value = value.and_then(|value| (entry.payload)(value, context.clone())); } value } diff --git a/crates/core/src/codec/anthropic.rs b/crates/core/src/codec/anthropic.rs index 9272aa062..88272d748 100644 --- a/crates/core/src/codec/anthropic.rs +++ b/crates/core/src/codec/anthropic.rs @@ -19,6 +19,7 @@ use serde::Deserialize; use crate::api::llm::LlmRequest; +use crate::api::runtime::{BuiltinLlmCodec, LlmCodecIdentity}; use crate::error::{FlowError, Result}; use crate::json::Json; @@ -922,6 +923,10 @@ fn anthropic_usage( // --------------------------------------------------------------------------- impl LlmResponseCodec for AnthropicMessagesCodec { + fn codec_identity(&self) -> LlmCodecIdentity { + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::AnthropicMessages) + } + fn decode_response(&self, response: &Json) -> Result { let raw: RawAnthropicResponse = serde_json::from_value(response.clone()) .map_err(|e| FlowError::Internal(format!("Anthropic Messages response decode: {e}")))?; @@ -968,6 +973,10 @@ impl LlmResponseCodec for AnthropicMessagesCodec { // --------------------------------------------------------------------------- impl LlmCodec for AnthropicMessagesCodec { + fn codec_identity(&self) -> LlmCodecIdentity { + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::AnthropicMessages) + } + fn decode(&self, request: &LlmRequest) -> Result { let obj = request .content diff --git a/crates/core/src/codec/openai_chat.rs b/crates/core/src/codec/openai_chat.rs index 4717426ac..dc3c46dca 100644 --- a/crates/core/src/codec/openai_chat.rs +++ b/crates/core/src/codec/openai_chat.rs @@ -9,6 +9,7 @@ use serde::Deserialize; use crate::api::llm::LlmRequest; +use crate::api::runtime::{BuiltinLlmCodec, LlmCodecIdentity}; use crate::error::{FlowError, Result}; use crate::json::Json; @@ -1116,6 +1117,10 @@ fn patch_chat_api_fields( // --------------------------------------------------------------------------- impl LlmResponseCodec for OpenAIChatCodec { + fn codec_identity(&self) -> LlmCodecIdentity { + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiChat) + } + fn decode_response(&self, response: &Json) -> Result { let raw: RawChatCompletion = serde_json::from_value(response.clone()) .map_err(|e| FlowError::Internal(format!("OpenAI Chat response decode: {e}")))?; @@ -1205,6 +1210,10 @@ impl LlmResponseCodec for OpenAIChatCodec { // --------------------------------------------------------------------------- impl LlmCodec for OpenAIChatCodec { + fn codec_identity(&self) -> LlmCodecIdentity { + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiChat) + } + fn decode(&self, request: &LlmRequest) -> Result { let obj = request .content diff --git a/crates/core/src/codec/openai_responses.rs b/crates/core/src/codec/openai_responses.rs index a36e5e1cd..279cf3d0f 100644 --- a/crates/core/src/codec/openai_responses.rs +++ b/crates/core/src/codec/openai_responses.rs @@ -18,6 +18,7 @@ use serde::Deserialize; use crate::api::llm::LlmRequest; +use crate::api::runtime::{BuiltinLlmCodec, LlmCodecIdentity}; use crate::error::{FlowError, Result}; use crate::json::Json; @@ -1194,6 +1195,10 @@ fn patch_responses_common_fields( // --------------------------------------------------------------------------- impl LlmResponseCodec for OpenAIResponsesCodec { + fn codec_identity(&self) -> LlmCodecIdentity { + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiResponses) + } + fn decode_response(&self, response: &Json) -> Result { let raw: RawResponsesResponse = serde_json::from_value(response.clone()) .map_err(|e| FlowError::Internal(format!("OpenAI Responses response decode: {e}")))?; @@ -1275,6 +1280,10 @@ impl LlmResponseCodec for OpenAIResponsesCodec { // --------------------------------------------------------------------------- impl LlmCodec for OpenAIResponsesCodec { + fn codec_identity(&self) -> LlmCodecIdentity { + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiResponses) + } + fn decode(&self, request: &LlmRequest) -> Result { let obj = request .content diff --git a/crates/core/src/codec/traits.rs b/crates/core/src/codec/traits.rs index 1c62647c7..379f3c263 100644 --- a/crates/core/src/codec/traits.rs +++ b/crates/core/src/codec/traits.rs @@ -4,6 +4,7 @@ //! LLM codec traits for bidirectional request translation. use crate::api::llm::LlmRequest; +use crate::api::runtime::LlmCodecIdentity; use crate::error::Result; use crate::json::Json; @@ -33,6 +34,15 @@ use super::response::AnnotatedLlmResponse; /// so the Rust core cannot know concrete types at compile time. /// Store as `Arc`. pub trait LlmCodec: Send + Sync { + /// Return this codec's identity for LLM sanitizer context. + /// + /// Custom codecs should keep the default [`LlmCodecIdentity::Opaque`] unless + /// they have a stable runtime registration ID. Callers must not infer a + /// provider surface from an opaque implementation or request shape. + fn codec_identity(&self) -> LlmCodecIdentity { + LlmCodecIdentity::Opaque + } + /// Parse opaque request content into structured form. fn decode(&self, request: &LlmRequest) -> Result; @@ -72,6 +82,14 @@ pub trait LlmCodec: Send + Sync { /// 1. Deserialize raw JSON into API-specific intermediate structs /// 2. Map intermediate structs into the normalized `AnnotatedLlmResponse` pub trait LlmResponseCodec: Send + Sync { + /// Return this codec's identity for LLM sanitizer context. + /// + /// Custom codecs should keep the default [`LlmCodecIdentity::Opaque`] unless + /// they have a stable runtime registration ID. + fn codec_identity(&self) -> LlmCodecIdentity { + LlmCodecIdentity::Opaque + } + /// Parse a raw JSON response into normalized structured form. /// /// Implementations should return `Err` only for genuinely unparseable input. diff --git a/crates/core/src/plugin/dynamic/native.rs b/crates/core/src/plugin/dynamic/native.rs index be1b658a9..f1217c746 100644 --- a/crates/core/src/plugin/dynamic/native.rs +++ b/crates/core/src/plugin/dynamic/native.rs @@ -3,9 +3,14 @@ //! Native dynamic plugin loader and host-side ABI adapter. +#[cfg(test)] +use std::cell::Cell; use std::cell::RefCell; +#[cfg(test)] +use std::collections::HashSet; use std::ffi::c_void; use std::future::Future; +use std::panic::{AssertUnwindSafe, catch_unwind}; use std::path::{Path, PathBuf}; use std::pin::Pin; use std::ptr; @@ -17,13 +22,16 @@ use libloading::{Library, Symbol}; use nemo_relay_plugin::{ NEMO_RELAY_NATIVE_ABI_VERSION, NemoRelayNativeEventSanitizeCb, NemoRelayNativeEventSubscriberCb, NemoRelayNativeFreeFn, NemoRelayNativeHostApiV1, - NemoRelayNativeJsonCb, NemoRelayNativeLlmConditionalCb, NemoRelayNativeLlmExecutionCb, - NemoRelayNativeLlmRequestCb, NemoRelayNativeLlmRequestInterceptCb, - NemoRelayNativeLlmStreamExecutionCb, NemoRelayNativeLlmStreamV1, NemoRelayNativePluginContext, - NemoRelayNativePluginEntry, NemoRelayNativePluginV1, NemoRelayNativeScopeHandle, - NemoRelayNativeScopeStack, NemoRelayNativeScopeStackBinding, NemoRelayNativeScopeType, - NemoRelayNativeString, NemoRelayNativeToolConditionalCb, NemoRelayNativeToolExecutionCb, - NemoRelayNativeToolJsonCb, NemoRelayNativeWithScopeStackCb, NemoRelayStatus, + NemoRelayNativeLlmCodecKind, NemoRelayNativeLlmConditionalCb, NemoRelayNativeLlmExecutionCb, + NemoRelayNativeLlmRequestCodec, NemoRelayNativeLlmRequestInterceptCb, + NemoRelayNativeLlmResponseCodec, NemoRelayNativeLlmSanitizeRequestCb, + NemoRelayNativeLlmSanitizeRequestContext, NemoRelayNativeLlmSanitizeResponseCb, + NemoRelayNativeLlmSanitizeResponseContext, NemoRelayNativeLlmStreamExecutionCb, + NemoRelayNativeLlmStreamV1, NemoRelayNativePluginContext, NemoRelayNativePluginEntry, + NemoRelayNativePluginV1, NemoRelayNativeScopeHandle, NemoRelayNativeScopeStack, + NemoRelayNativeScopeStackBinding, NemoRelayNativeScopeType, NemoRelayNativeString, + NemoRelayNativeToolConditionalCb, NemoRelayNativeToolExecutionCb, NemoRelayNativeToolJsonCb, + NemoRelayNativeWithScopeStackCb, NemoRelayStatus, }; use semver::{Version, VersionReq}; use serde_json::{Map, Value as Json}; @@ -34,10 +42,11 @@ use tokio_stream::{Stream, StreamExt}; use crate::api::event::{Event, EventSanitizeFields}; use crate::api::llm::{LlmRequest, LlmRequestInterceptOutcome}; use crate::api::runtime::{ - EventSanitizeFn, EventSubscriberFn, LlmConditionalFn, LlmExecutionFn, LlmExecutionNextFn, - LlmJsonStream, LlmRequestInterceptFn, LlmSanitizeRequestFn, LlmSanitizeResponseFn, - LlmStreamExecutionFn, LlmStreamExecutionNextFn, ToolConditionalFn, ToolExecutionFn, - ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, + EventSanitizeFn, EventSubscriberFn, LlmCodecIdentity, LlmConditionalFn, LlmExecutionFn, + LlmExecutionNextFn, LlmJsonStream, LlmRequestInterceptFn, LlmSanitizeRequestContext, + LlmSanitizeRequestFn, LlmSanitizeResponseContext, LlmSanitizeResponseFn, LlmStreamExecutionFn, + LlmStreamExecutionNextFn, ToolConditionalFn, ToolExecutionFn, ToolExecutionNextFn, + ToolInterceptFn, ToolSanitizeFn, }; use crate::api::runtime::{ ScopeStackHandle, ThreadScopeStackBinding, capture_thread_scope_stack, create_scope_stack, @@ -49,6 +58,8 @@ use crate::api::scope::{ }; use crate::api::scope::{event as emit_scope_mark, get_handle, pop_scope, push_scope}; use crate::api::tool::ToolExecutionInterceptOutcome; +use crate::codec::request::AnnotatedLlmRequest; +use crate::codec::traits::{LlmCodec, LlmResponseCodec}; use crate::error::{FlowError, Result as FlowResult}; use crate::plugin::{ ConfigDiagnostic, DiagnosticLevel, Plugin, PluginError, PluginRegistrationContext, @@ -491,6 +502,9 @@ struct NativeHostPluginContext { struct NativeHostString(Vec); +struct NativeHostLlmRequestCodec(Arc); +struct NativeHostLlmResponseCodec(Arc); + struct NativeHostScopeHandle(ScopeHandle); struct NativeHostScopeStack(ScopeStackHandle); @@ -499,6 +513,10 @@ struct NativeHostScopeStackBinding(ThreadScopeStackBinding); thread_local! { static NATIVE_LAST_ERROR: RefCell> = const { RefCell::new(None) }; + #[cfg(test)] + static NATIVE_STRING_LIVE_ALLOCATIONS: RefCell> = RefCell::new(HashSet::new()); + #[cfg(test)] + static NATIVE_STRING_FAIL_AFTER: Cell> = const { Cell::new(None) }; } fn set_native_last_error(message: impl Into) { @@ -513,6 +531,16 @@ fn native_last_error_message() -> Option { NATIVE_LAST_ERROR.with(|cell| cell.borrow().clone()) } +#[cfg(test)] +fn fail_native_string_allocation_after(successful_allocations: usize) { + NATIVE_STRING_FAIL_AFTER.with(|cell| cell.set(Some(successful_allocations))); +} + +#[cfg(test)] +fn native_string_live_allocations() -> usize { + NATIVE_STRING_LIVE_ALLOCATIONS.with(|allocations| allocations.borrow().len()) +} + unsafe extern "C" fn native_string_new( data: *const u8, len: usize, @@ -536,8 +564,29 @@ unsafe extern "C" fn native_string_new( set_native_last_error(format!("string data is not valid UTF-8: {err}")); return NemoRelayStatus::InvalidUtf8; } + #[cfg(test)] + let should_fail = NATIVE_STRING_FAIL_AFTER.with(|cell| match cell.get() { + Some(0) => { + cell.set(None); + true + } + Some(remaining) => { + cell.set(Some(remaining - 1)); + false + } + None => false, + }); + #[cfg(test)] + if should_fail { + set_native_last_error("injected native string allocation failure"); + return NemoRelayStatus::Internal; + } let handle = Box::new(NativeHostString(bytes.to_vec())); unsafe { *out = Box::into_raw(handle) as *mut NemoRelayNativeString }; + #[cfg(test)] + NATIVE_STRING_LIVE_ALLOCATIONS.with(|allocations| { + allocations.borrow_mut().insert(unsafe { *out } as usize); + }); NemoRelayStatus::Ok } @@ -560,6 +609,10 @@ unsafe extern "C" fn native_string_len(value: *const NemoRelayNativeString) -> u unsafe extern "C" fn native_string_free(value: *mut NemoRelayNativeString) { if !value.is_null() { drop(unsafe { Box::from_raw(value as *mut NativeHostString) }); + #[cfg(test)] + NATIVE_STRING_LIVE_ALLOCATIONS.with(|allocations| { + allocations.borrow_mut().remove(&(value as usize)); + }); } } @@ -574,6 +627,162 @@ unsafe extern "C" fn native_last_error_set(message: *const NemoRelayNativeString } } +unsafe extern "C" fn native_llm_request_codec_decode( + codec: *const NemoRelayNativeLlmRequestCodec, + request_json: *const NemoRelayNativeString, + out: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + clear_native_last_error(); + if out.is_null() { + set_native_last_error("request codec decode output pointer is null"); + return NemoRelayStatus::NullPointer; + } + unsafe { *out = ptr::null_mut() }; + if codec.is_null() { + set_native_last_error("request codec decode capability is null"); + return NemoRelayStatus::NullPointer; + } + if request_json.is_null() { + set_native_last_error("request codec decode request is null"); + return NemoRelayStatus::NullPointer; + } + let result = catch_unwind(AssertUnwindSafe(|| -> std::result::Result<_, String> { + let request: LlmRequest = serde_json::from_str( + &read_native_string(request_json).map_err(|error| error.to_string())?, + ) + .map_err(|error| format!("invalid request JSON: {error}"))?; + let codec = unsafe { &*(codec as *const NativeHostLlmRequestCodec) }; + let annotated = codec + .0 + .decode(&request) + .map_err(|error| error.to_string())?; + let annotated = serde_json::to_value(annotated).map_err(|error| error.to_string())?; + native_string_from_json(&annotated) + .ok_or_else(|| "failed to allocate decoded request".to_string()) + })); + match result { + Ok(Ok(value)) => { + unsafe { *out = value }; + NemoRelayStatus::Ok + } + Ok(Err(error)) => { + set_native_last_error(format!("request codec decode failed: {error}")); + NemoRelayStatus::Internal + } + Err(_) => { + set_native_last_error("request codec decode panicked"); + NemoRelayStatus::Internal + } + } +} + +unsafe extern "C" fn native_llm_request_codec_encode( + codec: *const NemoRelayNativeLlmRequestCodec, + annotated_json: *const NemoRelayNativeString, + original_json: *const NemoRelayNativeString, + out: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + clear_native_last_error(); + if out.is_null() { + set_native_last_error("request codec encode output pointer is null"); + return NemoRelayStatus::NullPointer; + } + unsafe { *out = ptr::null_mut() }; + if codec.is_null() { + set_native_last_error("request codec encode capability is null"); + return NemoRelayStatus::NullPointer; + } + if annotated_json.is_null() { + set_native_last_error("request codec encode annotated request is null"); + return NemoRelayStatus::NullPointer; + } + if original_json.is_null() { + set_native_last_error("request codec encode original request is null"); + return NemoRelayStatus::NullPointer; + } + let result = catch_unwind(AssertUnwindSafe(|| -> std::result::Result<_, String> { + let annotated: AnnotatedLlmRequest = serde_json::from_str( + &read_native_string(annotated_json).map_err(|error| error.to_string())?, + ) + .map_err(|error| format!("invalid annotated request JSON: {error}"))?; + let original: LlmRequest = serde_json::from_str( + &read_native_string(original_json).map_err(|error| error.to_string())?, + ) + .map_err(|error| format!("invalid original request JSON: {error}"))?; + let codec = unsafe { &*(codec as *const NativeHostLlmRequestCodec) }; + let request = codec + .0 + .encode(&annotated, &original) + .map_err(|error| error.to_string())?; + let request = serde_json::to_value(request).map_err(|error| error.to_string())?; + native_string_from_json(&request) + .ok_or_else(|| "failed to allocate encoded request".to_string()) + })); + match result { + Ok(Ok(value)) => { + unsafe { *out = value }; + NemoRelayStatus::Ok + } + Ok(Err(error)) => { + set_native_last_error(format!("request codec encode failed: {error}")); + NemoRelayStatus::Internal + } + Err(_) => { + set_native_last_error("request codec encode panicked"); + NemoRelayStatus::Internal + } + } +} + +unsafe extern "C" fn native_llm_response_codec_decode( + codec: *const NemoRelayNativeLlmResponseCodec, + response_json: *const NemoRelayNativeString, + out: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + clear_native_last_error(); + if out.is_null() { + set_native_last_error("response codec decode output pointer is null"); + return NemoRelayStatus::NullPointer; + } + unsafe { *out = ptr::null_mut() }; + if codec.is_null() { + set_native_last_error("response codec decode capability is null"); + return NemoRelayStatus::NullPointer; + } + if response_json.is_null() { + set_native_last_error("response codec decode response is null"); + return NemoRelayStatus::NullPointer; + } + let result = catch_unwind(AssertUnwindSafe(|| -> std::result::Result<_, String> { + let response: Json = serde_json::from_str( + &read_native_string(response_json).map_err(|error| error.to_string())?, + ) + .map_err(|error| format!("invalid response JSON: {error}"))?; + let codec = unsafe { &*(codec as *const NativeHostLlmResponseCodec) }; + let annotated = codec + .0 + .decode_response(&response) + .map_err(|error| error.to_string())?; + let annotated = serde_json::to_value(annotated).map_err(|error| error.to_string())?; + native_string_from_json(&annotated) + .ok_or_else(|| "failed to allocate decoded response".to_string()) + })); + match result { + Ok(Ok(value)) => { + unsafe { *out = value }; + NemoRelayStatus::Ok + } + Ok(Err(error)) => { + set_native_last_error(format!("response codec decode failed: {error}")); + NemoRelayStatus::Internal + } + Err(_) => { + set_native_last_error("response codec decode panicked"); + NemoRelayStatus::Internal + } + } +} + fn native_host_api() -> *const NemoRelayNativeHostApiV1 { static HOST_API: OnceLock = OnceLock::new(); static RELAY_VERSION: &[u8] = concat!(env!("CARGO_PKG_VERSION"), "\0").as_bytes(); @@ -587,6 +796,9 @@ fn native_host_api() -> *const NemoRelayNativeHostApiV1 { string_free: native_string_free, last_error_clear: native_last_error_clear, last_error_set: native_last_error_set, + llm_request_codec_decode: native_llm_request_codec_decode, + llm_request_codec_encode: native_llm_request_codec_encode, + llm_response_codec_decode: native_llm_response_codec_decode, plugin_context_register_subscriber: native_plugin_context_register_subscriber, plugin_context_register_tool_sanitize_request_guardrail: native_plugin_context_register_tool_sanitize_request_guardrail, @@ -681,6 +893,22 @@ fn take_json_from_native_string( result } +unsafe fn free_native_sanitizer_strings( + input: *mut NemoRelayNativeString, + codec_id: Option<*mut NemoRelayNativeString>, + output: *mut NemoRelayNativeString, +) { + unsafe { native_string_free(input) }; + if let Some(codec_id) = codec_id + && codec_id != input + { + unsafe { native_string_free(codec_id) }; + } + if !output.is_null() && output != input && Some(output) != codec_id { + unsafe { native_string_free(output) }; + } +} + fn optional_json_from_native_string( value: *const NemoRelayNativeString, field: &str, @@ -1261,7 +1489,7 @@ unsafe extern "C" fn native_plugin_context_register_llm_sanitize_request_guardra ctx: *mut NemoRelayNativePluginContext, name: *const NemoRelayNativeString, priority: i32, - cb: NemoRelayNativeLlmRequestCb, + cb: NemoRelayNativeLlmSanitizeRequestCb, user_data: *mut c_void, free_fn: NemoRelayNativeFreeFn, ) -> NemoRelayStatus { @@ -1279,7 +1507,7 @@ unsafe extern "C" fn native_plugin_context_register_llm_sanitize_request_guardra match ctx.register_llm_sanitize_request_guardrail( &name, priority, - wrap_llm_request_fn(instance, cb, user_data, free_fn), + wrap_llm_sanitize_request_fn(instance, cb, user_data, free_fn), ) { Ok(()) => NemoRelayStatus::Ok, Err(err) => status_from_plugin_error(err), @@ -1290,7 +1518,7 @@ unsafe extern "C" fn native_plugin_context_register_llm_sanitize_response_guardr ctx: *mut NemoRelayNativePluginContext, name: *const NemoRelayNativeString, priority: i32, - cb: NemoRelayNativeJsonCb, + cb: NemoRelayNativeLlmSanitizeResponseCb, user_data: *mut c_void, free_fn: NemoRelayNativeFreeFn, ) -> NemoRelayStatus { @@ -1308,7 +1536,7 @@ unsafe extern "C" fn native_plugin_context_register_llm_sanitize_response_guardr match ctx.register_llm_sanitize_response_guardrail( &name, priority, - wrap_json_fn(instance, cb, user_data, free_fn), + wrap_llm_sanitize_response_fn(instance, cb, user_data, free_fn), ) { Ok(()) => NemoRelayStatus::Ok, Err(err) => status_from_plugin_error(err), @@ -1725,74 +1953,161 @@ unsafe extern "C" fn native_tool_next( } } -fn wrap_llm_request_fn( +fn wrap_llm_sanitize_request_fn( instance: Arc, - cb: NemoRelayNativeLlmRequestCb, + cb: NemoRelayNativeLlmSanitizeRequestCb, user_data: *mut c_void, free_fn: NemoRelayNativeFreeFn, ) -> LlmSanitizeRequestFn { let user_data = make_user_data(instance, user_data, free_fn); - Arc::new(move |request| { - call_llm_request_callback(cb, user_data.ptr, &request).unwrap_or_else(|_| LlmRequest { - headers: Map::new(), - content: Json::Null, - }) + Arc::new(move |request, context| { + call_llm_sanitize_request_callback(cb, user_data.ptr, &request, context) + .ok() + .flatten() }) } -fn call_llm_request_callback( - cb: NemoRelayNativeLlmRequestCb, +fn wrap_llm_sanitize_response_fn( + instance: Arc, + cb: NemoRelayNativeLlmSanitizeResponseCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, +) -> LlmSanitizeResponseFn { + let user_data = make_user_data(instance, user_data, free_fn); + Arc::new(move |payload, context| { + call_llm_sanitize_response_callback(cb, user_data.ptr, &payload, context) + .ok() + .flatten() + }) +} + +fn call_llm_sanitize_request_callback( + cb: NemoRelayNativeLlmSanitizeRequestCb, user_data: *mut c_void, request: &LlmRequest, -) -> FlowResult { + context: LlmSanitizeRequestContext, +) -> FlowResult> { clear_native_last_error(); - let request_json = serde_json::to_value(request) - .map_err(|err| FlowError::Internal(format!("failed to serialize LLM request: {err}")))?; - let request_string = native_string_from_json(&request_json) - .ok_or_else(|| FlowError::Internal("failed to allocate native LLM request".into()))?; + let codec = context.resolve_codec().map(NativeHostLlmRequestCodec); + let (codec_kind, context_id) = native_llm_codec_identity(context.codec())?; + let request_json = match serde_json::to_value(request) { + Ok(request_json) => request_json, + Err(err) => { + if let Some(context_id) = context_id { + unsafe { native_string_free(context_id) }; + } + return Err(FlowError::Internal(format!( + "failed to serialize LLM request: {err}" + ))); + } + }; + let request_string = match native_string_from_json(&request_json) { + Some(request_string) => request_string, + None => { + if let Some(context_id) = context_id { + unsafe { native_string_free(context_id) }; + } + return Err(FlowError::Internal( + "failed to allocate native LLM request".into(), + )); + } + }; + let context = NemoRelayNativeLlmSanitizeRequestContext { + codec_kind, + codec_id: context_id.map_or(ptr::null(), |value| value.cast_const()), + codec: codec + .as_ref() + .map_or(ptr::null(), |value| std::ptr::from_ref(value).cast()), + }; let mut out = ptr::null_mut(); - let status = unsafe { cb(user_data, request_string, &mut out) }; - unsafe { native_string_free(request_string) }; + let status = unsafe { cb(user_data, request_string, context, &mut out) }; if status != NemoRelayStatus::Ok { - if !out.is_null() { - unsafe { native_string_free(out) }; - } + unsafe { free_native_sanitizer_strings(request_string, context_id, out) }; return Err(flow_error_from_status( status, - "native LLM request callback failed", + "native LLM sanitize-request callback failed", )); } + if out.is_null() { + unsafe { free_native_sanitizer_strings(request_string, context_id, out) }; + return Ok(None); + } let result_json = - take_json_from_native_string(out, "native LLM request callback returned null")?; + json_from_native_string(out, "native LLM sanitize-request returned invalid JSON"); + unsafe { free_native_sanitizer_strings(request_string, context_id, out) }; + let result_json = result_json?; serde_json::from_value(result_json) + .map(Some) .map_err(|err| FlowError::Internal(format!("invalid LLM request JSON: {err}"))) } -fn wrap_json_fn( - instance: Arc, - cb: NemoRelayNativeJsonCb, +fn call_llm_sanitize_response_callback( + cb: NemoRelayNativeLlmSanitizeResponseCb, user_data: *mut c_void, - free_fn: NemoRelayNativeFreeFn, -) -> LlmSanitizeResponseFn { - let user_data = make_user_data(instance, user_data, free_fn); - Arc::new(move |payload| { - clear_native_last_error(); - let payload_string = native_string_from_json(&payload); - let Some(payload_string) = payload_string else { - return Json::Null; - }; - let mut out = ptr::null_mut(); - let status = unsafe { cb(user_data.ptr, payload_string, &mut out) }; - unsafe { native_string_free(payload_string) }; - if status != NemoRelayStatus::Ok { - if !out.is_null() { - unsafe { native_string_free(out) }; + payload: &Json, + context: LlmSanitizeResponseContext, +) -> FlowResult> { + clear_native_last_error(); + let codec = context.resolve_codec().map(NativeHostLlmResponseCodec); + let (codec_kind, context_id) = native_llm_codec_identity(context.codec())?; + let payload_string = match native_string_from_json(payload) { + Some(payload_string) => payload_string, + None => { + if let Some(context_id) = context_id { + unsafe { native_string_free(context_id) }; } - return Json::Null; + return Err(FlowError::Internal( + "failed to allocate native LLM response".into(), + )); } - take_json_from_native_string(out, "native JSON callback returned null") - .unwrap_or(Json::Null) - }) + }; + let context = NemoRelayNativeLlmSanitizeResponseContext { + codec_kind, + codec_id: context_id.map_or(ptr::null(), |value| value.cast_const()), + codec: codec + .as_ref() + .map_or(ptr::null(), |value| std::ptr::from_ref(value).cast()), + }; + let mut out = ptr::null_mut(); + let status = unsafe { cb(user_data, payload_string, context, &mut out) }; + if status != NemoRelayStatus::Ok { + unsafe { free_native_sanitizer_strings(payload_string, context_id, out) }; + return Err(flow_error_from_status( + status, + "native LLM sanitize-response callback failed", + )); + } + if out.is_null() { + unsafe { free_native_sanitizer_strings(payload_string, context_id, out) }; + return Ok(None); + } + let result = json_from_native_string(out, "native LLM sanitize-response returned invalid JSON"); + unsafe { free_native_sanitizer_strings(payload_string, context_id, out) }; + result.map(Some) +} + +fn native_llm_codec_identity( + context: &LlmCodecIdentity, +) -> FlowResult<( + NemoRelayNativeLlmCodecKind, + Option<*mut NemoRelayNativeString>, +)> { + let (codec_kind, codec_id) = match context { + LlmCodecIdentity::None => (NemoRelayNativeLlmCodecKind::None, None), + LlmCodecIdentity::BuiltIn(codec) => { + (NemoRelayNativeLlmCodecKind::BuiltIn, Some(codec.id())) + } + LlmCodecIdentity::Runtime(id) => (NemoRelayNativeLlmCodecKind::Runtime, Some(id.as_str())), + LlmCodecIdentity::Opaque => (NemoRelayNativeLlmCodecKind::Opaque, None), + }; + let codec_id = + match codec_id { + Some(codec_id) => Some(native_string_from_str(codec_id).ok_or_else(|| { + FlowError::Internal("failed to allocate native LLM codec ID".into()) + })?), + None => None, + }; + Ok((codec_kind, codec_id)) } fn wrap_llm_conditional_fn( diff --git a/crates/core/src/plugin/dynamic/worker.rs b/crates/core/src/plugin/dynamic/worker.rs index ce25e4935..79e2e70ce 100644 --- a/crates/core/src/plugin/dynamic/worker.rs +++ b/crates/core/src/plugin/dynamic/worker.rs @@ -19,10 +19,14 @@ use nemo_relay_worker_proto::v1::relay_host_runtime_server::{ use nemo_relay_worker_proto::v1::{ CancelInvocationRequest, CreateScopeStackRequest, CreateScopeStackResponse, DropScopeStackRequest, EmitMarkRequest, GuardrailResult, HandshakeRequest, HealthRequest, - HostAck, InvokeRequest, InvokeResponse, JsonEnvelope, JsonResult, LlmInvocation, - LlmNextRequest, LlmStreamNextRequest, PopScopeRequest, PushScopeRequest, PushScopeResponse, - RegisterRequest, RegisterResponse, Registration, RegistrationSurface, ScopeContext, - ShutdownRequest, StreamChunk, ToolInvocation, ToolNextRequest, ValidateRequest, WorkerError, + HostAck, InvokeRequest, InvokeResponse, JsonEnvelope, JsonResult, LlmCodecDecodeRequest, + LlmCodecDecodeResponse, LlmCodecEncodeRequest, LlmCodecIdentity as ProtoLlmCodecIdentity, + LlmCodecKind, LlmInvocation, LlmNextRequest, + LlmSanitizeRequestContext as ProtoLlmSanitizeRequestContext, + LlmSanitizeResponseContext as ProtoLlmSanitizeResponseContext, LlmStreamNextRequest, + PopScopeRequest, PushScopeRequest, PushScopeResponse, RegisterRequest, RegisterResponse, + Registration, RegistrationSurface, ScopeContext, ShutdownRequest, StreamChunk, ToolInvocation, + ToolNextRequest, ValidateRequest, WorkerError, }; use nemo_relay_worker_proto::{WORKER_PROTOCOL_GRPC_V1, decode_json_envelope, json_envelope}; use semver::{Version, VersionReq}; @@ -55,8 +59,9 @@ use tower::service_fn; use crate::api::event::{Event, EventSanitizeFields}; use crate::api::llm::{LLM_REQUEST_INTERCEPT_OUTCOME_SCHEMA, LlmRequest}; use crate::api::runtime::{ - LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn, ToolExecutionNextFn, - current_scope_stack, with_scope_stack, + LlmCodecIdentity, LlmExecutionNextFn, LlmJsonStream, LlmSanitizeRequestContext, + LlmSanitizeResponseContext, LlmStreamExecutionNextFn, ToolExecutionNextFn, current_scope_stack, + with_scope_stack, }; use crate::api::scope::{ EmitMarkEventParams, PopScopeParams, PushScopeParams, ScopeAttributes, ScopeHandle, ScopeType, @@ -64,6 +69,7 @@ use crate::api::scope::{ }; use crate::api::tool::ToolExecutionInterceptOutcome; use crate::codec::request::{ANNOTATED_LLM_REQUEST_SCHEMA, AnnotatedLlmRequest}; +use crate::codec::traits::{LlmCodec, LlmResponseCodec}; use crate::error::{FlowError, Result as FlowResult}; use crate::plugin::{ ConfigDiagnostic, DiagnosticLevel, Plugin, PluginError, PluginRegistrationContext, @@ -1048,271 +1054,278 @@ impl WorkerPluginInstance { self.plugin_kind, registration.surface )) })?; - let name = registration.local_name.clone(); - let priority = registration.priority; - let break_chain = registration.break_chain; match surface { RegistrationSurface::Subscriber => { - let instance = Arc::new(self.clone_for_callback()); - let callback_name = name.clone(); - ctx.register_subscriber( - &name, - Arc::new(move |event| { - if instance.invoke_subscriber(&callback_name, event).is_err() { - instance.log_callback_fallback( - &callback_name, - RegistrationSurface::Subscriber, - ); - } - }), - )?; + self.install_subscriber_registration(ctx, ®istration.local_name)? } RegistrationSurface::MarkSanitizeGuardrail | RegistrationSurface::ScopeSanitizeStartGuardrail - | RegistrationSurface::ScopeSanitizeEndGuardrail => { - let instance = Arc::new(self.clone_for_callback()); - let callback_name = name.clone(); - let callback = Arc::new(move |event: &Event, _fields: EventSanitizeFields| { - instance - .invoke_event_sanitize(&callback_name, surface, event) - .unwrap_or_else(|_| { - instance.log_callback_fallback(&callback_name, surface); - EventSanitizeFields::default() - }) - }); - match surface { - RegistrationSurface::MarkSanitizeGuardrail => { - ctx.register_mark_sanitize_guardrail(&name, priority, callback)? - } - RegistrationSurface::ScopeSanitizeStartGuardrail => { - ctx.register_scope_sanitize_start_guardrail(&name, priority, callback)? - } - RegistrationSurface::ScopeSanitizeEndGuardrail => { - ctx.register_scope_sanitize_end_guardrail(&name, priority, callback)? - } - _ => unreachable!(), - } + | RegistrationSurface::ScopeSanitizeEndGuardrail => self + .install_event_sanitize_registration( + ctx, + ®istration.local_name, + registration.priority, + surface, + )?, + RegistrationSurface::ToolSanitizeRequestGuardrail + | RegistrationSurface::ToolSanitizeResponseGuardrail + | RegistrationSurface::ToolConditionalExecutionGuardrail + | RegistrationSurface::ToolRequestIntercept + | RegistrationSurface::ToolExecutionIntercept => { + self.install_tool_registration(ctx, registration, surface)? } - RegistrationSurface::ToolSanitizeRequestGuardrail => { - let instance = Arc::new(self.clone_for_callback()); - let callback_name = name.clone(); - ctx.register_tool_sanitize_request_guardrail( - &name, - priority, - Arc::new(move |tool_name, value| { - instance - .invoke_tool_json( - &callback_name, - RegistrationSurface::ToolSanitizeRequestGuardrail, - tool_name, - value.clone(), - None, - ) - .unwrap_or_else(|_| { - instance.log_callback_fallback( - &callback_name, - RegistrationSurface::ToolSanitizeRequestGuardrail, - ); - value - }) - }), - )?; + RegistrationSurface::LlmSanitizeRequestGuardrail + | RegistrationSurface::LlmSanitizeResponseGuardrail + | RegistrationSurface::LlmConditionalExecutionGuardrail + | RegistrationSurface::LlmRequestIntercept + | RegistrationSurface::LlmExecutionIntercept + | RegistrationSurface::LlmStreamExecutionIntercept => { + self.install_llm_registration(ctx, registration, surface)? } - RegistrationSurface::ToolSanitizeResponseGuardrail => { - let instance = Arc::new(self.clone_for_callback()); - let callback_name = name.clone(); - ctx.register_tool_sanitize_response_guardrail( - &name, - priority, - Arc::new(move |tool_name, value| { - instance - .invoke_tool_json( - &callback_name, - RegistrationSurface::ToolSanitizeResponseGuardrail, - tool_name, - value.clone(), - None, - ) - .unwrap_or_else(|_| { - instance.log_callback_fallback( - &callback_name, - RegistrationSurface::ToolSanitizeResponseGuardrail, - ); - value - }) - }), - )?; + RegistrationSurface::Unspecified => { + return Err(PluginError::RegistrationFailed(format!( + "worker plugin '{}' returned unspecified registration surface", + self.plugin_kind + ))); } - RegistrationSurface::ToolConditionalExecutionGuardrail => { - let instance = Arc::new(self.clone_for_callback()); - let callback_name = name.clone(); - ctx.register_tool_conditional_execution_guardrail( - &name, - priority, - Arc::new(move |tool_name, value| { - instance.invoke_tool_guardrail(&callback_name, tool_name, value.clone()) - }), - )?; + } + } + Ok(()) + } + + fn install_subscriber_registration( + &self, + ctx: &mut PluginRegistrationContext, + name: &str, + ) -> crate::plugin::Result<()> { + let instance = Arc::new(self.clone_for_callback()); + let callback_name = name.to_owned(); + ctx.register_subscriber( + name, + Arc::new(move |event| { + if instance.invoke_subscriber(&callback_name, event).is_err() { + instance.log_callback_fallback(&callback_name, RegistrationSurface::Subscriber); } - RegistrationSurface::ToolRequestIntercept => { - let instance = Arc::new(self.clone_for_callback()); - let callback_name = name.clone(); - ctx.register_tool_request_intercept( - &name, - priority, - break_chain, - Arc::new(move |tool_name, value| { - instance.invoke_tool_json( + }), + ) + } + + fn install_event_sanitize_registration( + &self, + ctx: &mut PluginRegistrationContext, + name: &str, + priority: i32, + surface: RegistrationSurface, + ) -> crate::plugin::Result<()> { + let instance = Arc::new(self.clone_for_callback()); + let callback_name = name.to_owned(); + let callback = Arc::new(move |event: &Event, _fields: EventSanitizeFields| { + instance + .invoke_event_sanitize(&callback_name, surface, event) + .unwrap_or_else(|_| { + instance.log_callback_fallback(&callback_name, surface); + EventSanitizeFields::default() + }) + }); + match surface { + RegistrationSurface::MarkSanitizeGuardrail => { + ctx.register_mark_sanitize_guardrail(name, priority, callback) + } + RegistrationSurface::ScopeSanitizeStartGuardrail => { + ctx.register_scope_sanitize_start_guardrail(name, priority, callback) + } + RegistrationSurface::ScopeSanitizeEndGuardrail => { + ctx.register_scope_sanitize_end_guardrail(name, priority, callback) + } + _ => unreachable!("event sanitizer surface was pre-filtered"), + } + } + + fn install_tool_registration( + &self, + ctx: &mut PluginRegistrationContext, + registration: &Registration, + surface: RegistrationSurface, + ) -> crate::plugin::Result<()> { + let name = registration.local_name.as_str(); + let priority = registration.priority; + let instance = Arc::new(self.clone_for_callback()); + let callback_name = name.to_owned(); + match surface { + RegistrationSurface::ToolSanitizeRequestGuardrail => ctx + .register_tool_sanitize_request_guardrail( + name, + priority, + Arc::new(move |tool_name, value| { + instance + .invoke_tool_json( &callback_name, - RegistrationSurface::ToolRequestIntercept, + surface, tool_name, - value, + value.clone(), None, ) - }), - )?; - } - RegistrationSurface::ToolExecutionIntercept => { - let instance = Arc::new(self.clone_for_callback()); - let callback_name = name.clone(); - ctx.register_tool_execution_intercept( - &name, - priority, - Arc::new(move |tool_name, value, next| { - let instance = instance.clone(); - let name = callback_name.clone(); - let tool_name = tool_name.to_string(); - Box::pin(async move { - instance - .invoke_tool_execution(&name, &tool_name, value, next) - .await + .unwrap_or_else(|_| { + instance.log_callback_fallback(&callback_name, surface); + value }) - }), - )?; - } - RegistrationSurface::LlmSanitizeRequestGuardrail => { - let instance = Arc::new(self.clone_for_callback()); - let callback_name = name.clone(); - ctx.register_llm_sanitize_request_guardrail( - &name, - priority, - Arc::new(move |request| { - instance - .invoke_llm_request_json( - &callback_name, - RegistrationSurface::LlmSanitizeRequestGuardrail, - "", - request.clone(), - None, - None, - ) - .unwrap_or_else(|_| { - instance.log_callback_fallback( - &callback_name, - RegistrationSurface::LlmSanitizeRequestGuardrail, - ); - request - }) - }), - )?; - } - RegistrationSurface::LlmSanitizeResponseGuardrail => { - let instance = Arc::new(self.clone_for_callback()); - let callback_name = name.clone(); - ctx.register_llm_sanitize_response_guardrail( - &name, - priority, - Arc::new(move |value| { - instance - .invoke_llm_response_json( - &callback_name, - RegistrationSurface::LlmSanitizeResponseGuardrail, - "", - value.clone(), - ) - .unwrap_or_else(|_| { - instance.log_callback_fallback( - &callback_name, - RegistrationSurface::LlmSanitizeResponseGuardrail, - ); - value - }) - }), - )?; - } - RegistrationSurface::LlmConditionalExecutionGuardrail => { - let instance = Arc::new(self.clone_for_callback()); - let callback_name = name.clone(); - ctx.register_llm_conditional_execution_guardrail( - &name, - priority, - Arc::new(move |request| { - instance.invoke_llm_guardrail(&callback_name, request.clone()) - }), - )?; - } - RegistrationSurface::LlmRequestIntercept => { - let instance = Arc::new(self.clone_for_callback()); - let callback_name = name.clone(); - ctx.register_llm_request_intercept( - &name, - priority, - break_chain, - Arc::new(move |model_name, request, annotated| { - instance.invoke_llm_request_intercept( + }), + ), + RegistrationSurface::ToolSanitizeResponseGuardrail => ctx + .register_tool_sanitize_response_guardrail( + name, + priority, + Arc::new(move |tool_name, value| { + instance + .invoke_tool_json( &callback_name, - model_name, - request, - annotated, + surface, + tool_name, + value.clone(), + None, ) - }), - )?; - } - RegistrationSurface::LlmExecutionIntercept => { - let instance = Arc::new(self.clone_for_callback()); - let callback_name = name.clone(); - ctx.register_llm_execution_intercept( - &name, - priority, - Arc::new(move |model_name, request, next| { - let instance = instance.clone(); - let name = callback_name.clone(); - let model_name = model_name.to_string(); - Box::pin(async move { - instance - .invoke_llm_execution(&name, &model_name, request, next) - .await + .unwrap_or_else(|_| { + instance.log_callback_fallback(&callback_name, surface); + value }) - }), - )?; - } - RegistrationSurface::LlmStreamExecutionIntercept => { - let instance = Arc::new(self.clone_for_callback()); - let callback_name = name.clone(); - ctx.register_llm_stream_execution_intercept( - &name, - priority, - Arc::new(move |model_name, request, next| { - let instance = instance.clone(); - let name = callback_name.clone(); - let model_name = model_name.to_string(); - Box::pin(async move { - instance - .invoke_llm_stream_execution(&name, &model_name, request, next) - .await + }), + ), + RegistrationSurface::ToolConditionalExecutionGuardrail => ctx + .register_tool_conditional_execution_guardrail( + name, + priority, + Arc::new(move |tool_name, value| { + instance.invoke_tool_guardrail(&callback_name, tool_name, value.clone()) + }), + ), + RegistrationSurface::ToolRequestIntercept => ctx.register_tool_request_intercept( + name, + priority, + registration.break_chain, + Arc::new(move |tool_name, value| { + instance.invoke_tool_json(&callback_name, surface, tool_name, value, None) + }), + ), + RegistrationSurface::ToolExecutionIntercept => ctx.register_tool_execution_intercept( + name, + priority, + Arc::new(move |tool_name, value, next| { + let instance = instance.clone(); + let callback_name = callback_name.clone(); + let tool_name = tool_name.to_owned(); + Box::pin(async move { + instance + .invoke_tool_execution(&callback_name, &tool_name, value, next) + .await + }) + }), + ), + _ => Err(PluginError::RegistrationFailed(format!( + "worker plugin '{}' cannot install registration surface {} as a tool callback", + self.plugin_kind, + surface.as_str_name() + ))), + } + } + + fn install_llm_registration( + &self, + ctx: &mut PluginRegistrationContext, + registration: &Registration, + surface: RegistrationSurface, + ) -> crate::plugin::Result<()> { + let name = registration.local_name.as_str(); + let priority = registration.priority; + let instance = Arc::new(self.clone_for_callback()); + let callback_name = name.to_owned(); + match surface { + RegistrationSurface::LlmSanitizeRequestGuardrail => ctx + .register_llm_sanitize_request_guardrail( + name, + priority, + Arc::new(move |request, context| { + instance + .invoke_llm_sanitize_request(&callback_name, request.clone(), context) + .unwrap_or_else(|_| { + instance.log_callback_fallback(&callback_name, surface); + None }) - }), - )?; - } - RegistrationSurface::Unspecified => { - return Err(PluginError::RegistrationFailed(format!( - "worker plugin '{}' returned unspecified registration surface", - self.plugin_kind - ))); - } - } + }), + ), + RegistrationSurface::LlmSanitizeResponseGuardrail => ctx + .register_llm_sanitize_response_guardrail( + name, + priority, + Arc::new(move |value, context| { + instance + .invoke_llm_sanitize_response(&callback_name, value.clone(), context) + .unwrap_or_else(|_| { + instance.log_callback_fallback(&callback_name, surface); + None + }) + }), + ), + RegistrationSurface::LlmConditionalExecutionGuardrail => ctx + .register_llm_conditional_execution_guardrail( + name, + priority, + Arc::new(move |request| { + instance.invoke_llm_guardrail(&callback_name, request.clone()) + }), + ), + RegistrationSurface::LlmRequestIntercept => ctx.register_llm_request_intercept( + name, + priority, + registration.break_chain, + Arc::new(move |model_name, request, annotated| { + instance.invoke_llm_request_intercept( + &callback_name, + model_name, + request, + annotated, + ) + }), + ), + RegistrationSurface::LlmExecutionIntercept => ctx.register_llm_execution_intercept( + name, + priority, + Arc::new(move |model_name, request, next| { + let instance = instance.clone(); + let callback_name = callback_name.clone(); + let model_name = model_name.to_owned(); + Box::pin(async move { + instance + .invoke_llm_execution(&callback_name, &model_name, request, next) + .await + }) + }), + ), + RegistrationSurface::LlmStreamExecutionIntercept => ctx + .register_llm_stream_execution_intercept( + name, + priority, + Arc::new(move |model_name, request, next| { + let instance = instance.clone(); + let callback_name = callback_name.clone(); + let model_name = model_name.to_owned(); + Box::pin(async move { + instance + .invoke_llm_stream_execution( + &callback_name, + &model_name, + request, + next, + ) + .await + }) + }), + ), + _ => Err(PluginError::RegistrationFailed(format!( + "worker plugin '{}' cannot install registration surface {} as an LLM callback", + self.plugin_kind, + surface.as_str_name() + ))), } - Ok(()) } fn clone_for_callback(&self) -> WorkerPluginCallback { @@ -1555,51 +1568,99 @@ impl WorkerPluginCallback { } } - fn invoke_llm_request_json( + fn invoke_llm_sanitize_request( &self, registration_name: &str, - surface: RegistrationSurface, - model_name: &str, request: LlmRequest, - annotated: Option, - continuation_id: Option, - ) -> FlowResult { - let invoke = self.base_request( + context: LlmSanitizeRequestContext, + ) -> FlowResult> { + let mut invoke = self.base_request( registration_name, - surface, - continuation_id, - Some(invoke_request_payload_llm( - model_name, + RegistrationSurface::LlmSanitizeRequestGuardrail, + None, + Some(invoke_request_payload_llm_context( + "", Some(request), - annotated, None, + None, + llm_invocation::SanitizeContext::RequestSanitizeContext( + ProtoLlmSanitizeRequestContext { + codec: Some(codec_identity_to_proto(context.codec())), + codec_capability_id: None, + }, + ), )), ); - let value = json_from_invoke_response(self.invoke_blocking(invoke)?)?; - serde_json::from_value(value).map_err(|err| { - FlowError::Internal(format!("worker returned invalid LLM request: {err}")) - }) + let capability_id = context.resolve_codec().map(|codec| { + let capability_id = self + .host_state + .insert_request_codec(&invoke.invocation_id, codec); + let Some(invoke_request_payload::Payload::Llm(llm)) = invoke.payload.as_mut() else { + unreachable!("LLM sanitizer invocation must have an LLM payload"); + }; + let Some(llm_invocation::SanitizeContext::RequestSanitizeContext(context)) = + llm.sanitize_context.as_mut() + else { + unreachable!("request sanitizer invocation must have a request context"); + }; + context.codec_capability_id = Some(capability_id.clone()); + capability_id + }); + let response = self.invoke_blocking(invoke); + if let Some(capability_id) = capability_id { + self.host_state.remove_codec(&capability_id); + } + optional_json_from_invoke_response(response?)? + .map(serde_json::from_value) + .transpose() + .map_err(|err| { + FlowError::Internal(format!("worker returned invalid LLM request: {err}")) + }) } - fn invoke_llm_response_json( + fn invoke_llm_sanitize_response( &self, registration_name: &str, - surface: RegistrationSurface, - model_name: &str, response: Json, - ) -> FlowResult { - let invoke = self.base_request( + context: LlmSanitizeResponseContext, + ) -> FlowResult> { + let mut invoke = self.base_request( registration_name, - surface, + RegistrationSurface::LlmSanitizeResponseGuardrail, None, - Some(invoke_request_payload_llm( - model_name, + Some(invoke_request_payload_llm_context( + "", None, None, Some(response), + llm_invocation::SanitizeContext::ResponseSanitizeContext( + ProtoLlmSanitizeResponseContext { + codec: Some(codec_identity_to_proto(context.codec())), + codec_capability_id: None, + }, + ), )), ); - json_from_invoke_response(self.invoke_blocking(invoke)?) + let capability_id = context.resolve_codec().map(|codec| { + let capability_id = self + .host_state + .insert_response_codec(&invoke.invocation_id, codec); + let Some(invoke_request_payload::Payload::Llm(llm)) = invoke.payload.as_mut() else { + unreachable!("LLM sanitizer invocation must have an LLM payload"); + }; + let Some(llm_invocation::SanitizeContext::ResponseSanitizeContext(context)) = + llm.sanitize_context.as_mut() + else { + unreachable!("response sanitizer invocation must have a response context"); + }; + context.codec_capability_id = Some(capability_id.clone()); + capability_id + }); + let response = self.invoke_blocking(invoke); + if let Some(capability_id) = capability_id { + self.host_state.remove_codec(&capability_id); + } + optional_json_from_invoke_response(response?) } fn invoke_llm_guardrail( @@ -1973,6 +2034,17 @@ struct WorkerHostRuntimeState { scope_stack_cleanup_complete: Condvar, scope_handles: Mutex>, continuations: Mutex>, + codecs: Mutex>, +} + +struct WorkerCodecCapability { + invocation_id: String, + direction: WorkerCodecDirection, +} + +enum WorkerCodecDirection { + Request(Arc), + Response(Arc), } struct StoredScopeStack { @@ -2024,6 +2096,85 @@ impl WorkerHostRuntimeState { scope_stack_cleanup_complete: Condvar::new(), scope_handles: Mutex::new(HashMap::new()), continuations: Mutex::new(HashMap::new()), + codecs: Mutex::new(HashMap::new()), + } + } + + fn insert_request_codec(&self, invocation_id: &str, codec: Arc) -> String { + self.insert_codec(invocation_id, WorkerCodecDirection::Request(codec)) + } + + fn insert_response_codec( + &self, + invocation_id: &str, + codec: Arc, + ) -> String { + self.insert_codec(invocation_id, WorkerCodecDirection::Response(codec)) + } + + fn insert_codec(&self, invocation_id: &str, direction: WorkerCodecDirection) -> String { + let id = format!("codec-{}", Uuid::now_v7()); + if let Ok(mut codecs) = self.codecs.lock() { + codecs.insert( + id.clone(), + WorkerCodecCapability { + invocation_id: invocation_id.to_owned(), + direction, + }, + ); + } + id + } + + fn remove_codec(&self, id: &str) { + if let Ok(mut codecs) = self.codecs.lock() { + codecs.remove(id); + } + } + + fn request_codec(&self, id: &str, invocation_id: &str) -> Result, Status> { + let codecs = self + .codecs + .lock() + .map_err(|err| Status::internal(format!("codec lock poisoned: {err}")))?; + let capability = codecs + .get(id) + .ok_or_else(|| Status::not_found("codec capability is unavailable"))?; + if capability.invocation_id != invocation_id { + return Err(Status::permission_denied( + "codec capability does not belong to this invocation", + )); + } + match &capability.direction { + WorkerCodecDirection::Request(codec) => Ok(codec.clone()), + WorkerCodecDirection::Response(_) => Err(Status::invalid_argument( + "codec capability is response-only", + )), + } + } + + fn response_codec( + &self, + id: &str, + invocation_id: &str, + ) -> Result, Status> { + let codecs = self + .codecs + .lock() + .map_err(|err| Status::internal(format!("codec lock poisoned: {err}")))?; + let capability = codecs + .get(id) + .ok_or_else(|| Status::not_found("codec capability is unavailable"))?; + if capability.invocation_id != invocation_id { + return Err(Status::permission_denied( + "codec capability does not belong to this invocation", + )); + } + match &capability.direction { + WorkerCodecDirection::Response(codec) => Ok(codec.clone()), + WorkerCodecDirection::Request(_) => { + Err(Status::invalid_argument("codec capability is request-only")) + } } } @@ -2425,6 +2576,79 @@ impl RelayHostRuntime for WorkerHostRuntimeService { }); Ok(Response::new(Box::pin(mapped))) } + + async fn decode_llm_codec_request( + &self, + request: Request, + ) -> Result, Status> { + let request = request.into_inner(); + self.state + .authorize(&request.activation_id, &request.auth_token)?; + let codec = self + .state + .request_codec(&request.codec_capability_id, &request.invocation_id)?; + let request = required_envelope(request.request, "codec request") + .and_then(|value| { + decode_json_envelope::(&value) + .map_err(|err| FlowError::Internal(err.to_string())) + }) + .map_err(status_from_flow)?; + Ok(Response::new(typed_json_result( + ANNOTATED_LLM_REQUEST_SCHEMA, + codec.decode(&request), + ))) + } + + async fn encode_llm_codec_request( + &self, + request: Request, + ) -> Result, Status> { + let request = request.into_inner(); + self.state + .authorize(&request.activation_id, &request.auth_token)?; + let codec = self + .state + .request_codec(&request.codec_capability_id, &request.invocation_id)?; + let annotated = required_envelope(request.annotated_request, "annotated codec request") + .and_then(|value| { + decode_json_envelope::(&value) + .map_err(|err| FlowError::Internal(err.to_string())) + }) + .map_err(status_from_flow)?; + let original = required_envelope(request.original_request, "original codec request") + .and_then(|value| { + decode_json_envelope::(&value) + .map_err(|err| FlowError::Internal(err.to_string())) + }) + .map_err(status_from_flow)?; + Ok(Response::new(typed_json_result( + LLM_REQUEST_SCHEMA, + codec.encode(&annotated, &original), + ))) + } + + async fn decode_llm_codec_response( + &self, + request: Request, + ) -> Result, Status> { + let request = request.into_inner(); + self.state + .authorize(&request.activation_id, &request.auth_token)?; + let codec = self + .state + .response_codec(&request.codec_capability_id, &request.invocation_id)?; + let response = required_envelope(request.response, "codec response") + .and_then(|value| { + decode_json_envelope::(&value) + .map_err(|err| FlowError::Internal(err.to_string())) + }) + .map_err(status_from_flow)?; + Ok(Response::new(json_result( + codec.decode_response(&response).and_then(|value| { + serde_json::to_value(value).map_err(|err| FlowError::Internal(err.to_string())) + }), + ))) + } } impl WorkerHostRuntimeService { @@ -2451,6 +2675,10 @@ mod invoke_request_payload { pub(crate) use nemo_relay_worker_proto::v1::invoke_request::Payload; } +mod llm_invocation { + pub(crate) use nemo_relay_worker_proto::v1::llm_invocation::SanitizeContext; +} + mod invoke_response_result { pub(crate) use nemo_relay_worker_proto::v1::invoke_response::Result; } @@ -2487,9 +2715,44 @@ fn invoke_request_payload_llm( response: response .as_ref() .map(|response| json_envelope_infallible(JSON_SCHEMA, response)), + sanitize_context: None, + }) +} + +fn invoke_request_payload_llm_context( + model_name: &str, + request: Option, + annotated_request: Option, + response: Option, + context: impl Into, +) -> invoke_request_payload::Payload { + invoke_request_payload::Payload::Llm(LlmInvocation { + model_name: model_name.into(), + request: request + .as_ref() + .map(|request| json_envelope_infallible(LLM_REQUEST_SCHEMA, request)), + annotated_request: annotated_request + .as_ref() + .map(|request| json_envelope_infallible(ANNOTATED_LLM_REQUEST_SCHEMA, request)), + response: response + .as_ref() + .map(|response| json_envelope_infallible(JSON_SCHEMA, response)), + sanitize_context: Some(context.into()), }) } +fn codec_identity_to_proto(identity: &LlmCodecIdentity) -> ProtoLlmCodecIdentity { + let (kind, id) = match identity { + LlmCodecIdentity::None => (LlmCodecKind::Unspecified as i32, None), + LlmCodecIdentity::BuiltIn(codec) => { + (LlmCodecKind::Builtin as i32, Some(codec.id().to_owned())) + } + LlmCodecIdentity::Runtime(id) => (LlmCodecKind::Runtime as i32, Some(id.clone())), + LlmCodecIdentity::Opaque => (LlmCodecKind::Opaque as i32, None), + }; + ProtoLlmCodecIdentity { kind, id } +} + fn json_envelope_infallible(schema: &str, value: &T) -> JsonEnvelope { json_envelope(schema, value).expect("Relay DTO JSON serialization should be infallible") } @@ -2512,6 +2775,27 @@ fn json_from_invoke_response(response: InvokeResponse) -> FlowResult { } } +fn optional_json_from_invoke_response(response: InvokeResponse) -> FlowResult> { + match response.result { + Some(invoke_response_result::Result::Empty(_)) => Ok(None), + Some(invoke_response_result::Result::Json(result)) => { + if let Some(error) = result.error { + return Err(worker_error_to_flow(error)); + } + let envelope = required_envelope(result.value, "worker JSON result")?; + decode_json_envelope::(&envelope) + .map(Some) + .map_err(|err| { + FlowError::Internal(format!("worker returned invalid JSON result: {err}")) + }) + } + Some(invoke_response_result::Result::Error(error)) => Err(worker_error_to_flow(error)), + _ => Err(FlowError::Internal( + "worker returned unexpected LLM sanitizer result".into(), + )), + } +} + fn guardrail_from_invoke_response(response: InvokeResponse) -> FlowResult> { match response.result { Some(invoke_response_result::Result::Guardrail(GuardrailResult { block_reason })) => { @@ -2572,6 +2856,19 @@ fn json_result(result: FlowResult) -> JsonResult { } } +fn typed_json_result(schema: &str, result: FlowResult) -> JsonResult { + match result { + Ok(value) => JsonResult { + value: Some(json_envelope_infallible(schema, &value)), + error: None, + }, + Err(err) => JsonResult { + value: None, + error: Some(flow_error_to_worker(err)), + }, + } +} + fn flow_error_to_worker(err: FlowError) -> WorkerError { WorkerError { code: "host.runtime_error".into(), diff --git a/crates/core/src/stream.rs b/crates/core/src/stream.rs index 1124b3d0d..3fee8dd3e 100644 --- a/crates/core/src/stream.rs +++ b/crates/core/src/stream.rs @@ -36,6 +36,7 @@ use crate::api::event::{BaseEvent, MarkEvent}; use crate::api::llm::LlmHandle; use crate::api::llm::emit_optimization_marks; use crate::api::optimization::finalize_optimization_summary; +use crate::api::runtime::LlmSanitizeResponseContext; use crate::api::runtime::NemoRelayContextState; use crate::api::runtime::global_context; use crate::api::runtime::{ @@ -71,6 +72,7 @@ pub struct LlmStreamWrapper { collector: Box Result<()> + Send>, finalizer: Option Json + Send>>, response_codec: Option>, + sanitize_context: LlmSanitizeResponseContext, metadata: Option, subscribers: Vec, chunk_index: u64, @@ -140,6 +142,8 @@ impl LlmStreamWrapper { subscribers: Vec, ) -> Self { let scope_stack = handle.captured_scope_stack().clone(); + let sanitize_context = + LlmSanitizeResponseContext::for_response_codec(response_codec.clone()); Self { inner, handle, @@ -147,6 +151,7 @@ impl LlmStreamWrapper { collector, finalizer: Some(finalizer), response_codec, + sanitize_context, metadata, subscribers, chunk_index: 0, @@ -204,6 +209,12 @@ impl LlmStreamWrapper { Some(finalizer) => finalizer(), None => Json::Null, }; + let response_was_null_without_fallback = aggregated.is_null() && self.handle.data.is_none(); + let response = if aggregated.is_null() { + self.handle.data.clone().unwrap_or(aggregated) + } else { + aggregated + }; let snapshot = { let ss_guard = self.scope_stack.read().expect("scope stack lock poisoned"); @@ -222,19 +233,27 @@ impl LlmStreamWrapper { let Some(entries) = snapshot else { return; }; - let sanitized = - NemoRelayContextState::llm_sanitize_response_snapshot_chain(aggregated, &entries); - let data = if sanitized.is_null() { - self.handle.data.clone() - } else { - Some(sanitized) + let sanitized = NemoRelayContextState::llm_sanitize_response_snapshot_chain( + response, + self.sanitize_context.clone(), + &entries, + ); + let data = match sanitized { + Some(response) if response_was_null_without_fallback && response.is_null() => None, + response => response, }; - let mut annotated_response: Option = - self.response_codec.as_ref().and_then(|codec| { - let mut decoded = codec.decode_response(data.as_ref()?).ok()?; - attach_estimated_cost_for_provider(&mut decoded, Some(&self.handle.name)); - Some(decoded) - }); + let annotation_omitted = data.as_ref().is_none_or(Json::is_null); + let mut annotated_response: Option = (!annotation_omitted) + .then(|| { + data.as_ref().and_then(|response| { + self.response_codec.as_ref().and_then(|codec| { + let mut decoded = codec.decode_response(response).ok()?; + attach_estimated_cost_for_provider(&mut decoded, Some(&self.handle.name)); + Some(decoded) + }) + }) + }) + .flatten(); let interruption = (interrupted && !has_authoritative_final_usage(annotated_response.as_ref())) .then_some("stream_interrupted"); @@ -249,7 +268,8 @@ impl LlmStreamWrapper { self.handle.model_name.as_deref(), &pricing, ); - if annotated_response.is_none() + if !annotation_omitted + && annotated_response.is_none() && let Some(summary) = summary { annotated_response = Some(AnnotatedLlmResponse { diff --git a/crates/core/tests/fixtures/native_plugin/src/lib.rs b/crates/core/tests/fixtures/native_plugin/src/lib.rs index 08b3191d8..350cf1414 100644 --- a/crates/core/tests/fixtures/native_plugin/src/lib.rs +++ b/crates/core/tests/fixtures/native_plugin/src/lib.rs @@ -130,12 +130,12 @@ impl NativePlugin for FixtureNativePlugin { ctx.register_llm_sanitize_request_guardrail( "fixture_llm_sanitize_request", 0, - |request| mark_llm_request(request, "native_plugin_llm_sanitize_request"), + |request, _context| Some(mark_llm_request(request, "native_plugin_llm_sanitize_request")), )?; ctx.register_llm_sanitize_response_guardrail( "fixture_llm_sanitize_response", 0, - |response| mark_json(response, "native_plugin_llm_sanitize_response"), + |response, _context| Some(mark_json(response, "native_plugin_llm_sanitize_response")), )?; ctx.register_llm_conditional_execution_guardrail( "fixture_llm_conditional", diff --git a/crates/core/tests/fixtures/worker_plugin/src/main.rs b/crates/core/tests/fixtures/worker_plugin/src/main.rs index c0320f2d5..a9dd41a67 100644 --- a/crates/core/tests/fixtures/worker_plugin/src/main.rs +++ b/crates/core/tests/fixtures/worker_plugin/src/main.rs @@ -2,11 +2,11 @@ // SPDX-License-Identifier: Apache-2.0 use nemo_relay_worker::{ - JsonStream, LlmNext, LlmStreamNext, PluginContext, ScopeType, ToolNext, WorkerPlugin, - ToolExecutionInterceptOutcome, WorkerSdkError, serve_plugin, + ConfigDiagnostic, DiagnosticLevel, EventSanitizeFields, Json, LlmRequest, PendingMarkSpec, }; use nemo_relay_worker::{ - ConfigDiagnostic, DiagnosticLevel, EventSanitizeFields, Json, LlmRequest, PendingMarkSpec, + JsonStream, LlmNext, LlmStreamNext, PluginContext, ScopeType, ToolExecutionInterceptOutcome, + ToolNext, WorkerPlugin, WorkerSdkError, serve_plugin, }; use serde_json::json; @@ -61,22 +61,26 @@ impl WorkerPlugin for FixtureWorkerPlugin { let runtime = ctx .runtime() .ok_or_else(|| WorkerSdkError::Callback("runtime handle missing".into()))?; - ctx.register_mark_sanitize_guardrail("fixture_mark_sanitize", 0, |_, fields| { - mark_event_fields(fields, "worker_plugin_mark") - }); - ctx.register_mark_sanitize_guardrail("fixture_mark_sanitize_data", 1, |_, mut fields| { - fields.data = Some(json!({"worker_plugin_mark_data": true})); - fields + ctx.register_mark_sanitize_guardrail("fixture_mark_sanitize", 0, |_, fields| async move { + Ok(mark_event_fields(fields, "worker_plugin_mark")) }); + ctx.register_mark_sanitize_guardrail( + "fixture_mark_sanitize_data", + 1, + |_, mut fields| async move { + fields.data = Some(json!({"worker_plugin_mark_data": true})); + Ok(fields) + }, + ); ctx.register_scope_sanitize_start_guardrail( "fixture_scope_start_sanitize", 0, - |_, fields| mark_event_fields(fields, "worker_plugin_scope_start"), + |_, fields| async move { Ok(mark_event_fields(fields, "worker_plugin_scope_start")) }, ); ctx.register_scope_sanitize_end_guardrail( "fixture_scope_end_sanitize", 0, - |_, fields| mark_event_fields(fields, "worker_plugin_scope_end"), + |_, fields| async move { Ok(mark_event_fields(fields, "worker_plugin_scope_end")) }, ); register_fixture_subscriber(ctx, runtime.clone()); register_fixture_tool_hooks( @@ -135,12 +139,14 @@ fn register_fixture_tool_hooks( ctx.register_tool_sanitize_request_guardrail( "fixture_tool_sanitize_request", 0, - |_name, args| mark_json(args, "worker_plugin_tool_sanitize_request"), + |_name, args| async move { Ok(mark_json(args, "worker_plugin_tool_sanitize_request")) }, ); ctx.register_tool_sanitize_response_guardrail( "fixture_tool_sanitize_response", 0, - |_name, result| mark_json(result, "worker_plugin_tool_sanitize_response"), + |_name, result| async move { + Ok(mark_json(result, "worker_plugin_tool_sanitize_response")) + }, ); ctx.register_tool_conditional_execution_guardrail( "fixture_tool_conditional", @@ -175,17 +181,15 @@ fn register_fixture_tool_hooks( let result = next .call(mark_json(args, "worker_plugin_tool_execution_request")) .await?; - Ok( - ToolExecutionInterceptOutcome::new(mark_json( - result, - "worker_plugin_tool_execution", - )) - .with_pending_mark( - PendingMarkSpec::builder() - .name("fixture.worker.tool_execution.mark") - .build(), - ), - ) + Ok(ToolExecutionInterceptOutcome::new(mark_json( + result, + "worker_plugin_tool_execution", + )) + .with_pending_mark( + PendingMarkSpec::builder() + .name("fixture.worker.tool_execution.mark") + .build(), + )) }, ); } @@ -198,18 +202,26 @@ fn register_fixture_llm_hooks( ctx.register_llm_sanitize_request_guardrail( "fixture_llm_sanitize_request", 0, - |request| mark_llm_request(request, "worker_plugin_llm_sanitize_request"), + |request, _context| async move { + Ok(Some(mark_llm_request( + request, + "worker_plugin_llm_sanitize_request", + ))) + }, ); ctx.register_llm_sanitize_response_guardrail( "fixture_llm_sanitize_response", 0, - |response| mark_json(response, "worker_plugin_llm_sanitize_response"), - ); - ctx.register_llm_conditional_execution_guardrail( - "fixture_llm_conditional", - 0, - |_request| Ok(None), + |response, _context| async move { + Ok(Some(mark_json( + response, + "worker_plugin_llm_sanitize_response", + ))) + }, ); + ctx.register_llm_conditional_execution_guardrail("fixture_llm_conditional", 0, |_request| { + Ok(None) + }); ctx.register_llm_request_intercept( "fixture_llm_request_intercept", 0, @@ -232,14 +244,16 @@ fn register_fixture_llm_hooks( None, ), }; - Ok(nemo_relay_worker::LlmRequestInterceptOutcome::new(request, annotated) - .with_pending_mark( - PendingMarkSpec::builder() - .name("fixture.worker.llm_request.mark") - .data(json!({ "source": "worker_request_intercept" })) - .metadata(json!({ "fixture": true })) - .build(), - )) + Ok( + nemo_relay_worker::LlmRequestInterceptOutcome::new(request, annotated) + .with_pending_mark( + PendingMarkSpec::builder() + .name("fixture.worker.llm_request.mark") + .data(json!({ "source": "worker_request_intercept" })) + .metadata(json!({ "fixture": true })) + .build(), + ), + ) }, ); ctx.register_llm_execution_intercept( @@ -278,7 +292,9 @@ fn register_fixture_llm_hooks( ); } -async fn emit_runtime_events(runtime: nemo_relay_worker::PluginRuntime) -> nemo_relay_worker::Result<()> { +async fn emit_runtime_events( + runtime: nemo_relay_worker::PluginRuntime, +) -> nemo_relay_worker::Result<()> { runtime .emit_mark("fixture.worker.mark", Some(json!("current")), None) .await?; @@ -311,7 +327,11 @@ async fn emit_runtime_events(runtime: nemo_relay_worker::PluginRuntime) -> nemo_ runtime .with_scope_stack(&isolated, || async move { isolated_runtime - .emit_mark("fixture.worker.isolated.mark", Some(json!("isolated")), None) + .emit_mark( + "fixture.worker.isolated.mark", + Some(json!("isolated")), + None, + ) .await }) .await?; diff --git a/crates/core/tests/integration/api_surface_tests.rs b/crates/core/tests/integration/api_surface_tests.rs index 6b3d83bd9..7a6045ac0 100644 --- a/crates/core/tests/integration/api_surface_tests.rs +++ b/crates/core/tests/integration/api_surface_tests.rs @@ -903,14 +903,18 @@ fn test_global_registry_and_subscriber_wrappers_cover_success_and_duplicates() { .unwrap(); assert!(deregister_tool_execution_intercept("tool-execution").unwrap()); - register_llm_sanitize_request_guardrail("llm-sanitize-request", 1, Arc::new(|request| request)) - .unwrap(); + register_llm_sanitize_request_guardrail( + "llm-sanitize-request", + 1, + Arc::new(|request, _context| Some(request)), + ) + .unwrap(); assert!(deregister_llm_sanitize_request_guardrail("llm-sanitize-request").unwrap()); register_llm_sanitize_response_guardrail( "llm-sanitize-response", 1, - Arc::new(|response| response), + Arc::new(|response, _context| Some(response)), ) .unwrap(); assert!(deregister_llm_sanitize_response_guardrail("llm-sanitize-response").unwrap()); @@ -1134,7 +1138,7 @@ fn test_scope_registry_and_subscriber_wrappers_cover_success_duplicates_and_miss &scope.uuid, "llm-sanitize-request", 1, - Arc::new(|request| request), + Arc::new(|request, _context| Some(request)), ) .unwrap(); assert!( @@ -1146,7 +1150,7 @@ fn test_scope_registry_and_subscriber_wrappers_cover_success_duplicates_and_miss &scope.uuid, "llm-sanitize-response", 1, - Arc::new(|response| response), + Arc::new(|response, _context| Some(response)), ) .unwrap(); assert!( @@ -1486,21 +1490,21 @@ async fn test_llm_api_emits_sanitized_events_and_covers_error_paths() { register_llm_sanitize_request_guardrail( "llm-sanitize-request", 1, - Arc::new(|mut request| { + Arc::new(|mut request, _context| { request.headers.insert("x-sanitized".into(), json!(true)); - request + Some(request) }), ) .unwrap(); register_llm_sanitize_response_guardrail( "llm-sanitize-response", 1, - Arc::new(|mut response| { + Arc::new(|mut response, _context| { response .as_object_mut() .unwrap() .insert("sanitized_response".into(), json!(true)); - response + Some(response) }), ) .unwrap(); diff --git a/crates/core/tests/integration/middleware_tests.rs b/crates/core/tests/integration/middleware_tests.rs index 5acc08b63..41093cdcb 100644 --- a/crates/core/tests/integration/middleware_tests.rs +++ b/crates/core/tests/integration/middleware_tests.rs @@ -2638,10 +2638,10 @@ async fn test_llm_middleware_callbacks_run_without_registry_or_scope_locks() { register_llm_sanitize_request_guardrail( "lock_global_llm_sanitize_request", 1, - Arc::new(move |request| { + Arc::new(move |request, _context| { record_middleware_callback(&tracked, "llm_sanitize_request_global"); assert_middleware_callback_locks_are_free(); - request + Some(request) }), ) .unwrap(); @@ -2650,10 +2650,10 @@ async fn test_llm_middleware_callbacks_run_without_registry_or_scope_locks() { &scope.uuid, "lock_scope_llm_sanitize_request", 2, - Arc::new(move |request| { + Arc::new(move |request, _context| { record_middleware_callback(&tracked, "llm_sanitize_request_scope"); assert_middleware_callback_locks_are_free(); - request + Some(request) }), ) .unwrap(); @@ -2707,10 +2707,10 @@ async fn test_llm_middleware_callbacks_run_without_registry_or_scope_locks() { register_llm_sanitize_response_guardrail( "lock_global_llm_sanitize_response", 1, - Arc::new(move |response| { + Arc::new(move |response, _context| { record_middleware_callback(&tracked, "llm_sanitize_response_global"); assert_middleware_callback_locks_are_free(); - response + Some(response) }), ) .unwrap(); @@ -2719,10 +2719,10 @@ async fn test_llm_middleware_callbacks_run_without_registry_or_scope_locks() { &scope.uuid, "lock_scope_llm_sanitize_response", 2, - Arc::new(move |response| { + Arc::new(move |response, _context| { record_middleware_callback(&tracked, "llm_sanitize_response_scope"); assert_middleware_callback_locks_are_free(); - response + Some(response) }), ) .unwrap(); diff --git a/crates/core/tests/integration/native_plugin_tests.rs b/crates/core/tests/integration/native_plugin_tests.rs index ef2604f58..ee8a860be 100644 --- a/crates/core/tests/integration/native_plugin_tests.rs +++ b/crates/core/tests/integration/native_plugin_tests.rs @@ -640,12 +640,17 @@ async fn sdk_cdylib_registers_tool_request_intercept() { flush_subscribers().expect("llm response annotation event should flush"); let llm_events = events.lock().unwrap().clone(); let llm_end = find_event(&llm_events, "native-fixture-llm", Some(ScopeCategory::End)); - let annotated = llm_end - .annotated_response() - .expect("native plugin should preserve response annotation"); - assert_eq!(annotated.id.as_deref(), Some("annotation-before-plugin")); - assert_eq!(annotated.extra["preexisting_annotation"], json!("kept")); - assert!(annotated.extra.get("native_plugin_annotation").is_none()); + assert_eq!( + llm_end.output().unwrap()["native_plugin_llm_sanitize_response"], + true + ); + assert!( + llm_end.annotated_response().is_none(), + "a changed response without an active codec must discard the stale caller annotation" + ); + let serialized = serde_json::to_string(llm_end).expect("LLM end event should serialize"); + assert!(!serialized.contains("annotation-before-plugin")); + assert!(!serialized.contains("preexisting_annotation")); drop(cleanup); activation.clear(); diff --git a/crates/core/tests/integration/pipeline_tests.rs b/crates/core/tests/integration/pipeline_tests.rs index 918407663..1bc192b4b 100644 --- a/crates/core/tests/integration/pipeline_tests.rs +++ b/crates/core/tests/integration/pipeline_tests.rs @@ -1448,7 +1448,7 @@ async fn test_response_codec_annotation_uses_sanitized_managed_response() { register_llm_sanitize_response_guardrail( "sanitize_resp_codec_annotation", 1, - Arc::new(|_response| make_openai_chat_response("Sanitized")), + Arc::new(|_response, _context| Some(make_openai_chat_response("Sanitized"))), ) .unwrap(); @@ -1648,9 +1648,11 @@ async fn test_request_codec_annotation_uses_sanitized_start_payload() { register_llm_sanitize_request_guardrail( "sanitize_req_codec_annotation", 1, - Arc::new(|request| LlmRequest { - headers: request.headers, - content: make_openai_chat_request("Sanitized").content, + Arc::new(|request, _context| { + Some(LlmRequest { + headers: request.headers, + content: make_openai_chat_request("Sanitized").content, + }) }), ) .unwrap(); @@ -1914,7 +1916,7 @@ async fn test_stream_response_codec_annotation_uses_sanitized_aggregated_respons register_llm_sanitize_response_guardrail( "stream_sanitize_resp_codec_annotation", 1, - Arc::new(|_response| make_openai_chat_response("Sanitized")), + Arc::new(|_response, _context| Some(make_openai_chat_response("Sanitized"))), ) .unwrap(); diff --git a/crates/core/tests/unit/context_tests.rs b/crates/core/tests/unit/context_tests.rs index 07e4ea6b6..49e57b2cb 100644 --- a/crates/core/tests/unit/context_tests.rs +++ b/crates/core/tests/unit/context_tests.rs @@ -320,8 +320,12 @@ fn context_state_supports_extensions_events_and_builders() { content: json!({"messages": []}), }; let entries = state.llm_sanitize_request_entries(&[]); - let sanitized = - NemoRelayContextState::llm_sanitize_request_snapshot_chain(request.clone(), &entries); + let sanitized = NemoRelayContextState::llm_sanitize_request_snapshot_chain( + request.clone(), + crate::api::runtime::LlmSanitizeRequestContext::default(), + &entries, + ) + .expect("an empty sanitizer chain must retain the request"); assert!(sanitized.headers.is_empty()); let events = Arc::new(Mutex::new(Vec::::new())); diff --git a/crates/core/tests/unit/dynamic_worker_tests.rs b/crates/core/tests/unit/dynamic_worker_tests.rs index 1e28358ad..7754238dc 100644 --- a/crates/core/tests/unit/dynamic_worker_tests.rs +++ b/crates/core/tests/unit/dynamic_worker_tests.rs @@ -4,7 +4,12 @@ use std::sync::{Arc, Mutex}; use crate::api::event::{BaseEvent, MarkEvent}; -use crate::api::runtime::NemoRelayContextState; +use crate::api::runtime::{ + BuiltinLlmCodec, LlmCodecIdentity, LlmSanitizeRequestContext, LlmSanitizeResponseContext, + NemoRelayContextState, +}; +use crate::codec::openai_chat::OpenAIChatCodec; +use crate::codec::traits::{LlmCodec, LlmResponseCodec}; use nemo_relay_worker_proto::json_envelope; use nemo_relay_worker_proto::v1::invoke_response::Result as InvokeResult; use nemo_relay_worker_proto::v1::plugin_worker_server::{PluginWorker, PluginWorkerServer}; @@ -12,10 +17,10 @@ use nemo_relay_worker_proto::v1::stream_chunk::Item as StreamItem; use nemo_relay_worker_proto::v1::{ CancelInvocationRequest, CreateScopeStackRequest, DropScopeStackRequest, EmitMarkRequest, EmptyResult, GuardrailResult, HandshakeRequest, HandshakeResponse, HealthRequest, - HealthResponse, JsonEnvelope, JsonResult, LlmNextRequest, LlmRequestInterceptResult, - LlmStreamNextRequest, PopScopeRequest, PushScopeRequest, Registration, ScopeContext, - ScopeType as ProtoScopeType, ShutdownRequest, StreamChunk, ToolNextRequest, ValidateRequest, - ValidateResponse, WorkerAck, + HealthResponse, JsonEnvelope, JsonResult, LlmCodecDecodeRequest, LlmCodecDecodeResponse, + LlmCodecEncodeRequest, LlmNextRequest, LlmRequestInterceptResult, LlmStreamNextRequest, + PopScopeRequest, PushScopeRequest, Registration, ScopeContext, ScopeType as ProtoScopeType, + ShutdownRequest, StreamChunk, ToolNextRequest, ValidateRequest, ValidateResponse, WorkerAck, }; use serde_json::json; use tokio_stream::StreamExt; @@ -464,13 +469,10 @@ async fn callback_helpers_cover_worker_response_edges() { ); let error = callback - .invoke_llm_request_json( + .invoke_llm_sanitize_request( "llm_json_invalid", - RegistrationSurface::LlmSanitizeRequestGuardrail, - "model", valid_llm_request(), - None, - None, + LlmSanitizeRequestContext::default(), ) .expect_err("invalid LLM JSON result should fail"); assert!(error.to_string().contains("invalid type")); @@ -537,6 +539,287 @@ async fn callback_helpers_cover_worker_response_edges() { ); } +#[tokio::test(flavor = "multi_thread")] +async fn llm_worker_sanitizers_forward_codec_context_and_omission() { + enable_operational_logs(); + let seen = Arc::new(Mutex::new(Vec::new())); + let (callback, _shutdown) = fake_callback_service({ + let seen = seen.clone(); + move |request| { + let Some(nemo_relay_worker_proto::v1::invoke_request::Payload::Llm(invocation)) = + request.payload + else { + panic!("LLM sanitizer must receive an LLM invocation"); + }; + let codec = match invocation.sanitize_context.as_ref() { + Some(nemo_relay_worker_proto::v1::llm_invocation::SanitizeContext::RequestSanitizeContext(context)) => context.codec.as_ref(), + Some(nemo_relay_worker_proto::v1::llm_invocation::SanitizeContext::ResponseSanitizeContext(context)) => context.codec.as_ref(), + None => None, + }; + seen.lock().unwrap().push(( + request.registration_name, + codec + .map(|codec| codec.kind) + .unwrap_or(LlmCodecKind::Unspecified as i32), + codec.and_then(|codec| codec.id.clone()), + invocation.request.is_some(), + invocation.response.is_some(), + )); + InvokeResponse { + result: Some(InvokeResult::Empty(EmptyResult {})), + } + } + }) + .await; + + let identities = [ + LlmCodecIdentity::None, + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiResponses), + LlmCodecIdentity::Runtime("com.example.chat.v1".into()), + LlmCodecIdentity::Opaque, + ]; + for identity in identities { + assert!( + callback + .invoke_llm_sanitize_request( + "request", + valid_llm_request(), + LlmSanitizeRequestContext::with_identity(identity.clone()), + ) + .expect("empty worker result must represent request omission") + .is_none() + ); + assert!( + callback + .invoke_llm_sanitize_response( + "response", + json!({"secret": "value"}), + LlmSanitizeResponseContext::with_identity(identity), + ) + .expect("empty worker result must represent response omission") + .is_none() + ); + } + + let seen = seen.lock().unwrap().clone(); + assert_eq!( + seen, + [ + ( + "request".into(), + LlmCodecKind::Unspecified as i32, + None, + true, + false, + ), + ( + "response".into(), + LlmCodecKind::Unspecified as i32, + None, + false, + true, + ), + ( + "request".into(), + LlmCodecKind::Builtin as i32, + Some("openai_responses".into()), + true, + false, + ), + ( + "response".into(), + LlmCodecKind::Builtin as i32, + Some("openai_responses".into()), + false, + true, + ), + ( + "request".into(), + LlmCodecKind::Runtime as i32, + Some("com.example.chat.v1".into()), + true, + false, + ), + ( + "response".into(), + LlmCodecKind::Runtime as i32, + Some("com.example.chat.v1".into()), + false, + true, + ), + ( + "request".into(), + LlmCodecKind::Opaque as i32, + None, + true, + false, + ), + ( + "response".into(), + LlmCodecKind::Opaque as i32, + None, + false, + true, + ), + ] + ); +} + +#[tokio::test(flavor = "multi_thread")] +async fn llm_worker_codec_capabilities_are_active_only_during_sanitizer_invocation() { + enable_operational_logs(); + type SeenCapability = (String, String, String); + + let host_state = Arc::new(Mutex::new(None::>)); + let seen = Arc::new(Mutex::new(Vec::::new())); + let (callback, _shutdown) = fake_callback_service({ + let host_state = host_state.clone(); + let seen = seen.clone(); + move |request| { + let invocation_id = request.invocation_id.clone(); + let registration_name = request.registration_name.clone(); + let Some(nemo_relay_worker_proto::v1::invoke_request::Payload::Llm(invocation)) = + request.payload + else { + panic!("LLM sanitizer must receive an LLM invocation"); + }; + let state = host_state + .lock() + .unwrap() + .clone() + .expect("host state must be installed before invoking the worker"); + let capability_id = match invocation + .sanitize_context + .as_ref() + .expect("codec context must be forwarded") + { + nemo_relay_worker_proto::v1::llm_invocation::SanitizeContext::RequestSanitizeContext( + context, + ) => { + let id = context + .codec_capability_id + .as_deref() + .expect("active request codec must receive a capability"); + let codec = state + .request_codec(id, &invocation_id) + .expect("request capability must resolve during invocation"); + let request: LlmRequest = decode_json_envelope( + invocation.request.as_ref().expect("request payload"), + ) + .expect("request payload must decode"); + let annotated = codec + .decode(&request) + .expect("resolved request codec must be usable"); + assert_eq!(annotated.model.as_deref(), Some("gpt-test")); + id.to_owned() + } + nemo_relay_worker_proto::v1::llm_invocation::SanitizeContext::ResponseSanitizeContext( + context, + ) => { + let id = context + .codec_capability_id + .as_deref() + .expect("active response codec must receive a capability"); + let codec = state + .response_codec(id, &invocation_id) + .expect("response capability must resolve during invocation"); + let response: Json = decode_json_envelope( + invocation.response.as_ref().expect("response payload"), + ) + .expect("response payload must decode"); + let annotated = codec + .decode_response(&response) + .expect("resolved response codec must be usable"); + assert_eq!(annotated.model.as_deref(), Some("gpt-test")); + id.to_owned() + } + }; + seen.lock().unwrap().push(( + registration_name.clone(), + capability_id, + invocation_id, + )); + if registration_name == "response-error" { + InvokeResponse { + result: Some(InvokeResult::Error( + nemo_relay_worker_proto::v1::WorkerError { + code: "worker.failed".into(), + message: "boom".into(), + retryable: false, + }, + )), + } + } else { + InvokeResponse { + result: Some(InvokeResult::Empty(EmptyResult {})), + } + } + } + }) + .await; + *host_state.lock().unwrap() = Some(callback.host_state.clone()); + + let codec = Arc::new(OpenAIChatCodec); + let request = LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "gpt-test", + "messages": [{"role": "user", "content": "secret"}] + }), + }; + assert!( + callback + .invoke_llm_sanitize_request( + "request-success", + request, + LlmSanitizeRequestContext::for_request_codec(Some(codec.clone())), + ) + .expect("request sanitizer must succeed") + .is_none() + ); + + let response = json!({ + "id": "chatcmpl-test", + "model": "gpt-test", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "secret"}, + "finish_reason": "stop" + }] + }); + let error = callback + .invoke_llm_sanitize_response( + "response-error", + response, + LlmSanitizeResponseContext::for_response_codec(Some(codec)), + ) + .expect_err("worker sanitizer error must surface"); + assert!(error.to_string().contains("worker.failed: boom")); + + let seen = seen.lock().unwrap().clone(); + assert_eq!(seen.len(), 2); + let (_, request_capability, request_invocation) = &seen[0]; + assert_eq!( + callback + .host_state + .request_codec(request_capability, request_invocation) + .err() + .expect("successful invocation must expire its request capability") + .code(), + tonic::Code::NotFound + ); + let (_, response_capability, response_invocation) = &seen[1]; + assert_eq!( + callback + .host_state + .response_codec(response_capability, response_invocation) + .err() + .expect("failed invocation must expire its response capability") + .code(), + tonic::Code::NotFound + ); +} + #[tokio::test(flavor = "multi_thread")] async fn callback_stream_transport_error_surfaces_to_host_stream() { enable_operational_logs(); @@ -1164,20 +1447,24 @@ async fn installed_callbacks_apply_surface_specific_fallbacks() { tool_response ); let entries = state.llm_sanitize_request_entries(&[]); - assert_eq!( + assert!( NemoRelayContextState::llm_sanitize_request_snapshot_chain( llm_request.clone(), + crate::api::runtime::LlmSanitizeRequestContext::default(), &entries, - ), - llm_request + ) + .is_none(), + "a worker request sanitizer failure must omit the observability payload" ); let entries = state.llm_sanitize_response_entries(&[]); - assert_eq!( + assert!( NemoRelayContextState::llm_sanitize_response_snapshot_chain( llm_response.clone(), + crate::api::runtime::LlmSanitizeResponseContext::default(), &entries, - ), - llm_response + ) + .is_none(), + "a worker response sanitizer failure must omit the observability payload" ); } crate::api::subscriber::flush_subscribers().expect("subscriber callback should flush"); @@ -1346,6 +1633,152 @@ async fn host_runtime_service_covers_auth_scope_and_ack_errors() { ); } +#[tokio::test] +async fn host_runtime_codec_capabilities_are_directional_authorized_and_ephemeral() { + let state = Arc::new(WorkerHostRuntimeState::new( + ACTIVATION_ID.into(), + AUTH_TOKEN.into(), + )); + let service = WorkerHostRuntimeService { + state: state.clone(), + }; + let codec = Arc::new(OpenAIChatCodec); + let request_codec: Arc = codec.clone(); + let response_codec: Arc = codec.clone(); + let invocation_id = "sanitize-invocation"; + let request_capability = state.insert_request_codec(invocation_id, request_codec); + let response_capability = state.insert_response_codec(invocation_id, response_codec); + let request = LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "gpt-test", + "messages": [{"role": "user", "content": "secret"}], + "preserve": true + }), + }; + + let unauthorized = service + .decode_llm_codec_request(Request::new(LlmCodecDecodeRequest { + activation_id: ACTIVATION_ID.into(), + auth_token: "wrong".into(), + codec_capability_id: request_capability.clone(), + invocation_id: invocation_id.into(), + request: Some(json_envelope("nemo.relay.LlmRequest@1", &request).unwrap()), + })) + .await + .expect_err("wrong activation credentials must be rejected"); + assert_eq!(unauthorized.code(), tonic::Code::PermissionDenied); + + let forged = service + .decode_llm_codec_request(Request::new(LlmCodecDecodeRequest { + activation_id: ACTIVATION_ID.into(), + auth_token: AUTH_TOKEN.into(), + codec_capability_id: "codec-forged".into(), + invocation_id: invocation_id.into(), + request: Some(json_envelope("nemo.relay.LlmRequest@1", &request).unwrap()), + })) + .await + .expect_err("forged capability must be rejected"); + assert_eq!(forged.code(), tonic::Code::NotFound); + + let wrong_direction = service + .decode_llm_codec_request(Request::new(LlmCodecDecodeRequest { + activation_id: ACTIVATION_ID.into(), + auth_token: AUTH_TOKEN.into(), + codec_capability_id: response_capability.clone(), + invocation_id: invocation_id.into(), + request: Some(json_envelope("nemo.relay.LlmRequest@1", &request).unwrap()), + })) + .await + .expect_err("response capability cannot decode requests"); + assert_eq!(wrong_direction.code(), tonic::Code::InvalidArgument); + + let wrong_invocation = service + .decode_llm_codec_request(Request::new(LlmCodecDecodeRequest { + activation_id: ACTIVATION_ID.into(), + auth_token: AUTH_TOKEN.into(), + codec_capability_id: request_capability.clone(), + invocation_id: "another-invocation".into(), + request: Some(json_envelope("nemo.relay.LlmRequest@1", &request).unwrap()), + })) + .await + .expect_err("a capability cannot be reused by another invocation"); + assert_eq!(wrong_invocation.code(), tonic::Code::PermissionDenied); + + let decoded = service + .decode_llm_codec_request(Request::new(LlmCodecDecodeRequest { + activation_id: ACTIVATION_ID.into(), + auth_token: AUTH_TOKEN.into(), + codec_capability_id: request_capability.clone(), + invocation_id: invocation_id.into(), + request: Some(json_envelope("nemo.relay.LlmRequest@1", &request).unwrap()), + })) + .await + .expect("active request capability decodes") + .into_inner(); + assert!(decoded.error.is_none()); + let decoded = decoded.value.expect("decoded request value"); + assert_eq!(decoded.schema, "nemo.relay.AnnotatedLlmRequest@2"); + let annotated: AnnotatedLlmRequest = decode_json_envelope(&decoded).unwrap(); + assert_eq!(annotated.model.as_deref(), Some("gpt-test")); + + let encoded = service + .encode_llm_codec_request(Request::new(LlmCodecEncodeRequest { + activation_id: ACTIVATION_ID.into(), + auth_token: AUTH_TOKEN.into(), + codec_capability_id: request_capability.clone(), + invocation_id: invocation_id.into(), + annotated_request: Some( + json_envelope("nemo.relay.AnnotatedLlmRequest@2", &annotated).unwrap(), + ), + original_request: Some(json_envelope("nemo.relay.LlmRequest@1", &request).unwrap()), + })) + .await + .expect("active request capability encodes") + .into_inner(); + assert!(encoded.error.is_none()); + let encoded = encoded.value.expect("encoded request value"); + assert_eq!(encoded.schema, "nemo.relay.LlmRequest@1"); + let encoded: LlmRequest = decode_json_envelope(&encoded).unwrap(); + assert_eq!(encoded.content["preserve"], json!(true)); + + let response = json!({ + "id": "chatcmpl-test", + "model": "gpt-test", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "secret"}, + "finish_reason": "stop" + }] + }); + let decoded = service + .decode_llm_codec_response(Request::new(LlmCodecDecodeResponse { + activation_id: ACTIVATION_ID.into(), + auth_token: AUTH_TOKEN.into(), + codec_capability_id: response_capability.clone(), + invocation_id: invocation_id.into(), + response: Some(json_envelope(JSON_SCHEMA, &response).unwrap()), + })) + .await + .expect("active response capability decodes") + .into_inner(); + assert!(decoded.error.is_none()); + + state.remove_codec(&request_capability); + state.remove_codec(&response_capability); + let expired = service + .decode_llm_codec_request(Request::new(LlmCodecDecodeRequest { + activation_id: ACTIVATION_ID.into(), + auth_token: AUTH_TOKEN.into(), + codec_capability_id: request_capability, + invocation_id: invocation_id.into(), + request: Some(json_envelope("nemo.relay.LlmRequest@1", &request).unwrap()), + })) + .await + .expect_err("removed capability must expire"); + assert_eq!(expired.code(), tonic::Code::NotFound); +} + #[tokio::test] async fn host_runtime_service_reports_poisoned_internal_locks() { enable_operational_logs(); diff --git a/crates/core/tests/unit/llm_api_tests.rs b/crates/core/tests/unit/llm_api_tests.rs index 5dd2e3914..733f5b2b7 100644 --- a/crates/core/tests/unit/llm_api_tests.rs +++ b/crates/core/tests/unit/llm_api_tests.rs @@ -12,25 +12,36 @@ use serde_json::json; use tokio_stream::StreamExt; use super::{ - LlmCallExecuteParams, LlmCallParams, LlmHandle, LlmRequest, LlmStreamCallExecuteParams, - emit_optimization_marks_with, llm_call, llm_call_execute, llm_stream_call_execute, - project_llm_request_to_current_user_turn, + CreateLlmHandleParams, LlmCallEndParams, LlmCallExecuteParams, LlmCallParams, LlmHandle, + LlmRequest, LlmStreamCallExecuteParams, create_llm_handle, emit_llm_start, + emit_optimization_marks_with, llm_call, llm_call_end, llm_call_execute, + llm_stream_call_execute, project_llm_request_to_current_user_turn, + sanitize_context_for_request_codec, sanitize_context_for_response_codec, }; use crate::api::event::{Event, ScopeCategory}; use crate::api::optimization::finalize_optimization_summary; -use crate::api::runtime::LlmJsonStream; +use crate::api::registry::{ + deregister_llm_sanitize_request_guardrail, deregister_llm_sanitize_response_guardrail, + register_llm_sanitize_request_guardrail, register_llm_sanitize_response_guardrail, +}; +use crate::api::runtime::{BuiltinLlmCodec, LlmCodecIdentity, LlmJsonStream}; use crate::api::runtime::{ NemoRelayContextState, create_scope_stack, global_context, set_thread_scope_stack, }; use crate::api::scope::{COMPACTION_EVENT_NAME, EmitMarkEventParams, event}; use crate::api::scope::{PopScopeParams, PushScopeParams, ScopeType, pop_scope, push_scope}; use crate::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; +use crate::codec::anthropic::AnthropicMessagesCodec; use crate::codec::openai_chat::OpenAIChatCodec; -use crate::codec::request::{AnnotatedLlmRequest, Message}; -use crate::codec::traits::LlmCodec; +use crate::codec::openai_responses::OpenAIResponsesCodec; +use crate::codec::request::{AnnotatedLlmRequest, Message, MessageContent}; +use crate::codec::traits::{LlmCodec, LlmResponseCodec}; use crate::error::FlowError; use crate::json::Json; -use crate::{codec::optimization::LlmOptimizationContribution, codec::response::PricingResolver}; +use crate::{ + codec::optimization::LlmOptimizationContribution, + codec::response::{AnnotatedLlmResponse, PricingResolver}, +}; fn reset_global() { let _ = spdlog::init_log_crate_proxy(); @@ -76,6 +87,115 @@ struct ProjectionFailingCodec { projection_attempts: Arc, } +struct RuntimeIdentityCodec; + +impl LlmCodec for RuntimeIdentityCodec { + fn codec_identity(&self) -> LlmCodecIdentity { + LlmCodecIdentity::Runtime("com.example.chat.v1".into()) + } + + fn decode(&self, request: &LlmRequest) -> crate::error::Result { + OpenAIChatCodec.decode(request) + } + + fn encode( + &self, + annotated: &AnnotatedLlmRequest, + original: &LlmRequest, + ) -> crate::error::Result { + OpenAIChatCodec.encode(annotated, original) + } +} + +impl LlmResponseCodec for RuntimeIdentityCodec { + fn codec_identity(&self) -> LlmCodecIdentity { + LlmCodecIdentity::Runtime("com.example.chat.v1".into()) + } + + fn decode_response(&self, response: &Json) -> crate::error::Result { + OpenAIChatCodec.decode_response(response) + } +} + +#[test] +fn sanitizer_context_preserves_all_codec_identity_states() { + let identity_only_request = crate::api::runtime::LlmSanitizeRequestContext::with_identity( + LlmCodecIdentity::Runtime("identity-only.request.v1".into()), + ); + assert_eq!( + identity_only_request.codec(), + &LlmCodecIdentity::Runtime("identity-only.request.v1".into()) + ); + assert!(identity_only_request.resolve_codec().is_none()); + + let identity_only_response = crate::api::runtime::LlmSanitizeResponseContext::with_identity( + LlmCodecIdentity::Runtime("identity-only.response.v1".into()), + ); + assert_eq!( + identity_only_response.codec(), + &LlmCodecIdentity::Runtime("identity-only.response.v1".into()) + ); + assert!(identity_only_response.resolve_codec().is_none()); + + assert_eq!( + sanitize_context_for_request_codec(None).codec(), + &LlmCodecIdentity::None + ); + assert_eq!( + sanitize_context_for_request_codec(Some(&OpenAIChatCodec)).codec(), + &LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiChat) + ); + assert_eq!( + sanitize_context_for_request_codec(Some(&OpenAIResponsesCodec)).codec(), + &LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiResponses) + ); + assert_eq!( + sanitize_context_for_request_codec(Some(&AnthropicMessagesCodec)).codec(), + &LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::AnthropicMessages) + ); + assert_eq!( + sanitize_context_for_request_codec(Some(&RuntimeIdentityCodec)).codec(), + &LlmCodecIdentity::Runtime("com.example.chat.v1".into()) + ); + assert_eq!( + sanitize_context_for_request_codec(Some(&ProjectionFailingCodec { + projection_attempts: Arc::new(AtomicUsize::new(0)), + })) + .codec(), + &LlmCodecIdentity::Opaque + ); + assert_eq!( + sanitize_context_for_response_codec(Some(&OpenAIChatCodec as &dyn LlmResponseCodec)) + .codec(), + &LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiChat) + ); + assert_eq!( + sanitize_context_for_response_codec(Some(&OpenAIResponsesCodec as &dyn LlmResponseCodec)) + .codec(), + &LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiResponses) + ); + assert_eq!( + sanitize_context_for_response_codec(Some(&AnthropicMessagesCodec as &dyn LlmResponseCodec)) + .codec(), + &LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::AnthropicMessages) + ); + assert_eq!( + sanitize_context_for_response_codec(Some(&RuntimeIdentityCodec)).codec(), + &LlmCodecIdentity::Runtime("com.example.chat.v1".into()) + ); + assert_eq!( + sanitize_context_for_response_codec(None).codec(), + &LlmCodecIdentity::None + ); + assert_eq!( + sanitize_context_for_response_codec(Some(&ProjectionFailingCodec { + projection_attempts: Arc::new(AtomicUsize::new(0)), + })) + .codec(), + &LlmCodecIdentity::Opaque + ); +} + impl LlmCodec for ProjectionFailingCodec { fn decode(&self, request: &LlmRequest) -> crate::error::Result { OpenAIChatCodec.decode(request) @@ -95,6 +215,15 @@ impl LlmCodec for ProjectionFailingCodec { } } +impl LlmResponseCodec for ProjectionFailingCodec { + fn decode_response( + &self, + response: &Json, + ) -> crate::error::Result { + OpenAIChatCodec.decode_response(response) + } +} + fn emit_compaction() { event( EmitMarkEventParams::builder() @@ -104,6 +233,486 @@ fn emit_compaction() { .unwrap(); } +fn secret_request() -> LlmRequest { + LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "SECRET"}] + }), + } +} + +fn redacted_request() -> LlmRequest { + LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "[REDACTED]"}] + }), + } +} + +fn secret_response() -> Json { + json!({ + "id": "chatcmpl-test", + "model": "gpt-4o-mini", + "choices": [{"message": {"role": "assistant", "content": "SECRET"}}] + }) +} + +fn redacted_response() -> Json { + json!({ + "id": "chatcmpl-test", + "model": "gpt-4o-mini", + "choices": [{"message": {"role": "assistant", "content": "[REDACTED]"}}] + }) +} + +#[test] +fn sanitization_invalidates_manual_annotations_without_a_codec() { + let _guard = lock_global_runtime(); + reset_global(); + set_thread_scope_stack(create_scope_stack()); + + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = events.clone(); + register_subscriber( + "manual-annotation-invalidation", + Arc::new(move |event| captured.lock().unwrap().push(event.clone())), + ) + .unwrap(); + register_llm_sanitize_request_guardrail( + "manual-annotation-invalidation-request", + 1, + Arc::new(|_request, _context| Some(redacted_request())), + ) + .unwrap(); + register_llm_sanitize_response_guardrail( + "manual-annotation-invalidation-response", + 1, + Arc::new(|_response, _context| Some(redacted_response())), + ) + .unwrap(); + + let request = secret_request(); + let handle = llm_call( + LlmCallParams::builder() + .name("manual") + .request(&request) + .annotated_request(Arc::new(OpenAIChatCodec.decode(&request).unwrap())) + .build(), + ) + .unwrap(); + let response = secret_response(); + llm_call_end( + LlmCallEndParams::builder() + .handle(&handle) + .response(response.clone()) + .annotated_response(Arc::new( + OpenAIChatCodec.decode_response(&response).unwrap(), + )) + .build(), + ) + .unwrap(); + + flush_subscribers().unwrap(); + let captured = events.lock().unwrap(); + assert_eq!(captured.len(), 2); + assert!( + captured + .iter() + .all(|event| !serde_json::to_string(event).unwrap().contains("SECRET")) + ); + assert!(captured[0].annotated_request().is_none()); + assert!(captured[1].annotated_response().is_none()); + + assert!( + deregister_llm_sanitize_request_guardrail("manual-annotation-invalidation-request") + .unwrap() + ); + assert!( + deregister_llm_sanitize_response_guardrail("manual-annotation-invalidation-response") + .unwrap() + ); + assert!(deregister_subscriber("manual-annotation-invalidation").unwrap()); +} + +#[test] +fn no_op_sanitizers_keep_manual_annotations() { + let _guard = lock_global_runtime(); + reset_global(); + set_thread_scope_stack(create_scope_stack()); + + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = events.clone(); + register_subscriber( + "manual-annotation-noop", + Arc::new(move |event| captured.lock().unwrap().push(event.clone())), + ) + .unwrap(); + register_llm_sanitize_request_guardrail( + "manual-annotation-noop-request", + 1, + Arc::new(|request, _context| Some(request)), + ) + .unwrap(); + register_llm_sanitize_response_guardrail( + "manual-annotation-noop-response", + 1, + Arc::new(|response, _context| Some(response)), + ) + .unwrap(); + + let request = secret_request(); + let handle = llm_call( + LlmCallParams::builder() + .name("manual") + .request(&request) + .annotated_request(Arc::new(OpenAIChatCodec.decode(&request).unwrap())) + .build(), + ) + .unwrap(); + let response = secret_response(); + llm_call_end( + LlmCallEndParams::builder() + .handle(&handle) + .response(response.clone()) + .annotated_response(Arc::new( + OpenAIChatCodec.decode_response(&response).unwrap(), + )) + .build(), + ) + .unwrap(); + + flush_subscribers().unwrap(); + let captured = events.lock().unwrap(); + assert_eq!(captured.len(), 2); + assert!(captured[0].annotated_request().is_some()); + assert!(captured[1].annotated_response().is_some()); + + assert!(deregister_llm_sanitize_request_guardrail("manual-annotation-noop-request").unwrap()); + assert!(deregister_llm_sanitize_response_guardrail("manual-annotation-noop-response").unwrap()); + assert!(deregister_subscriber("manual-annotation-noop").unwrap()); +} + +#[test] +fn sanitization_regenerates_annotations_with_active_codecs() { + let _guard = lock_global_runtime(); + reset_global(); + set_thread_scope_stack(create_scope_stack()); + + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = events.clone(); + register_subscriber( + "active-codec-annotation-regeneration", + Arc::new(move |event| captured.lock().unwrap().push(event.clone())), + ) + .unwrap(); + register_llm_sanitize_request_guardrail( + "active-codec-annotation-regeneration-request", + 1, + Arc::new(|_request, _context| Some(redacted_request())), + ) + .unwrap(); + register_llm_sanitize_response_guardrail( + "active-codec-annotation-regeneration-response", + 1, + Arc::new(|_response, _context| Some(redacted_response())), + ) + .unwrap(); + + let request = secret_request(); + let handle = create_llm_handle( + CreateLlmHandleParams::builder() + .name("manual-with-codec") + .build(), + ) + .unwrap(); + emit_llm_start( + &handle, + &request, + Some(Arc::new(OpenAIChatCodec.decode(&request).unwrap())), + Some(Arc::new(OpenAIChatCodec)), + ) + .unwrap(); + let response = secret_response(); + llm_call_end( + LlmCallEndParams::builder() + .handle(&handle) + .response(response.clone()) + .annotated_response(Arc::new( + OpenAIChatCodec.decode_response(&response).unwrap(), + )) + .response_codec(Arc::new(OpenAIChatCodec)) + .build(), + ) + .unwrap(); + + flush_subscribers().unwrap(); + let captured = events.lock().unwrap(); + assert_eq!(captured.len(), 2); + assert!( + captured + .iter() + .all(|event| !serde_json::to_string(event).unwrap().contains("SECRET")) + ); + assert_eq!( + captured[0].annotated_request().unwrap().messages, + vec![Message::User { + content: MessageContent::Text("[REDACTED]".to_string()), + name: None, + }] + ); + assert_eq!( + captured[1].annotated_response().unwrap().message, + Some(MessageContent::Text("[REDACTED]".to_string())) + ); + + assert!( + deregister_llm_sanitize_request_guardrail("active-codec-annotation-regeneration-request") + .unwrap() + ); + assert!( + deregister_llm_sanitize_response_guardrail("active-codec-annotation-regeneration-response") + .unwrap() + ); + assert!(deregister_subscriber("active-codec-annotation-regeneration").unwrap()); +} + +#[test] +fn buffered_null_fallback_is_sanitized_before_emission() { + let _guard = lock_global_runtime(); + reset_global(); + set_thread_scope_stack(create_scope_stack()); + + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = Arc::clone(&events); + register_subscriber( + "buffered-null-fallback", + Arc::new(move |event| captured.lock().unwrap().push(event.clone())), + ) + .unwrap(); + + let seen = Arc::new(Mutex::new(Vec::::new())); + let sanitizer_inputs = Arc::clone(&seen); + register_llm_sanitize_response_guardrail( + "buffered-null-fallback-null", + 1, + Arc::new(move |response, _context| { + sanitizer_inputs.lock().unwrap().push(response); + Some(Json::Null) + }), + ) + .unwrap(); + + let handle = create_llm_handle( + CreateLlmHandleParams::builder() + .name("buffered-null-fallback-null") + .build(), + ) + .unwrap(); + let fallback = secret_response(); + llm_call_end( + LlmCallEndParams::builder() + .handle(&handle) + .response(Json::Null) + .data(fallback.clone()) + .annotated_response(Arc::new( + OpenAIChatCodec.decode_response(&fallback).unwrap(), + )) + .response_codec(Arc::new(OpenAIChatCodec)) + .build(), + ) + .unwrap(); + assert!(deregister_llm_sanitize_response_guardrail("buffered-null-fallback-null").unwrap()); + + register_llm_sanitize_response_guardrail( + "buffered-null-fallback-redacted", + 1, + Arc::new(|_response, _context| Some(redacted_response())), + ) + .unwrap(); + let handle = create_llm_handle( + CreateLlmHandleParams::builder() + .name("buffered-null-fallback-redacted") + .build(), + ) + .unwrap(); + llm_call_end( + LlmCallEndParams::builder() + .handle(&handle) + .response(Json::Null) + .data(fallback.clone()) + .annotated_response(Arc::new( + OpenAIChatCodec.decode_response(&fallback).unwrap(), + )) + .response_codec(Arc::new(OpenAIChatCodec)) + .build(), + ) + .unwrap(); + assert!(deregister_llm_sanitize_response_guardrail("buffered-null-fallback-redacted").unwrap()); + + let handle = create_llm_handle( + CreateLlmHandleParams::builder() + .name("buffered-null-without-fallback") + .build(), + ) + .unwrap(); + llm_call_end( + LlmCallEndParams::builder() + .handle(&handle) + .response(Json::Null) + .build(), + ) + .unwrap(); + + flush_subscribers().unwrap(); + let captured = events.lock().unwrap(); + assert_eq!(*seen.lock().unwrap(), vec![fallback]); + assert_eq!(captured.len(), 3); + assert_eq!(captured[0].output(), Some(&Json::Null)); + assert!(captured[0].annotated_response().is_none()); + assert_eq!(captured[1].output(), Some(&redacted_response())); + assert_eq!( + captured[1].annotated_response().unwrap().message, + Some(MessageContent::Text("[REDACTED]".to_string())) + ); + assert!(captured[2].output().is_none()); + assert!(captured[2].annotated_response().is_none()); + assert!( + captured + .iter() + .all(|event| !serde_json::to_string(event).unwrap().contains("SECRET")) + ); + + assert!(deregister_subscriber("buffered-null-fallback").unwrap()); +} + +#[test] +fn streaming_null_fallback_is_sanitized_before_emission() { + let _guard = lock_global_runtime(); + reset_global(); + set_thread_scope_stack(create_scope_stack()); + + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = Arc::clone(&events); + register_subscriber( + "streaming-null-fallback", + Arc::new(move |event| { + if event.scope_category() == Some(ScopeCategory::End) { + captured.lock().unwrap().push(event.clone()); + } + }), + ) + .unwrap(); + + let seen = Arc::new(Mutex::new(Vec::::new())); + let sanitizer_inputs = Arc::clone(&seen); + register_llm_sanitize_response_guardrail( + "streaming-null-fallback-null", + 1, + Arc::new(move |response, _context| { + sanitizer_inputs.lock().unwrap().push(response); + Some(Json::Null) + }), + ) + .unwrap(); + + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut stream = llm_stream_call_execute( + LlmStreamCallExecuteParams::builder() + .name("streaming-null-fallback-null") + .request(request()) + .func(Arc::new(|_request| { + Box::pin(async { Ok(LlmJsonStream::new(tokio_stream::empty())) }) + })) + .collector(Box::new(|_chunk| Ok(()))) + .finalizer(Box::new(|| Json::Null)) + .data(secret_response()) + .response_codec(Arc::new(OpenAIChatCodec)) + .build(), + ) + .await + .unwrap(); + while let Some(chunk) = stream.next().await { + chunk.unwrap(); + } + }); + assert!(deregister_llm_sanitize_response_guardrail("streaming-null-fallback-null").unwrap()); + + register_llm_sanitize_response_guardrail( + "streaming-null-fallback-redacted", + 1, + Arc::new(|_response, _context| Some(redacted_response())), + ) + .unwrap(); + runtime.block_on(async { + let mut stream = llm_stream_call_execute( + LlmStreamCallExecuteParams::builder() + .name("streaming-null-fallback-redacted") + .request(request()) + .func(Arc::new(|_request| { + Box::pin(async { Ok(LlmJsonStream::new(tokio_stream::empty())) }) + })) + .collector(Box::new(|_chunk| Ok(()))) + .finalizer(Box::new(|| Json::Null)) + .data(secret_response()) + .response_codec(Arc::new(OpenAIChatCodec)) + .build(), + ) + .await + .unwrap(); + while let Some(chunk) = stream.next().await { + chunk.unwrap(); + } + }); + assert!( + deregister_llm_sanitize_response_guardrail("streaming-null-fallback-redacted").unwrap() + ); + + runtime.block_on(async { + let mut stream = llm_stream_call_execute( + LlmStreamCallExecuteParams::builder() + .name("streaming-null-without-fallback") + .request(request()) + .func(Arc::new(|_request| { + Box::pin(async { Ok(LlmJsonStream::new(tokio_stream::empty())) }) + })) + .collector(Box::new(|_chunk| Ok(()))) + .finalizer(Box::new(|| Json::Null)) + .build(), + ) + .await + .unwrap(); + while let Some(chunk) = stream.next().await { + chunk.unwrap(); + } + }); + + flush_subscribers().unwrap(); + let captured = events.lock().unwrap(); + assert_eq!(*seen.lock().unwrap(), vec![secret_response()]); + assert_eq!(captured.len(), 3); + assert_eq!(captured[0].output(), Some(&Json::Null)); + assert!(captured[0].annotated_response().is_none()); + assert_eq!(captured[1].output(), Some(&redacted_response())); + assert_eq!( + captured[1].annotated_response().unwrap().message, + Some(MessageContent::Text("[REDACTED]".to_string())) + ); + assert!(captured[2].output().is_none()); + assert!(captured[2].annotated_response().is_none()); + assert!( + captured + .iter() + .all(|event| !serde_json::to_string(event).unwrap().contains("SECRET")) + ); + + assert!(deregister_subscriber("streaming-null-fallback").unwrap()); +} + #[test] fn freshness_culls_annotations_and_repeated_compactions_are_idempotent() { let _guard = lock_global_runtime(); @@ -746,6 +1355,99 @@ fn close_boundary_freezes_identical_mark_and_summary_contributions() { ); } +#[test] +fn failed_managed_calls_sanitize_fallback_end_data() { + let _guard = lock_global_runtime(); + reset_global(); + set_thread_scope_stack(create_scope_stack()); + + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = Arc::clone(&events); + register_subscriber( + "failed-managed-call-sanitization", + Arc::new(move |event| { + if event.scope_category() == Some(ScopeCategory::End) { + captured.lock().unwrap().push(event.clone()); + } + }), + ) + .unwrap(); + + let seen = Arc::new(Mutex::new(Vec::new())); + let sanitizer_inputs = Arc::clone(&seen); + register_llm_sanitize_response_guardrail( + "failed-managed-call-sanitization", + 1, + Arc::new(move |response, context| { + sanitizer_inputs + .lock() + .unwrap() + .push((response, context.codec().clone())); + Some(redacted_response()) + }), + ) + .unwrap(); + + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let buffered_error = llm_call_execute( + LlmCallExecuteParams::builder() + .name("failed-buffered-call") + .request(request()) + .func(Arc::new(|_request| { + Box::pin(async { Err(FlowError::Internal("buffered boom".to_string())) }) + })) + .data(secret_response()) + .response_codec(Arc::new(OpenAIChatCodec)) + .build(), + ) + .await + .unwrap_err(); + assert!(buffered_error.to_string().contains("buffered boom")); + + let stream_error = match llm_stream_call_execute( + LlmStreamCallExecuteParams::builder() + .name("failed-stream-call") + .request(request()) + .func(Arc::new(|_request| { + Box::pin(async { Err(FlowError::Internal("stream setup boom".to_string())) }) + })) + .collector(Box::new(|_chunk| Ok(()))) + .finalizer(Box::new(|| Json::Null)) + .data(secret_response()) + .response_codec(Arc::new(OpenAIChatCodec)) + .build(), + ) + .await + { + Ok(_) => panic!("stream setup should fail"), + Err(error) => error, + }; + assert!(stream_error.to_string().contains("stream setup boom")); + }); + + flush_subscribers().unwrap(); + assert!( + deregister_llm_sanitize_response_guardrail("failed-managed-call-sanitization").unwrap() + ); + assert!(deregister_subscriber("failed-managed-call-sanitization").unwrap()); + + let seen = seen.lock().unwrap(); + assert_eq!(seen.len(), 2); + assert!(seen.iter().all(|(response, codec)| { + response == &secret_response() + && codec == &LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiChat) + })); + + let events = events.lock().unwrap(); + assert_eq!(events.len(), 2); + assert!(events.iter().all(|event| { + event.output() == Some(&redacted_response()) + && event.annotated_response().is_none() + && !serde_json::to_string(event).unwrap().contains("SECRET") + })); +} + #[test] fn llm_call_execute_adds_otel_status_metadata_to_end_events() { let _guard = lock_global_runtime(); diff --git a/crates/core/tests/unit/native_plugin_tests.rs b/crates/core/tests/unit/native_plugin_tests.rs index 3e3243b16..beec626af 100644 --- a/crates/core/tests/unit/native_plugin_tests.rs +++ b/crates/core/tests/unit/native_plugin_tests.rs @@ -10,11 +10,18 @@ use std::panic::{AssertUnwindSafe, catch_unwind}; use std::sync::atomic::{AtomicUsize, Ordering}; use nemo_relay_plugin::{ - NemoRelayNativeLlmNextFn, NemoRelayNativeLlmStreamNextFn, NemoRelayNativeToolNextFn, + NemoRelayNativeLlmNextFn, NemoRelayNativeLlmSanitizeRequestContext, + NemoRelayNativeLlmSanitizeResponseContext, NemoRelayNativeLlmStreamNextFn, + NemoRelayNativeToolNextFn, }; use serde_json::json; -use crate::api::runtime::{NemoRelayContextState, global_context}; +use crate::api::runtime::{ + BuiltinLlmCodec, LlmSanitizeRequestContext, LlmSanitizeResponseContext, NemoRelayContextState, + global_context, +}; +use crate::codec::openai_chat::OpenAIChatCodec; +use crate::codec::response::AnnotatedLlmResponse; struct ThreadScopeStackRestore(Option); @@ -65,6 +72,50 @@ fn assert_last_error_contains(expected: &str) { ); } +struct FailingNativeCodec; + +impl LlmCodec for FailingNativeCodec { + fn decode(&self, _request: &LlmRequest) -> FlowResult { + Err(FlowError::Internal("request decode rejected".into())) + } + + fn encode( + &self, + _annotated: &AnnotatedLlmRequest, + _original: &LlmRequest, + ) -> FlowResult { + Err(FlowError::Internal("request encode rejected".into())) + } +} + +impl LlmResponseCodec for FailingNativeCodec { + fn decode_response(&self, _response: &Json) -> FlowResult { + Err(FlowError::Internal("response decode rejected".into())) + } +} + +struct PanickingNativeCodec; + +impl LlmCodec for PanickingNativeCodec { + fn decode(&self, _request: &LlmRequest) -> FlowResult { + panic!("request decode panic") + } + + fn encode( + &self, + _annotated: &AnnotatedLlmRequest, + _original: &LlmRequest, + ) -> FlowResult { + panic!("request encode panic") + } +} + +impl LlmResponseCodec for PanickingNativeCodec { + fn decode_response(&self, _response: &Json) -> FlowResult { + panic!("response decode panic") + } +} + #[test] fn native_string_and_json_helpers_cover_abi_boundaries() { clear_native_last_error(); @@ -512,6 +563,7 @@ unsafe extern "C" fn noop_tool_execution( unsafe extern "C" fn noop_llm_request( _user_data: *mut c_void, _request_json: *const NemoRelayNativeString, + _context: NemoRelayNativeLlmSanitizeRequestContext, _out_request_json: *mut *mut NemoRelayNativeString, ) -> NemoRelayStatus { NemoRelayStatus::Ok @@ -520,6 +572,7 @@ unsafe extern "C" fn noop_llm_request( unsafe extern "C" fn noop_json( _user_data: *mut c_void, _payload_json: *const NemoRelayNativeString, + _context: NemoRelayNativeLlmSanitizeResponseContext, _out_json: *mut *mut NemoRelayNativeString, ) -> NemoRelayStatus { NemoRelayStatus::Ok @@ -705,6 +758,239 @@ fn native_registration_entrypoints_reject_null_contexts() { assert_last_error_contains("plugin context is null"); } +#[test] +fn native_codec_operations_report_json_and_codec_failures() { + let openai_request_codec = + NativeHostLlmRequestCodec(Arc::new(OpenAIChatCodec) as Arc); + let openai_response_codec = + NativeHostLlmResponseCodec(Arc::new(OpenAIChatCodec) as Arc); + let failing_request_codec = + NativeHostLlmRequestCodec(Arc::new(FailingNativeCodec) as Arc); + let failing_response_codec = + NativeHostLlmResponseCodec(Arc::new(FailingNativeCodec) as Arc); + let invalid_json = native_string("not-json"); + let mut output = ptr::null_mut(); + + assert_eq!( + unsafe { + native_llm_request_codec_decode( + ptr::from_ref(&openai_request_codec).cast(), + invalid_json, + &mut output, + ) + }, + NemoRelayStatus::Internal + ); + assert!(output.is_null()); + assert_last_error_contains("invalid request JSON"); + + assert_eq!( + unsafe { + native_llm_response_codec_decode( + ptr::from_ref(&openai_response_codec).cast(), + invalid_json, + &mut output, + ) + }, + NemoRelayStatus::Internal + ); + assert!(output.is_null()); + assert_last_error_contains("invalid response JSON"); + unsafe { native_string_free(invalid_json) }; + + let request = LlmRequest { + headers: Map::new(), + content: json!({ + "model": "gpt-test", + "messages": [{"role": "user", "content": "secret"}] + }), + }; + let request_json = native_string(&serde_json::to_string(&request).unwrap()); + assert_eq!( + unsafe { + native_llm_request_codec_decode( + ptr::from_ref(&failing_request_codec).cast(), + request_json, + &mut output, + ) + }, + NemoRelayStatus::Internal + ); + assert!(output.is_null()); + assert_last_error_contains("request decode rejected"); + + let annotated = OpenAIChatCodec.decode(&request).unwrap(); + let annotated_json = native_string(&serde_json::to_string(&annotated).unwrap()); + assert_eq!( + unsafe { + native_llm_request_codec_encode( + ptr::from_ref(&failing_request_codec).cast(), + annotated_json, + request_json, + &mut output, + ) + }, + NemoRelayStatus::Internal + ); + assert!(output.is_null()); + assert_last_error_contains("request encode rejected"); + + let response_json = native_string( + r#"{"id":"chatcmpl-test","model":"gpt-test","choices":[{"index":0,"message":{"role":"assistant","content":"secret"},"finish_reason":"stop"}]}"#, + ); + assert_eq!( + unsafe { + native_llm_response_codec_decode( + ptr::from_ref(&failing_response_codec).cast(), + response_json, + &mut output, + ) + }, + NemoRelayStatus::Internal + ); + assert!(output.is_null()); + assert_last_error_contains("response decode rejected"); + + unsafe { + native_string_free(request_json); + native_string_free(annotated_json); + native_string_free(response_json); + } +} + +#[test] +fn native_codec_operations_contain_codec_panics() { + let request_codec = + NativeHostLlmRequestCodec(Arc::new(PanickingNativeCodec) as Arc); + let response_codec = + NativeHostLlmResponseCodec(Arc::new(PanickingNativeCodec) as Arc); + let request = LlmRequest { + headers: Map::new(), + content: json!({ + "model": "gpt-test", + "messages": [{"role": "user", "content": "secret"}] + }), + }; + let request_json = native_string(&serde_json::to_string(&request).unwrap()); + let annotated = OpenAIChatCodec.decode(&request).unwrap(); + let annotated_json = native_string(&serde_json::to_string(&annotated).unwrap()); + let response_json = native_string( + r#"{"id":"chatcmpl-test","model":"gpt-test","choices":[{"index":0,"message":{"role":"assistant","content":"secret"},"finish_reason":"stop"}]}"#, + ); + let mut output = ptr::null_mut(); + + assert_eq!( + unsafe { + native_llm_request_codec_decode( + ptr::from_ref(&request_codec).cast(), + request_json, + &mut output, + ) + }, + NemoRelayStatus::Internal + ); + assert!(output.is_null()); + assert_last_error_contains("request codec decode panicked"); + + assert_eq!( + unsafe { + native_llm_request_codec_encode( + ptr::from_ref(&request_codec).cast(), + annotated_json, + request_json, + &mut output, + ) + }, + NemoRelayStatus::Internal + ); + assert!(output.is_null()); + assert_last_error_contains("request codec encode panicked"); + + assert_eq!( + unsafe { + native_llm_response_codec_decode( + ptr::from_ref(&response_codec).cast(), + response_json, + &mut output, + ) + }, + NemoRelayStatus::Internal + ); + assert!(output.is_null()); + assert_last_error_contains("response codec decode panicked"); + + unsafe { + native_string_free(request_json); + native_string_free(annotated_json); + native_string_free(response_json); + } +} + +#[test] +fn native_codec_operations_clear_output_slots_on_null_arguments() { + let request_codec = NativeHostLlmRequestCodec(Arc::new(OpenAIChatCodec) as Arc); + let response_codec = + NativeHostLlmResponseCodec(Arc::new(OpenAIChatCodec) as Arc); + let request_json = native_string( + r#"{"headers":{},"content":{"model":"gpt-test","messages":[{"role":"user","content":"hello"}]}}"#, + ); + + let request_decode_sentinel = native_string("request-decode-sentinel"); + let mut output = request_decode_sentinel; + set_native_last_error("stale request decode error"); + assert_eq!( + unsafe { + native_llm_request_codec_decode( + ptr::from_ref(&request_codec).cast(), + ptr::null(), + &mut output, + ) + }, + NemoRelayStatus::NullPointer + ); + assert!(output.is_null()); + assert_last_error_contains("request codec decode request is null"); + unsafe { native_string_free(request_decode_sentinel) }; + + let request_encode_sentinel = native_string("request-encode-sentinel"); + output = request_encode_sentinel; + set_native_last_error("stale request encode error"); + assert_eq!( + unsafe { + native_llm_request_codec_encode( + ptr::from_ref(&request_codec).cast(), + ptr::null(), + request_json, + &mut output, + ) + }, + NemoRelayStatus::NullPointer + ); + assert!(output.is_null()); + assert_last_error_contains("request codec encode annotated request is null"); + unsafe { native_string_free(request_encode_sentinel) }; + + let response_decode_sentinel = native_string("response-decode-sentinel"); + output = response_decode_sentinel; + set_native_last_error("stale response decode error"); + assert_eq!( + unsafe { + native_llm_response_codec_decode( + ptr::from_ref(&response_codec).cast(), + ptr::null(), + &mut output, + ) + }, + NemoRelayStatus::NullPointer + ); + assert!(output.is_null()); + assert_last_error_contains("response codec decode response is null"); + unsafe { + native_string_free(response_decode_sentinel); + native_string_free(request_json); + } +} + unsafe extern "C" fn tool_json_echo( _user_data: *mut c_void, _name: *const NemoRelayNativeString, @@ -730,6 +1016,7 @@ unsafe extern "C" fn tool_json_error( unsafe extern "C" fn llm_request_echo( _user_data: *mut c_void, request_json: *const NemoRelayNativeString, + _context: NemoRelayNativeLlmSanitizeRequestContext, out_request_json: *mut *mut NemoRelayNativeString, ) -> NemoRelayStatus { let request = read_native_string(request_json).unwrap(); @@ -737,6 +1024,65 @@ unsafe extern "C" fn llm_request_echo( NemoRelayStatus::Ok } +unsafe extern "C" fn llm_request_alias( + _user_data: *mut c_void, + request_json: *const NemoRelayNativeString, + _context: NemoRelayNativeLlmSanitizeRequestContext, + out_request_json: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + unsafe { *out_request_json = request_json.cast_mut() }; + NemoRelayStatus::Ok +} + +unsafe extern "C" fn llm_request_codec_round_trip( + _user_data: *mut c_void, + request_json: *const NemoRelayNativeString, + context: NemoRelayNativeLlmSanitizeRequestContext, + out_request_json: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + assert!(!context.codec.is_null()); + let mut annotated = ptr::null_mut(); + let status = + unsafe { native_llm_request_codec_decode(context.codec, request_json, &mut annotated) }; + if status != NemoRelayStatus::Ok { + return status; + } + let status = unsafe { + native_llm_request_codec_encode(context.codec, annotated, request_json, out_request_json) + }; + unsafe { native_string_free(annotated) }; + status +} + +unsafe extern "C" fn llm_response_codec_decode_and_echo( + _user_data: *mut c_void, + response_json: *const NemoRelayNativeString, + context: NemoRelayNativeLlmSanitizeResponseContext, + out_response_json: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + assert!(!context.codec.is_null()); + let mut annotated = ptr::null_mut(); + let status = + unsafe { native_llm_response_codec_decode(context.codec, response_json, &mut annotated) }; + if status != NemoRelayStatus::Ok { + return status; + } + unsafe { native_string_free(annotated) }; + let response = read_native_string(response_json).unwrap(); + unsafe { *out_response_json = native_string(&response) }; + NemoRelayStatus::Ok +} + +unsafe extern "C" fn llm_response_alias( + _user_data: *mut c_void, + response_json: *const NemoRelayNativeString, + _context: NemoRelayNativeLlmSanitizeResponseContext, + out_response_json: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + unsafe { *out_response_json = response_json.cast_mut() }; + NemoRelayStatus::Ok +} + #[test] fn native_callback_helpers_cover_success_error_and_invalid_output() { assert_eq!( @@ -755,9 +1101,172 @@ fn native_callback_helpers_cover_success_error_and_invalid_output() { content: json!({"model": "test"}), }; assert_eq!( - call_llm_request_callback(llm_request_echo, ptr::null_mut(), &request).unwrap(), - request + call_llm_sanitize_request_callback( + llm_request_echo, + ptr::null_mut(), + &request, + LlmSanitizeRequestContext::default(), + ) + .unwrap(), + Some(request) + ); + + let request = LlmRequest { + headers: Map::new(), + content: json!({"model": "alias"}), + }; + assert_eq!( + call_llm_sanitize_request_callback( + llm_request_alias, + ptr::null_mut(), + &request, + LlmSanitizeRequestContext::default(), + ) + .unwrap(), + Some(request) + ); + + let response = json!({"message": "alias"}); + assert_eq!( + call_llm_sanitize_response_callback( + llm_response_alias, + ptr::null_mut(), + &response, + LlmSanitizeResponseContext::default(), + ) + .unwrap(), + Some(response) + ); +} + +#[test] +fn native_sanitizer_context_resolves_directional_codecs() { + let codec = Arc::new(OpenAIChatCodec); + let request = LlmRequest { + headers: Map::new(), + content: json!({ + "model": "gpt-test", + "messages": [{"role": "user", "content": "secret"}], + "preserve": true + }), + }; + let sanitized = call_llm_sanitize_request_callback( + llm_request_codec_round_trip, + ptr::null_mut(), + &request, + LlmSanitizeRequestContext::for_request_codec(Some(codec.clone())), + ) + .unwrap() + .expect("native request sanitizer returns a request"); + assert_eq!(sanitized, request); + + let response = json!({ + "id": "chatcmpl-test", + "model": "gpt-test", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "secret"}, + "finish_reason": "stop" + }] + }); + let sanitized = call_llm_sanitize_response_callback( + llm_response_codec_decode_and_echo, + ptr::null_mut(), + &response, + LlmSanitizeResponseContext::for_response_codec(Some(codec)), + ) + .unwrap() + .expect("native response sanitizer returns a response"); + assert_eq!(sanitized, response); +} + +#[test] +fn native_llm_sanitize_context_preserves_all_codec_identity_states() { + let cases = [ + ( + LlmCodecIdentity::None, + NemoRelayNativeLlmCodecKind::None, + None, + ), + ( + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiChat), + NemoRelayNativeLlmCodecKind::BuiltIn, + Some("openai_chat"), + ), + ( + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiResponses), + NemoRelayNativeLlmCodecKind::BuiltIn, + Some("openai_responses"), + ), + ( + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::AnthropicMessages), + NemoRelayNativeLlmCodecKind::BuiltIn, + Some("anthropic_messages"), + ), + ( + LlmCodecIdentity::Runtime("com.example.chat.v1".into()), + NemoRelayNativeLlmCodecKind::Runtime, + Some("com.example.chat.v1"), + ), + ( + LlmCodecIdentity::Opaque, + NemoRelayNativeLlmCodecKind::Opaque, + None, + ), + ]; + + for (codec, expected_kind, expected_id) in cases { + let (codec_kind, context_id) = + native_llm_codec_identity(&codec).expect("native context conversion should succeed"); + assert_eq!(codec_kind, expected_kind); + assert_eq!( + context_id.map(|id| read_native_string(id).unwrap()), + expected_id.map(str::to_owned), + ); + if let Some(context_id) = context_id { + unsafe { native_string_free(context_id) }; + } + } +} + +#[test] +fn native_llm_sanitizer_input_allocation_failures_release_codec_ids() { + let request = LlmRequest { + headers: Map::new(), + content: json!({"model": "test"}), + }; + let identity = LlmCodecIdentity::Runtime("com.example.chat.v1".into()); + let live_before = native_string_live_allocations(); + + fail_native_string_allocation_after(1); + let request_error = call_llm_sanitize_request_callback( + llm_request_alias, + ptr::null_mut(), + &request, + LlmSanitizeRequestContext::with_identity(identity.clone()), + ) + .unwrap_err(); + assert!( + request_error + .to_string() + .contains("failed to allocate native LLM request") + ); + assert_eq!(native_string_live_allocations(), live_before); + + fail_native_string_allocation_after(1); + let response_error = call_llm_sanitize_response_callback( + llm_response_alias, + ptr::null_mut(), + &json!({"model": "test"}), + LlmSanitizeResponseContext::with_identity(identity), + ) + .unwrap_err(); + assert!( + response_error + .to_string() + .contains("failed to allocate native LLM response") ); + assert_eq!(native_string_live_allocations(), live_before); } fn tool_next(output: FlowResult) -> ToolExecutionNextFn { diff --git a/crates/core/tests/unit/plugin_tests.rs b/crates/core/tests/unit/plugin_tests.rs index f3d1824a9..b39abb9d2 100644 --- a/crates/core/tests/unit/plugin_tests.rs +++ b/crates/core/tests/unit/plugin_tests.rs @@ -1511,13 +1511,13 @@ fn test_plugin_registration_context_supports_guardrail_helpers() { ctx.register_llm_sanitize_request_guardrail( "llm_sanitize_request", 1, - Arc::new(|request| request), + Arc::new(|request, _context| Some(request)), ) .unwrap(); ctx.register_llm_sanitize_response_guardrail( "llm_sanitize_response", 1, - Arc::new(|response| response), + Arc::new(|response, _context| Some(response)), ) .unwrap(); ctx.register_llm_conditional_execution_guardrail( @@ -1651,14 +1651,14 @@ fn test_plugin_registration_context_maps_duplicate_registration_errors() { ctx.register_llm_sanitize_request_guardrail( "llm-sanitize-request", 1, - Arc::new(|request| request), + Arc::new(|request, _context| Some(request)), ) .unwrap(); expect_registration_failed( ctx.register_llm_sanitize_request_guardrail( "llm-sanitize-request", 1, - Arc::new(|request| request), + Arc::new(|request, _context| Some(request)), ), "llm sanitize request guardrail:", ); @@ -1666,14 +1666,14 @@ fn test_plugin_registration_context_maps_duplicate_registration_errors() { ctx.register_llm_sanitize_response_guardrail( "llm-sanitize-response", 1, - Arc::new(|response| response), + Arc::new(|response, _context| Some(response)), ) .unwrap(); expect_registration_failed( ctx.register_llm_sanitize_response_guardrail( "llm-sanitize-response", 1, - Arc::new(|response| response), + Arc::new(|response, _context| Some(response)), ), "llm sanitize response guardrail:", ); @@ -1815,13 +1815,13 @@ fn test_plugin_registration_context_maps_deregistration_errors() { ctx.register_llm_sanitize_request_guardrail( "llm-sanitize-request", 1, - Arc::new(|request| request), + Arc::new(|request, _context| Some(request)), ) .unwrap(); ctx.register_llm_sanitize_response_guardrail( "llm-sanitize-response", 1, - Arc::new(|response| response), + Arc::new(|response, _context| Some(response)), ) .unwrap(); ctx.register_llm_conditional_execution_guardrail("llm-conditional", 1, Arc::new(|_| Ok(None))) diff --git a/crates/ffi/nemo_relay.h b/crates/ffi/nemo_relay.h index f5acfa076..3c381015a 100644 --- a/crates/ffi/nemo_relay.h +++ b/crates/ffi/nemo_relay.h @@ -58,6 +58,29 @@ enum NemoRelayStatus { }; typedef int32_t NemoRelayStatus; +/** + * Codec identity kind supplied to an LLM sanitizer. + */ +enum NemoRelayLlmSanitizeCodecKind { + /** + * No codec was active. + */ + NEMO_RELAY_LLM_SANITIZE_CODEC_KIND_NONE = 0, + /** + * A Relay built-in codec was active. + */ + NEMO_RELAY_LLM_SANITIZE_CODEC_KIND_BUILT_IN = 1, + /** + * A runtime-registered codec was active. + */ + NEMO_RELAY_LLM_SANITIZE_CODEC_KIND_RUNTIME = 2, + /** + * A codec was active but has no registered identity. + */ + NEMO_RELAY_LLM_SANITIZE_CODEC_KIND_OPAQUE = 3, +}; +typedef uint32_t NemoRelayLlmSanitizeCodecKind; + /** * The type of scope in the agent execution hierarchy. */ @@ -149,6 +172,16 @@ typedef struct FfiLLMHandle FfiLLMHandle; */ typedef struct FfiLLMRequest FfiLLMRequest; +/** + * Borrowed, callback-scoped request codec capability supplied to an LLM sanitizer. + */ +typedef struct FfiLlmSanitizeRequestCodec FfiLlmSanitizeRequestCodec; + +/** + * Borrowed, callback-scoped response codec capability supplied to an LLM sanitizer. + */ +typedef struct FfiLlmSanitizeResponseCodec FfiLlmSanitizeResponseCodec; + /** * Opaque OpenInference subscriber handle. */ @@ -248,17 +281,61 @@ typedef char *(*NemoRelayCodecEncodeFn)(void *user_data, const struct FfiLLMRequest *original_request); /** - * Callback for LLM request sanitization. Receives an `FfiLLMRequest` and returns - * a new (possibly modified) `FfiLLMRequest`. Return null to use defaults. + * Codec identity supplied to an LLM sanitizer. `codec_id` is null for + * `None` and `Opaque`, and is valid only for the duration of the callback. */ -typedef struct FfiLLMRequest *(*NemoRelayLlmRequestCb)(void *user_data, - const struct FfiLLMRequest *request); +typedef struct NemoRelayLlmSanitizeRequestContext { + /** + * Kind of active codec identity. + */ + NemoRelayLlmSanitizeCodecKind codec_kind; + /** + * Built-in or runtime codec ID, when applicable. + */ + const char *codec_id; + /** + * Borrowed request codec capability, or null when no codec is active. + */ + const struct FfiLlmSanitizeRequestCodec *codec; +} NemoRelayLlmSanitizeRequestContext; /** - * Generic JSON-to-JSON callback, used for LLM response sanitization and intercepts. - * The returned string must be allocated with `malloc` or equivalent. + * LLM request sanitizer. It receives the request first and its codec context + * second. Return null to omit the observability payload. The request is + * borrowed, but returning that same pointer is supported as a pass-through. + * Any other non-null result transfers ownership to Relay. */ -typedef char *(*NemoRelayJsonCb)(void *user_data, const char *json); +typedef struct FfiLLMRequest *(*NemoRelayLlmSanitizeRequestCb)(void *user_data, + const struct FfiLLMRequest *request, + struct NemoRelayLlmSanitizeRequestContext context); + +/** + * Directional codec context supplied to an LLM response sanitizer. + */ +typedef struct NemoRelayLlmSanitizeResponseContext { + /** + * Kind of active codec identity. + */ + NemoRelayLlmSanitizeCodecKind codec_kind; + /** + * Built-in or runtime codec ID, when applicable. + */ + const char *codec_id; + /** + * Borrowed response codec capability, or null when no codec is active. + */ + const struct FfiLlmSanitizeResponseCodec *codec; +} NemoRelayLlmSanitizeResponseContext; + +/** + * LLM response sanitizer. It receives response JSON first and its codec + * context second. Return null to omit the observability payload. The response + * is borrowed, but returning that same pointer is supported as a pass-through. + * Any other non-null result transfers ownership to Relay. + */ +typedef char *(*NemoRelayLlmSanitizeResponseCb)(void *user_data, + const char *response_json, + struct NemoRelayLlmSanitizeResponseContext context); /** * Callback for LLM conditional execution guardrails. @@ -703,6 +780,41 @@ NemoRelayStatus nemo_relay_scope_register_scope_sanitize_end_guardrail(const cha NemoRelayStatus nemo_relay_scope_deregister_scope_sanitize_end_guardrail(const char *scope_uuid, const char *name); +/** + * Decode a request through a callback-scoped sanitizer codec capability. + * + * The returned JSON string must be freed with `nemo_relay_string_free`. + * + * # Safety + * Both pointers must be non-null and valid only during the sanitizer callback. + */ +char *nemo_relay_llm_sanitize_request_codec_decode(const struct FfiLlmSanitizeRequestCodec *codec, + const struct FfiLLMRequest *request); + +/** + * Encode normalized request changes through a callback-scoped codec capability. + * + * The returned request is owned by the caller and must be freed with + * `nemo_relay_llm_request_free`. Returns null on failure. + * + * # Safety + * All pointers must be non-null and valid only during the sanitizer callback. + */ +struct FfiLLMRequest *nemo_relay_llm_sanitize_request_codec_encode(const struct FfiLlmSanitizeRequestCodec *codec, + const char *annotated_json, + const struct FfiLLMRequest *original); + +/** + * Decode a response through a callback-scoped sanitizer codec capability. + * + * The returned JSON string must be freed with `nemo_relay_string_free`. + * + * # Safety + * All pointers must be non-null and valid only during the sanitizer callback. + */ +char *nemo_relay_llm_sanitize_response_codec_decode(const struct FfiLlmSanitizeResponseCodec *codec, + const char *response_json); + /** * Begin a manual LLM call lifecycle span. * @@ -942,8 +1054,8 @@ int32_t nemo_relay_stream_next(struct FfiStream *stream, char **out_chunk); void nemo_relay_stream_free(struct FfiStream *stream); /** - * Register an LLM request sanitization guardrail. The callback can modify or - * replace the LLM request before it is sent. + * Register an LLM request sanitizer. The callback receives the emitted + * request first and per-call codec context second; null omits observability. * * # Parameters * - `name`: Unique guardrail name. @@ -957,7 +1069,7 @@ void nemo_relay_stream_free(struct FfiStream *stream); */ NemoRelayStatus nemo_relay_register_llm_sanitize_request_guardrail(const char *name, int32_t priority, - NemoRelayLlmRequestCb cb, + NemoRelayLlmSanitizeRequestCb cb, void *user_data, NemoRelayFreeFn free_fn); @@ -985,7 +1097,7 @@ NemoRelayStatus nemo_relay_deregister_llm_sanitize_request_guardrail(const char */ NemoRelayStatus nemo_relay_register_llm_sanitize_response_guardrail(const char *name, int32_t priority, - NemoRelayJsonCb cb, + NemoRelayLlmSanitizeResponseCb cb, void *user_data, NemoRelayFreeFn free_fn); @@ -1704,7 +1816,7 @@ NemoRelayStatus nemo_relay_plugin_context_register_tool_conditional_execution_gu NemoRelayStatus nemo_relay_plugin_context_register_llm_sanitize_request_guardrail(struct FfiPluginContext *ctx, const char *name, int32_t priority, - NemoRelayLlmRequestCb cb, + NemoRelayLlmSanitizeRequestCb cb, void *user_data, NemoRelayFreeFn free_fn); @@ -1718,7 +1830,7 @@ NemoRelayStatus nemo_relay_plugin_context_register_llm_sanitize_request_guardrai NemoRelayStatus nemo_relay_plugin_context_register_llm_sanitize_response_guardrail(struct FfiPluginContext *ctx, const char *name, int32_t priority, - NemoRelayJsonCb cb, + NemoRelayLlmSanitizeResponseCb cb, void *user_data, NemoRelayFreeFn free_fn); @@ -2007,7 +2119,7 @@ NemoRelayStatus nemo_relay_scope_deregister_tool_execution_intercept(const char NemoRelayStatus nemo_relay_scope_register_llm_sanitize_request_guardrail(const char *scope_uuid, const char *name, int32_t priority, - NemoRelayLlmRequestCb cb, + NemoRelayLlmSanitizeRequestCb cb, void *user_data, NemoRelayFreeFn free_fn); @@ -2037,7 +2149,7 @@ NemoRelayStatus nemo_relay_scope_deregister_llm_sanitize_request_guardrail(const NemoRelayStatus nemo_relay_scope_register_llm_sanitize_response_guardrail(const char *scope_uuid, const char *name, int32_t priority, - NemoRelayJsonCb cb, + NemoRelayLlmSanitizeResponseCb cb, void *user_data, NemoRelayFreeFn free_fn); diff --git a/crates/ffi/src/api/llm.rs b/crates/ffi/src/api/llm.rs index 2cacc1c7b..62060d846 100644 --- a/crates/ffi/src/api/llm.rs +++ b/crates/ffi/src/api/llm.rs @@ -2,17 +2,141 @@ // SPDX-License-Identifier: Apache-2.0 use super::{ - Arc, FfiCodecHandle, FfiLLMHandle, FfiScopeHandle, FlowResult, LlmAttributes, - LlmExecutionNextFn, LlmRequest, LlmStreamExecutionNextFn, NemoRelayCodecDecodeFn, - NemoRelayCodecEncodeFn, NemoRelayCollectorCb, NemoRelayFinalizerCb, NemoRelayFreeFn, - NemoRelayLlmExecCb, NemoRelayStatus, TASK_SCOPE_STACK, c_char, c_str_to_json, - c_str_to_opt_json, c_str_to_string, clear_last_error, core_llm_api, current_scope_stack, - json_to_c_string, set_last_error, status_from_error, tokio_runtime, - unix_micros_to_opt_timestamp, wrap_codec_fn, wrap_collector_fn, wrap_finalizer_fn, - wrap_llm_exec_fn, wrap_llm_stream_exec_fn, + Arc, FfiCodecHandle, FfiLLMHandle, FfiLLMRequest, FfiLlmSanitizeRequestCodec, + FfiLlmSanitizeResponseCodec, FfiScopeHandle, FlowResult, LlmAttributes, LlmExecutionNextFn, + LlmRequest, LlmStreamExecutionNextFn, NemoRelayCodecDecodeFn, NemoRelayCodecEncodeFn, + NemoRelayCollectorCb, NemoRelayFinalizerCb, NemoRelayFreeFn, NemoRelayLlmExecCb, + NemoRelayStatus, TASK_SCOPE_STACK, c_char, c_str_to_json, c_str_to_opt_json, c_str_to_string, + clear_last_error, core_llm_api, current_scope_stack, json_to_c_string, set_last_error, + status_from_error, tokio_runtime, unix_micros_to_opt_timestamp, wrap_codec_fn, + wrap_collector_fn, wrap_finalizer_fn, wrap_llm_exec_fn, wrap_llm_stream_exec_fn, }; +use std::panic::{AssertUnwindSafe, catch_unwind}; use tokio_stream::StreamExt; +/// Decode a request through a callback-scoped sanitizer codec capability. +/// +/// The returned JSON string must be freed with `nemo_relay_string_free`. +/// +/// # Safety +/// Both pointers must be non-null and valid only during the sanitizer callback. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_llm_sanitize_request_codec_decode( + codec: *const FfiLlmSanitizeRequestCodec, + request: *const FfiLLMRequest, +) -> *mut c_char { + clear_last_error(); + if codec.is_null() || request.is_null() { + set_last_error("null sanitizer request codec argument"); + return std::ptr::null_mut(); + } + let result = catch_unwind(AssertUnwindSafe( + || -> std::result::Result<*mut c_char, String> { + let annotated = unsafe { &*codec } + .0 + .decode(&unsafe { &*request }.0) + .map_err(|error| error.to_string())?; + let value = serde_json::to_value(annotated).map_err(|error| error.to_string())?; + Ok(json_to_c_string(&value)) + }, + )); + match result { + Ok(Ok(value)) => value, + Ok(Err(error)) => { + set_last_error(&error); + std::ptr::null_mut() + } + Err(_) => { + set_last_error("sanitizer request codec decode panicked"); + std::ptr::null_mut() + } + } +} + +/// Encode normalized request changes through a callback-scoped codec capability. +/// +/// The returned request is owned by the caller and must be freed with +/// `nemo_relay_llm_request_free`. Returns null on failure. +/// +/// # Safety +/// All pointers must be non-null and valid only during the sanitizer callback. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_llm_sanitize_request_codec_encode( + codec: *const FfiLlmSanitizeRequestCodec, + annotated_json: *const c_char, + original: *const FfiLLMRequest, +) -> *mut FfiLLMRequest { + clear_last_error(); + if codec.is_null() || annotated_json.is_null() || original.is_null() { + set_last_error("null sanitizer request codec argument"); + return std::ptr::null_mut(); + } + let result = catch_unwind(AssertUnwindSafe( + || -> std::result::Result<*mut FfiLLMRequest, String> { + let annotated = c_str_to_json(annotated_json) + .and_then(|value| serde_json::from_value(value).ok()) + .ok_or_else(|| "invalid annotated request JSON".to_string())?; + let request = unsafe { &*codec } + .0 + .encode(&annotated, &unsafe { &*original }.0) + .map_err(|error| error.to_string())?; + Ok(Box::into_raw(Box::new(FfiLLMRequest(request)))) + }, + )); + match result { + Ok(Ok(value)) => value, + Ok(Err(error)) => { + set_last_error(&error); + std::ptr::null_mut() + } + Err(_) => { + set_last_error("sanitizer request codec encode panicked"); + std::ptr::null_mut() + } + } +} + +/// Decode a response through a callback-scoped sanitizer codec capability. +/// +/// The returned JSON string must be freed with `nemo_relay_string_free`. +/// +/// # Safety +/// All pointers must be non-null and valid only during the sanitizer callback. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_llm_sanitize_response_codec_decode( + codec: *const FfiLlmSanitizeResponseCodec, + response_json: *const c_char, +) -> *mut c_char { + clear_last_error(); + if codec.is_null() || response_json.is_null() { + set_last_error("null sanitizer response codec argument"); + return std::ptr::null_mut(); + } + let result = catch_unwind(AssertUnwindSafe( + || -> std::result::Result<*mut c_char, String> { + let response = + c_str_to_json(response_json).ok_or_else(|| "invalid response JSON".to_string())?; + let annotated = unsafe { &*codec } + .0 + .decode_response(&response) + .map_err(|error| error.to_string())?; + let value = serde_json::to_value(annotated).map_err(|error| error.to_string())?; + Ok(json_to_c_string(&value)) + }, + )); + match result { + Ok(Ok(value)) => value, + Ok(Err(error)) => { + set_last_error(&error); + std::ptr::null_mut() + } + Err(_) => { + set_last_error("sanitizer response codec decode panicked"); + std::ptr::null_mut() + } + } +} + // --------------------------------------------------------------------------- // LLM lifecycle // --------------------------------------------------------------------------- diff --git a/crates/ffi/src/api/llm_registry.rs b/crates/ffi/src/api/llm_registry.rs index 1d53e4ce4..38883ecee 100644 --- a/crates/ffi/src/api/llm_registry.rs +++ b/crates/ffi/src/api/llm_registry.rs @@ -2,20 +2,20 @@ // SPDX-License-Identifier: Apache-2.0 use super::{ - NemoRelayEventSubscriberCb, NemoRelayFreeFn, NemoRelayJsonCb, NemoRelayLlmConditionalCb, - NemoRelayLlmExecInterceptCb, NemoRelayLlmRequestCb, NemoRelayLlmRequestInterceptCb, - NemoRelayStatus, c_char, c_str_to_string, clear_last_error, core_registry_api, - core_subscriber_api, status_from_error, wrap_event_subscriber, wrap_llm_conditional_fn, - wrap_llm_exec_intercept_fn, wrap_llm_request_intercept_fn, wrap_llm_response_fn, - wrap_llm_sanitize_request_fn, wrap_llm_stream_exec_intercept_fn, + NemoRelayEventSubscriberCb, NemoRelayFreeFn, NemoRelayLlmConditionalCb, + NemoRelayLlmExecInterceptCb, NemoRelayLlmRequestInterceptCb, NemoRelayLlmSanitizeRequestCb, + NemoRelayLlmSanitizeResponseCb, NemoRelayStatus, c_char, c_str_to_string, clear_last_error, + core_registry_api, core_subscriber_api, status_from_error, wrap_event_subscriber, + wrap_llm_conditional_fn, wrap_llm_exec_intercept_fn, wrap_llm_request_intercept_fn, + wrap_llm_sanitize_request_fn, wrap_llm_sanitize_response_fn, wrap_llm_stream_exec_intercept_fn, }; // --------------------------------------------------------------------------- // LLM guardrail registrations // --------------------------------------------------------------------------- -/// Register an LLM request sanitization guardrail. The callback can modify or -/// replace the LLM request before it is sent. +/// Register an LLM request sanitizer. The callback receives the emitted +/// request first and per-call codec context second; null omits observability. /// /// # Parameters /// - `name`: Unique guardrail name. @@ -30,7 +30,7 @@ use super::{ pub unsafe extern "C" fn nemo_relay_register_llm_sanitize_request_guardrail( name: *const c_char, priority: i32, - cb: NemoRelayLlmRequestCb, + cb: NemoRelayLlmSanitizeRequestCb, user_data: *mut libc::c_void, free_fn: NemoRelayFreeFn, ) -> NemoRelayStatus { @@ -81,7 +81,7 @@ pub unsafe extern "C" fn nemo_relay_deregister_llm_sanitize_request_guardrail( pub unsafe extern "C" fn nemo_relay_register_llm_sanitize_response_guardrail( name: *const c_char, priority: i32, - cb: NemoRelayJsonCb, + cb: NemoRelayLlmSanitizeResponseCb, user_data: *mut libc::c_void, free_fn: NemoRelayFreeFn, ) -> NemoRelayStatus { @@ -90,7 +90,7 @@ pub unsafe extern "C" fn nemo_relay_register_llm_sanitize_response_guardrail( Ok(s) => s, Err(status) => return status, }; - let wrapped = wrap_llm_response_fn(cb, user_data, free_fn); + let wrapped = wrap_llm_sanitize_response_fn(cb, user_data, free_fn); match core_registry_api::register_llm_sanitize_response_guardrail(&name, priority, wrapped) { Ok(()) => NemoRelayStatus::Ok, Err(e) => status_from_error(&e), diff --git a/crates/ffi/src/api/mod.rs b/crates/ffi/src/api/mod.rs index 7759b6e3b..42f40bdca 100644 --- a/crates/ffi/src/api/mod.rs +++ b/crates/ffi/src/api/mod.rs @@ -15,14 +15,14 @@ use std::time::Duration; use crate::callable::{ NemoRelayCodecDecodeFn, NemoRelayCodecEncodeFn, NemoRelayCollectorCb, NemoRelayEventSanitizeCb, - NemoRelayEventSubscriberCb, NemoRelayFinalizerCb, NemoRelayFreeFn, NemoRelayJsonCb, - NemoRelayLlmConditionalCb, NemoRelayLlmExecCb, NemoRelayLlmExecInterceptCb, - NemoRelayLlmRequestCb, NemoRelayLlmRequestInterceptCb, NemoRelayPluginRegisterCb, + NemoRelayEventSubscriberCb, NemoRelayFinalizerCb, NemoRelayFreeFn, NemoRelayLlmConditionalCb, + NemoRelayLlmExecCb, NemoRelayLlmExecInterceptCb, NemoRelayLlmRequestInterceptCb, + NemoRelayLlmSanitizeRequestCb, NemoRelayLlmSanitizeResponseCb, NemoRelayPluginRegisterCb, NemoRelayPluginValidateCb, NemoRelayToolConditionalCb, NemoRelayToolExecCb, NemoRelayToolExecInterceptCb, NemoRelayToolSanitizeCb, wrap_codec_fn, wrap_collector_fn, wrap_event_sanitize_fn, wrap_event_subscriber, wrap_finalizer_fn, wrap_llm_conditional_fn, wrap_llm_exec_fn, wrap_llm_exec_intercept_fn, wrap_llm_request_intercept_fn, - wrap_llm_response_fn, wrap_llm_sanitize_request_fn, wrap_llm_stream_exec_fn, + wrap_llm_sanitize_request_fn, wrap_llm_sanitize_response_fn, wrap_llm_stream_exec_fn, wrap_llm_stream_exec_intercept_fn, wrap_tool_conditional_fn, wrap_tool_exec_fn, wrap_tool_exec_intercept_fn, wrap_tool_request_intercept_fn, wrap_tool_sanitize_fn, }; @@ -36,8 +36,9 @@ use crate::error::{ }; use crate::types::{ FfiAtifExporter, FfiAtofExporter, FfiCodecHandle, FfiLLMHandle, FfiLLMRequest, - FfiOpenInferenceSubscriber, FfiOpenTelemetrySubscriber, FfiPluginActivation, FfiPluginContext, - FfiScopeHandle, FfiScopeStack, FfiThreadScopeStackBinding, FfiToolHandle, NemoRelayScopeType, + FfiLlmSanitizeRequestCodec, FfiLlmSanitizeResponseCodec, FfiOpenInferenceSubscriber, + FfiOpenTelemetrySubscriber, FfiPluginActivation, FfiPluginContext, FfiScopeHandle, + FfiScopeStack, FfiThreadScopeStackBinding, FfiToolHandle, NemoRelayScopeType, }; pub use crate::types::{nemo_relay_openinference_subscriber_free, nemo_relay_otel_subscriber_free}; use libc::c_char; diff --git a/crates/ffi/src/api/plugin.rs b/crates/ffi/src/api/plugin.rs index 9130a2f94..09cd151f8 100644 --- a/crates/ffi/src/api/plugin.rs +++ b/crates/ffi/src/api/plugin.rs @@ -4,19 +4,19 @@ use super::{ Arc, CStr, ConfigDiagnostic, DiagnosticLevel, DynamicPluginActivationSpec, FfiPluginActivation, FfiPluginContext, Future, NemoRelayEventSanitizeCb, NemoRelayEventSubscriberCb, - NemoRelayFreeFn, NemoRelayJsonCb, NemoRelayLlmConditionalCb, NemoRelayLlmExecInterceptCb, - NemoRelayLlmRequestCb, NemoRelayLlmRequestInterceptCb, NemoRelayPluginRegisterCb, - NemoRelayPluginValidateCb, NemoRelayStatus, NemoRelayToolConditionalCb, - NemoRelayToolExecInterceptCb, NemoRelayToolSanitizeCb, Pin, Plugin, PluginConfig, PluginError, - PluginHostActivation, PluginRegistrationContext, active_plugin_report, c_char, c_str_to_json, - c_str_to_string, clear_last_error, clear_plugin_configuration, deregister_plugin, - initialize_plugins, json_to_c_string, last_error_message, list_plugin_kinds, - nemo_relay_string_free, register_adaptive_component, register_plugin, set_last_error, - status_from_plugin_error, tokio_runtime, validate_plugin_config, wrap_event_sanitize_fn, - wrap_event_subscriber, wrap_llm_conditional_fn, wrap_llm_exec_intercept_fn, - wrap_llm_request_intercept_fn, wrap_llm_response_fn, wrap_llm_sanitize_request_fn, - wrap_llm_stream_exec_intercept_fn, wrap_tool_conditional_fn, wrap_tool_exec_intercept_fn, - wrap_tool_request_intercept_fn, wrap_tool_sanitize_fn, + NemoRelayFreeFn, NemoRelayLlmConditionalCb, NemoRelayLlmExecInterceptCb, + NemoRelayLlmRequestInterceptCb, NemoRelayLlmSanitizeRequestCb, NemoRelayLlmSanitizeResponseCb, + NemoRelayPluginRegisterCb, NemoRelayPluginValidateCb, NemoRelayStatus, + NemoRelayToolConditionalCb, NemoRelayToolExecInterceptCb, NemoRelayToolSanitizeCb, Pin, Plugin, + PluginConfig, PluginError, PluginHostActivation, PluginRegistrationContext, + active_plugin_report, c_char, c_str_to_json, c_str_to_string, clear_last_error, + clear_plugin_configuration, deregister_plugin, initialize_plugins, json_to_c_string, + last_error_message, list_plugin_kinds, nemo_relay_string_free, register_adaptive_component, + register_plugin, set_last_error, status_from_plugin_error, tokio_runtime, + validate_plugin_config, wrap_event_sanitize_fn, wrap_event_subscriber, wrap_llm_conditional_fn, + wrap_llm_exec_intercept_fn, wrap_llm_request_intercept_fn, wrap_llm_sanitize_request_fn, + wrap_llm_sanitize_response_fn, wrap_llm_stream_exec_intercept_fn, wrap_tool_conditional_fn, + wrap_tool_exec_intercept_fn, wrap_tool_request_intercept_fn, wrap_tool_sanitize_fn, }; use crate::api::event_registry::Surface; use nemo_relay_pii_redaction::component::register_pii_redaction_component; @@ -731,7 +731,7 @@ pub unsafe extern "C" fn nemo_relay_plugin_context_register_llm_sanitize_request ctx: *mut FfiPluginContext, name: *const c_char, priority: i32, - cb: NemoRelayLlmRequestCb, + cb: NemoRelayLlmSanitizeRequestCb, user_data: *mut libc::c_void, free_fn: NemoRelayFreeFn, ) -> NemoRelayStatus { @@ -763,7 +763,7 @@ pub unsafe extern "C" fn nemo_relay_plugin_context_register_llm_sanitize_respons ctx: *mut FfiPluginContext, name: *const c_char, priority: i32, - cb: NemoRelayJsonCb, + cb: NemoRelayLlmSanitizeResponseCb, user_data: *mut libc::c_void, free_fn: NemoRelayFreeFn, ) -> NemoRelayStatus { @@ -776,7 +776,7 @@ pub unsafe extern "C" fn nemo_relay_plugin_context_register_llm_sanitize_respons Ok(value) => value, Err(status) => return status, }; - let wrapped = wrap_llm_response_fn(cb, user_data, free_fn); + let wrapped = wrap_llm_sanitize_response_fn(cb, user_data, free_fn); match unsafe { &mut *((*ctx).0) } .register_llm_sanitize_response_guardrail(&name, priority, wrapped) { diff --git a/crates/ffi/src/api/scope_registry.rs b/crates/ffi/src/api/scope_registry.rs index 0b2887155..50efd644b 100644 --- a/crates/ffi/src/api/scope_registry.rs +++ b/crates/ffi/src/api/scope_registry.rs @@ -2,15 +2,15 @@ // SPDX-License-Identifier: Apache-2.0 use super::{ - NemoRelayEventSubscriberCb, NemoRelayFreeFn, NemoRelayJsonCb, NemoRelayLlmConditionalCb, - NemoRelayLlmExecInterceptCb, NemoRelayLlmRequestCb, NemoRelayLlmRequestInterceptCb, - NemoRelayStatus, NemoRelayToolConditionalCb, NemoRelayToolExecInterceptCb, - NemoRelayToolSanitizeCb, c_char, c_str_to_string, clear_last_error, core_registry_api, - core_subscriber_api, set_last_error, status_from_error, wrap_event_subscriber, - wrap_llm_conditional_fn, wrap_llm_exec_intercept_fn, wrap_llm_request_intercept_fn, - wrap_llm_response_fn, wrap_llm_sanitize_request_fn, wrap_llm_stream_exec_intercept_fn, - wrap_tool_conditional_fn, wrap_tool_exec_intercept_fn, wrap_tool_request_intercept_fn, - wrap_tool_sanitize_fn, + NemoRelayEventSubscriberCb, NemoRelayFreeFn, NemoRelayLlmConditionalCb, + NemoRelayLlmExecInterceptCb, NemoRelayLlmRequestInterceptCb, NemoRelayLlmSanitizeRequestCb, + NemoRelayLlmSanitizeResponseCb, NemoRelayStatus, NemoRelayToolConditionalCb, + NemoRelayToolExecInterceptCb, NemoRelayToolSanitizeCb, c_char, c_str_to_string, + clear_last_error, core_registry_api, core_subscriber_api, set_last_error, status_from_error, + wrap_event_subscriber, wrap_llm_conditional_fn, wrap_llm_exec_intercept_fn, + wrap_llm_request_intercept_fn, wrap_llm_sanitize_request_fn, wrap_llm_sanitize_response_fn, + wrap_llm_stream_exec_intercept_fn, wrap_tool_conditional_fn, wrap_tool_exec_intercept_fn, + wrap_tool_request_intercept_fn, wrap_tool_sanitize_fn, }; // --------------------------------------------------------------------------- @@ -363,7 +363,7 @@ pub unsafe extern "C" fn nemo_relay_scope_register_llm_sanitize_request_guardrai scope_uuid: *const c_char, name: *const c_char, priority: i32, - cb: NemoRelayLlmRequestCb, + cb: NemoRelayLlmSanitizeRequestCb, user_data: *mut libc::c_void, free_fn: NemoRelayFreeFn, ) -> NemoRelayStatus { @@ -426,7 +426,7 @@ pub unsafe extern "C" fn nemo_relay_scope_register_llm_sanitize_response_guardra scope_uuid: *const c_char, name: *const c_char, priority: i32, - cb: NemoRelayJsonCb, + cb: NemoRelayLlmSanitizeResponseCb, user_data: *mut libc::c_void, free_fn: NemoRelayFreeFn, ) -> NemoRelayStatus { @@ -439,7 +439,7 @@ pub unsafe extern "C" fn nemo_relay_scope_register_llm_sanitize_response_guardra Ok(s) => s, Err(status) => return status, }; - let wrapped = wrap_llm_response_fn(cb, user_data, free_fn); + let wrapped = wrap_llm_sanitize_response_fn(cb, user_data, free_fn); match core_registry_api::scope_register_llm_sanitize_response_guardrail( &uuid, &name, priority, wrapped, ) { diff --git a/crates/ffi/src/callable.rs b/crates/ffi/src/callable.rs index 3b968a193..f327e01d2 100644 --- a/crates/ffi/src/callable.rs +++ b/crates/ffi/src/callable.rs @@ -23,9 +23,10 @@ use std::sync::Arc; use libc::c_char; use nemo_relay::api::runtime::{ - EventSanitizeFn, EventSubscriberFn, LlmConditionalFn, LlmExecutionNextFn, LlmJsonStream, - LlmRequestInterceptFn, LlmSanitizeRequestFn, LlmSanitizeResponseFn, LlmStreamExecutionNextFn, - ToolConditionalFn, ToolExecutionFn, ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, + EventSanitizeFn, EventSubscriberFn, LlmCodecIdentity, LlmConditionalFn, LlmExecutionNextFn, + LlmJsonStream, LlmRequestInterceptFn, LlmSanitizeRequestContext, LlmSanitizeRequestFn, + LlmSanitizeResponseContext, LlmSanitizeResponseFn, LlmStreamExecutionNextFn, ToolConditionalFn, + ToolExecutionFn, ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, }; use serde_json::Value as Json; use tokio_stream::StreamExt; @@ -37,7 +38,7 @@ use nemo_relay::codec::request::AnnotatedLlmRequest as AnnotatedLLMRequest; use nemo_relay::codec::traits::LlmCodec; use nemo_relay::error::{FlowError, Result}; -use crate::convert::json_to_c_string; +use crate::convert::{c_str_to_json, json_to_c_string}; use crate::error::{NemoRelayStatus, clear_last_error, last_error_message, set_last_error}; use crate::types::{FfiEvent, FfiLLMRequest, FfiPluginContext}; @@ -97,18 +98,63 @@ pub type NemoRelayToolExecInterceptCb = unsafe extern "C" fn( next_ctx: *mut libc::c_void, ) -> *mut c_char; -/// Generic JSON-to-JSON callback, used for LLM response sanitization and intercepts. -/// The returned string must be allocated with `malloc` or equivalent. -pub type NemoRelayJsonCb = - unsafe extern "C" fn(user_data: *mut libc::c_void, json: *const c_char) -> *mut c_char; +/// Codec identity kind supplied to an LLM sanitizer. +#[repr(u32)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum NemoRelayLlmSanitizeCodecKind { + /// No codec was active. + None = 0, + /// A Relay built-in codec was active. + BuiltIn = 1, + /// A runtime-registered codec was active. + Runtime = 2, + /// A codec was active but has no registered identity. + Opaque = 3, +} + +/// Codec identity supplied to an LLM sanitizer. `codec_id` is null for +/// `None` and `Opaque`, and is valid only for the duration of the callback. +#[repr(C)] +pub struct NemoRelayLlmSanitizeRequestContext { + /// Kind of active codec identity. + pub codec_kind: NemoRelayLlmSanitizeCodecKind, + /// Built-in or runtime codec ID, when applicable. + pub codec_id: *const c_char, + /// Borrowed request codec capability, or null when no codec is active. + pub codec: *const crate::types::FfiLlmSanitizeRequestCodec, +} -/// Callback for LLM request sanitization. Receives an `FfiLLMRequest` and returns -/// a new (possibly modified) `FfiLLMRequest`. Return null to use defaults. -pub type NemoRelayLlmRequestCb = unsafe extern "C" fn( +/// Directional codec context supplied to an LLM response sanitizer. +#[repr(C)] +pub struct NemoRelayLlmSanitizeResponseContext { + /// Kind of active codec identity. + pub codec_kind: NemoRelayLlmSanitizeCodecKind, + /// Built-in or runtime codec ID, when applicable. + pub codec_id: *const c_char, + /// Borrowed response codec capability, or null when no codec is active. + pub codec: *const crate::types::FfiLlmSanitizeResponseCodec, +} + +/// LLM request sanitizer. It receives the request first and its codec context +/// second. Return null to omit the observability payload. The request is +/// borrowed, but returning that same pointer is supported as a pass-through. +/// Any other non-null result transfers ownership to Relay. +pub type NemoRelayLlmSanitizeRequestCb = unsafe extern "C" fn( user_data: *mut libc::c_void, request: *const FfiLLMRequest, + context: NemoRelayLlmSanitizeRequestContext, ) -> *mut FfiLLMRequest; +/// LLM response sanitizer. It receives response JSON first and its codec +/// context second. Return null to omit the observability payload. The response +/// is borrowed, but returning that same pointer is supported as a pass-through. +/// Any other non-null result transfers ownership to Relay. +pub type NemoRelayLlmSanitizeResponseCb = unsafe extern "C" fn( + user_data: *mut libc::c_void, + response_json: *const c_char, + context: NemoRelayLlmSanitizeResponseContext, +) -> *mut c_char; + /// Callback for LLM conditional execution guardrails. /// Returns NULL to allow execution, or an error message string to reject. pub type NemoRelayLlmConditionalCb = unsafe extern "C" fn( @@ -560,23 +606,6 @@ pub fn wrap_llm_stream_exec_intercept_fn( ) } -/// Wrap a generic C JSON callback into a Rust closure. -pub fn wrap_json_fn( - cb: NemoRelayJsonCb, - user_data: *mut libc::c_void, - free_fn: NemoRelayFreeFn, -) -> Box Json + Send + Sync> { - let ud = make_user_data(user_data, free_fn); - Box::new(move |value: Json| { - let c_json = json_to_c_string(&value); - let result_ptr = unsafe { cb(ud.ptr, c_json) }; - unsafe { nemo_relay_string_free_internal(c_json) }; - let result = ptr_to_json(result_ptr); - unsafe { nemo_relay_string_free_internal(result_ptr) }; - result - }) -} - /// Wrap a C LLM request intercept callback (annotated-aware) into a Rust /// `LlmRequestInterceptFn` closure. The callback receives the intercept name, /// the opaque `FfiLLMRequest`, and the annotated JSON (or null). It writes one @@ -650,47 +679,107 @@ pub fn wrap_llm_request_intercept_fn( ) } -/// Wrap a C JSON callback into a `Fn(Json) -> Json` closure for LLM response -/// sanitization. The callback receives the response as a JSON string and -/// returns the (possibly modified) JSON string. -pub fn wrap_llm_response_fn( - cb: NemoRelayJsonCb, +/// Wrap a C LLM request sanitizer into a Rust closure. +pub fn wrap_llm_sanitize_request_fn( + cb: NemoRelayLlmSanitizeRequestCb, user_data: *mut libc::c_void, free_fn: NemoRelayFreeFn, -) -> LlmSanitizeResponseFn { +) -> LlmSanitizeRequestFn { let ud = make_user_data(user_data, free_fn); - Arc::new(move |response: Json| { - let c_json = json_to_c_string(&response); - let result_ptr = unsafe { cb(ud.ptr, c_json) }; - unsafe { nemo_relay_string_free_internal(c_json) }; - let result_json = ptr_to_json(result_ptr); - unsafe { nemo_relay_string_free_internal(result_ptr) }; - result_json - }) + Arc::new( + move |request: LlmRequest, context: LlmSanitizeRequestContext| { + clear_last_error(); + let (codec_kind, codec_id) = match ffi_codec_identity(context.codec()) { + Ok(identity) => identity, + Err(error) => { + set_last_error(&error.to_string()); + return None; + } + }; + let codec = context + .resolve_codec() + .map(crate::types::FfiLlmSanitizeRequestCodec); + let ffi_context = NemoRelayLlmSanitizeRequestContext { + codec_kind, + codec_id: codec_id + .as_ref() + .map_or(std::ptr::null(), |name| name.as_ptr()), + codec: codec.as_ref().map_or(std::ptr::null(), std::ptr::from_ref), + }; + let ffi_req = Box::into_raw(Box::new(FfiLLMRequest(request))); + let result_ptr = unsafe { cb(ud.ptr, ffi_req, ffi_context) }; + if result_ptr.is_null() { + unsafe { drop(Box::from_raw(ffi_req)) }; + return None; + } + if result_ptr == ffi_req { + return Some(unsafe { Box::from_raw(ffi_req) }.0); + } + unsafe { drop(Box::from_raw(ffi_req)) }; + Some(unsafe { Box::from_raw(result_ptr) }.0) + }, + ) } -/// Wrap a C LLM request sanitize callback into a Rust closure. -pub fn wrap_llm_sanitize_request_fn( - cb: NemoRelayLlmRequestCb, +/// Wrap a C LLM response sanitizer into a Rust closure. +pub fn wrap_llm_sanitize_response_fn( + cb: NemoRelayLlmSanitizeResponseCb, user_data: *mut libc::c_void, free_fn: NemoRelayFreeFn, -) -> LlmSanitizeRequestFn { +) -> LlmSanitizeResponseFn { let ud = make_user_data(user_data, free_fn); - Arc::new(move |request: LlmRequest| { - let ffi_req = Box::into_raw(Box::new(FfiLLMRequest(request))); - let result_ptr = unsafe { cb(ud.ptr, ffi_req) }; - // Free the input request - unsafe { drop(Box::from_raw(ffi_req)) }; + Arc::new(move |response: Json, context: LlmSanitizeResponseContext| { + clear_last_error(); + let (codec_kind, codec_id) = match ffi_codec_identity(context.codec()) { + Ok(identity) => identity, + Err(error) => { + set_last_error(&error.to_string()); + return None; + } + }; + let codec = context + .resolve_codec() + .map(crate::types::FfiLlmSanitizeResponseCodec); + let ffi_context = NemoRelayLlmSanitizeResponseContext { + codec_kind, + codec_id: codec_id + .as_ref() + .map_or(std::ptr::null(), |name| name.as_ptr()), + codec: codec.as_ref().map_or(std::ptr::null(), std::ptr::from_ref), + }; + let response_json = json_to_c_string(&response); + let result_ptr = unsafe { cb(ud.ptr, response_json, ffi_context) }; if result_ptr.is_null() { - // If callback returns null, return a default - LlmRequest { - headers: serde_json::Map::new(), - content: Json::Null, + unsafe { nemo_relay_string_free_internal(response_json) }; + return None; + } + let result = c_str_to_json(result_ptr); + unsafe { + nemo_relay_string_free_internal(response_json); + if result_ptr != response_json { + nemo_relay_string_free_internal(result_ptr); } - } else { - let result = unsafe { Box::from_raw(result_ptr) }; - result.0 } + result + }) +} + +fn ffi_codec_identity( + identity: &LlmCodecIdentity, +) -> Result<(NemoRelayLlmSanitizeCodecKind, Option)> { + Ok(match identity { + LlmCodecIdentity::None => (NemoRelayLlmSanitizeCodecKind::None, None), + LlmCodecIdentity::BuiltIn(codec) => ( + NemoRelayLlmSanitizeCodecKind::BuiltIn, + Some(CString::new(codec.id()).expect("built-in codec IDs never contain NUL")), + ), + LlmCodecIdentity::Runtime(id) => ( + NemoRelayLlmSanitizeCodecKind::Runtime, + Some(CString::new(id.as_str()).map_err(|_| { + FlowError::InvalidArgument("runtime codec ID contains an embedded NUL".to_string()) + })?), + ), + LlmCodecIdentity::Opaque => (NemoRelayLlmSanitizeCodecKind::Opaque, None), }) } diff --git a/crates/ffi/src/types/mod.rs b/crates/ffi/src/types/mod.rs index e66f0c4e1..26bf32577 100644 --- a/crates/ffi/src/types/mod.rs +++ b/crates/ffi/src/types/mod.rs @@ -47,6 +47,10 @@ pub struct FfiToolHandle(pub ToolHandle); pub struct FfiLLMHandle(pub LlmHandle); /// Opaque wrapper around an LLM request (headers, content). pub struct FfiLLMRequest(pub LlmRequest); +/// Borrowed, callback-scoped request codec capability supplied to an LLM sanitizer. +pub struct FfiLlmSanitizeRequestCodec(pub std::sync::Arc); +/// Borrowed, callback-scoped response codec capability supplied to an LLM sanitizer. +pub struct FfiLlmSanitizeResponseCodec(pub std::sync::Arc); /// Opaque wrapper around a lifecycle event emitted by the runtime. pub struct FfiEvent(pub Event); /// Opaque handle to an isolated scope stack for per-request/per-task isolation. diff --git a/crates/ffi/tests/integration/api_tests.rs b/crates/ffi/tests/integration/api_tests.rs index fd6a26951..140f089ce 100644 --- a/crates/ffi/tests/integration/api_tests.rs +++ b/crates/ffi/tests/integration/api_tests.rs @@ -14,7 +14,10 @@ use nemo_relay::plugin::PluginRegistrationContext; use serde_json::{Value as Json, json}; use uuid::Uuid; -use nemo_relay_ffi::callable::{NemoRelayLlmExecNextFn, NemoRelayToolExecNextFn}; +use nemo_relay_ffi::callable::{ + NemoRelayLlmExecNextFn, NemoRelayLlmSanitizeCodecKind, NemoRelayLlmSanitizeRequestContext, + NemoRelayLlmSanitizeResponseContext, NemoRelayToolExecNextFn, +}; use nemo_relay_ffi::convert::nemo_relay_string_free; use nemo_relay_ffi::error::{NemoRelayStatus, nemo_relay_last_error, set_last_error}; use nemo_relay_ffi::types::{ @@ -343,6 +346,7 @@ unsafe extern "C" fn tool_exec_intercept_cb( unsafe extern "C" fn llm_request_cb( _user_data: *mut libc::c_void, request: *const FfiLLMRequest, + _context: NemoRelayLlmSanitizeRequestContext, ) -> *mut FfiLLMRequest { let request = unsafe { &*request }; let mut content = request.0.content.clone(); @@ -356,6 +360,7 @@ unsafe extern "C" fn llm_request_cb( unsafe extern "C" fn llm_response_cb( _user_data: *mut libc::c_void, response_json: *const c_char, + _context: NemoRelayLlmSanitizeResponseContext, ) -> *mut c_char { let mut response: Json = serde_json::from_str( unsafe { CStr::from_ptr(response_json) } @@ -388,7 +393,17 @@ unsafe extern "C" fn llm_request_intercept_cb( _annotated_json: *const c_char, out_outcome_json: *mut *mut c_char, ) -> NemoRelayStatus { - let transformed = unsafe { Box::from_raw(llm_request_cb(ptr::null_mut(), request)) }; + let transformed = unsafe { + Box::from_raw(llm_request_cb( + ptr::null_mut(), + request, + NemoRelayLlmSanitizeRequestContext { + codec_kind: NemoRelayLlmSanitizeCodecKind::None, + codec_id: ptr::null(), + codec: ptr::null(), + }, + )) + }; let outcome = json!({ "request": transformed.0, "annotated_request": null, diff --git a/crates/ffi/tests/integration/callable_extra_tests.rs b/crates/ffi/tests/integration/callable_extra_tests.rs index ca7eaf50d..b176f4988 100644 --- a/crates/ffi/tests/integration/callable_extra_tests.rs +++ b/crates/ffi/tests/integration/callable_extra_tests.rs @@ -69,10 +69,52 @@ unsafe extern "C" fn llm_request_intercept_invalid_annotated_cb( unsafe extern "C" fn llm_request_passthrough_cb( _user_data: *mut libc::c_void, request: *const FfiLLMRequest, + _context: NemoRelayLlmSanitizeRequestContext, ) -> *mut FfiLLMRequest { Box::into_raw(Box::new(FfiLLMRequest(unsafe { (&*request).0.clone() }))) } +unsafe extern "C" fn llm_request_codec_round_trip_cb( + _user_data: *mut libc::c_void, + request: *const FfiLLMRequest, + context: NemoRelayLlmSanitizeRequestContext, +) -> *mut FfiLLMRequest { + assert_eq!(context.codec_kind, NemoRelayLlmSanitizeCodecKind::BuiltIn); + assert_eq!( + unsafe { CStr::from_ptr(context.codec_id) } + .to_str() + .unwrap(), + "openai_chat" + ); + let annotated = unsafe { nemo_relay_llm_sanitize_request_codec_decode(context.codec, request) }; + assert!(!annotated.is_null()); + let normalized: Json = + serde_json::from_str(unsafe { CStr::from_ptr(annotated) }.to_str().unwrap()).unwrap(); + assert_eq!(normalized["model"], json!("gpt-test")); + let encoded = + unsafe { nemo_relay_llm_sanitize_request_codec_encode(context.codec, annotated, request) }; + unsafe { nemo_relay_string_free_internal(annotated) }; + encoded +} + +unsafe extern "C" fn llm_response_codec_decode_cb( + _user_data: *mut libc::c_void, + response: *const c_char, + context: NemoRelayLlmSanitizeResponseContext, +) -> *mut c_char { + assert_eq!(context.codec_kind, NemoRelayLlmSanitizeCodecKind::BuiltIn); + let annotated = + unsafe { nemo_relay_llm_sanitize_response_codec_decode(context.codec, response) }; + assert!(!annotated.is_null()); + let normalized: Json = + serde_json::from_str(unsafe { CStr::from_ptr(annotated) }.to_str().unwrap()).unwrap(); + assert_eq!(normalized["model"], json!("gpt-test")); + unsafe { nemo_relay_string_free_internal(annotated) }; + CString::new(unsafe { CStr::from_ptr(response) }.to_bytes()) + .unwrap() + .into_raw() +} + unsafe extern "C" fn llm_conditional_error_cb( _user_data: *mut libc::c_void, _request: *const FfiLLMRequest, @@ -227,7 +269,11 @@ fn test_callable_extra_request_intercept_and_codec_paths() { ); let sanitize = wrap_llm_sanitize_request_fn(llm_request_passthrough_cb, ptr::null_mut(), None); - let sanitized = sanitize(request.clone()); + let sanitized = sanitize( + request.clone(), + nemo_relay::api::runtime::LlmSanitizeRequestContext::default(), + ) + .expect("non-null sanitizer result"); assert_eq!(sanitized.content, request.content); let conditional = wrap_llm_conditional_fn(llm_conditional_error_cb, ptr::null_mut(), None); @@ -298,3 +344,43 @@ fn test_callable_extra_request_intercept_and_codec_paths() { let encode_err = invalid_encode.encode(&annotated, &request).unwrap_err(); assert!(encode_err.to_string().contains("invalid result JSON")); } + +#[test] +fn test_sanitizer_context_resolves_directional_ffi_codecs() { + use nemo_relay::api::runtime::{LlmSanitizeRequestContext, LlmSanitizeResponseContext}; + use nemo_relay::codec::openai_chat::OpenAIChatCodec; + + let codec = Arc::new(OpenAIChatCodec); + let request = LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "gpt-test", + "messages": [{"role": "user", "content": "secret"}], + "preserve": true + }), + }; + let sanitized = + wrap_llm_sanitize_request_fn(llm_request_codec_round_trip_cb, ptr::null_mut(), None)( + request.clone(), + LlmSanitizeRequestContext::for_request_codec(Some(codec.clone())), + ) + .expect("codec round trip returns a request"); + assert_eq!(sanitized.content, request.content); + + let response = json!({ + "id": "chatcmpl-test", + "model": "gpt-test", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "secret"}, + "finish_reason": "stop" + }] + }); + let sanitized = + wrap_llm_sanitize_response_fn(llm_response_codec_decode_cb, ptr::null_mut(), None)( + response.clone(), + LlmSanitizeResponseContext::for_response_codec(Some(codec)), + ) + .expect("codec decode returns a response"); + assert_eq!(sanitized, response); +} diff --git a/crates/ffi/tests/unit/api_tests.rs b/crates/ffi/tests/unit/api_tests.rs index 4c8c19c7b..b7661aae1 100644 --- a/crates/ffi/tests/unit/api_tests.rs +++ b/crates/ffi/tests/unit/api_tests.rs @@ -6,31 +6,37 @@ use super::*; use std::ffi::{CStr, CString}; use std::ptr; -use std::sync::{Mutex, OnceLock}; +use std::sync::{Arc, Mutex, OnceLock}; +use nemo_relay::codec::request::AnnotatedLlmRequest; +use nemo_relay::codec::response::AnnotatedLlmResponse; +use nemo_relay::codec::traits::{LlmCodec, LlmResponseCodec}; use nemo_relay::plugin::PluginRegistrationContext; use serde_json::{Value as Json, json}; use uuid::Uuid; -use crate::callable::{NemoRelayLlmExecNextFn, NemoRelayToolExecNextFn}; +use crate::callable::{ + NemoRelayLlmExecNextFn, NemoRelayLlmSanitizeCodecKind, NemoRelayLlmSanitizeRequestContext, + NemoRelayLlmSanitizeResponseContext, NemoRelayToolExecNextFn, +}; use crate::convert::nemo_relay_string_free; use crate::error::{NemoRelayStatus, nemo_relay_last_error}; use crate::types::{ - FfiAtifExporter, FfiEvent, FfiLLMHandle, FfiLLMRequest, FfiOpenTelemetrySubscriber, - FfiPluginActivation, FfiScopeStack, FfiToolHandle, nemo_relay_atif_exporter_free, - nemo_relay_event_data, nemo_relay_event_input, nemo_relay_event_metadata, - nemo_relay_event_model_name, nemo_relay_event_name, nemo_relay_event_output, - nemo_relay_event_parent_uuid, nemo_relay_event_scope_type, nemo_relay_event_timestamp, - nemo_relay_event_tool_call_id, nemo_relay_event_uuid, nemo_relay_llm_handle_attributes, - nemo_relay_llm_handle_free, nemo_relay_llm_handle_name, nemo_relay_llm_handle_parent_uuid, - nemo_relay_llm_handle_uuid, nemo_relay_llm_request_content, nemo_relay_llm_request_free, - nemo_relay_llm_request_headers, nemo_relay_llm_request_new, nemo_relay_otel_subscriber_free, - nemo_relay_scope_handle_attributes, nemo_relay_scope_handle_data, nemo_relay_scope_handle_free, - nemo_relay_scope_handle_metadata, nemo_relay_scope_handle_name, - nemo_relay_scope_handle_parent_uuid, nemo_relay_scope_handle_scope_type, - nemo_relay_scope_handle_uuid, nemo_relay_scope_stack_free, nemo_relay_tool_handle_attributes, - nemo_relay_tool_handle_free, nemo_relay_tool_handle_name, nemo_relay_tool_handle_parent_uuid, - nemo_relay_tool_handle_uuid, + FfiAtifExporter, FfiEvent, FfiLLMHandle, FfiLLMRequest, FfiLlmSanitizeRequestCodec, + FfiLlmSanitizeResponseCodec, FfiOpenTelemetrySubscriber, FfiPluginActivation, FfiScopeStack, + FfiToolHandle, nemo_relay_atif_exporter_free, nemo_relay_event_data, nemo_relay_event_input, + nemo_relay_event_metadata, nemo_relay_event_model_name, nemo_relay_event_name, + nemo_relay_event_output, nemo_relay_event_parent_uuid, nemo_relay_event_scope_type, + nemo_relay_event_timestamp, nemo_relay_event_tool_call_id, nemo_relay_event_uuid, + nemo_relay_llm_handle_attributes, nemo_relay_llm_handle_free, nemo_relay_llm_handle_name, + nemo_relay_llm_handle_parent_uuid, nemo_relay_llm_handle_uuid, nemo_relay_llm_request_content, + nemo_relay_llm_request_free, nemo_relay_llm_request_headers, nemo_relay_llm_request_new, + nemo_relay_otel_subscriber_free, nemo_relay_scope_handle_attributes, + nemo_relay_scope_handle_data, nemo_relay_scope_handle_free, nemo_relay_scope_handle_metadata, + nemo_relay_scope_handle_name, nemo_relay_scope_handle_parent_uuid, + nemo_relay_scope_handle_scope_type, nemo_relay_scope_handle_uuid, nemo_relay_scope_stack_free, + nemo_relay_tool_handle_attributes, nemo_relay_tool_handle_free, nemo_relay_tool_handle_name, + nemo_relay_tool_handle_parent_uuid, nemo_relay_tool_handle_uuid, }; use crate::{api, callable, types}; @@ -340,6 +346,7 @@ unsafe extern "C" fn tool_exec_intercept_cb( unsafe extern "C" fn llm_request_cb( _user_data: *mut libc::c_void, request: *const FfiLLMRequest, + _context: NemoRelayLlmSanitizeRequestContext, ) -> *mut FfiLLMRequest { let request = unsafe { &*request }; let mut content = request.0.content.clone(); @@ -353,6 +360,7 @@ unsafe extern "C" fn llm_request_cb( unsafe extern "C" fn llm_response_cb( _user_data: *mut libc::c_void, response_json: *const c_char, + _context: NemoRelayLlmSanitizeResponseContext, ) -> *mut c_char { let mut response: Json = serde_json::from_str( unsafe { CStr::from_ptr(response_json) } @@ -385,7 +393,17 @@ unsafe extern "C" fn llm_request_intercept_cb( _annotated_json: *const c_char, out_outcome_json: *mut *mut c_char, ) -> NemoRelayStatus { - let transformed = unsafe { Box::from_raw(llm_request_cb(ptr::null_mut(), request)) }; + let transformed = unsafe { + Box::from_raw(llm_request_cb( + ptr::null_mut(), + request, + NemoRelayLlmSanitizeRequestContext { + codec_kind: NemoRelayLlmSanitizeCodecKind::None, + codec_id: ptr::null(), + codec: ptr::null(), + }, + )) + }; let outcome = json!({ "request": transformed.0, "annotated_request": null, @@ -628,6 +646,86 @@ unsafe extern "C" fn plugin_register_fail_with_last_error( NemoRelayStatus::Internal } +struct PanickingFfiSanitizerCodec; + +impl LlmCodec for PanickingFfiSanitizerCodec { + fn decode( + &self, + _request: &nemo_relay::api::llm::LlmRequest, + ) -> nemo_relay::error::Result { + panic!("request decode panic") + } + + fn encode( + &self, + _annotated: &AnnotatedLlmRequest, + _original: &nemo_relay::api::llm::LlmRequest, + ) -> nemo_relay::error::Result { + panic!("request encode panic") + } +} + +impl LlmResponseCodec for PanickingFfiSanitizerCodec { + fn decode_response(&self, _response: &Json) -> nemo_relay::error::Result { + panic!("response decode panic") + } +} + +#[test] +fn ffi_sanitizer_codec_entrypoints_contain_codec_panics() { + let request_codec = + FfiLlmSanitizeRequestCodec(Arc::new(PanickingFfiSanitizerCodec) as Arc); + let response_codec = FfiLlmSanitizeResponseCodec( + Arc::new(PanickingFfiSanitizerCodec) as Arc + ); + let request = FfiLLMRequest(nemo_relay::api::llm::LlmRequest { + headers: serde_json::Map::new(), + content: json!({"model": "gpt-test", "messages": []}), + }); + let annotated_json = cstring(&serde_json::to_string(&AnnotatedLlmRequest::default()).unwrap()); + let response_json = cstring("{}"); + + let decoded = unsafe { + api::nemo_relay_llm_sanitize_request_codec_decode( + ptr::from_ref(&request_codec), + ptr::from_ref(&request), + ) + }; + assert!(decoded.is_null()); + assert!( + unsafe { read_last_error() } + .unwrap() + .contains("sanitizer request codec decode panicked") + ); + + let encoded = unsafe { + api::nemo_relay_llm_sanitize_request_codec_encode( + ptr::from_ref(&request_codec), + annotated_json.as_ptr(), + ptr::from_ref(&request), + ) + }; + assert!(encoded.is_null()); + assert!( + unsafe { read_last_error() } + .unwrap() + .contains("sanitizer request codec encode panicked") + ); + + let decoded = unsafe { + api::nemo_relay_llm_sanitize_response_codec_decode( + ptr::from_ref(&response_codec), + response_json.as_ptr(), + ) + }; + assert!(decoded.is_null()); + assert!( + unsafe { read_last_error() } + .unwrap() + .contains("sanitizer response codec decode panicked") + ); +} + #[path = "api/core_tests.rs"] mod core_tests; #[path = "api/coverage_sweeps_tests.rs"] diff --git a/crates/ffi/tests/unit/callable_tests.rs b/crates/ffi/tests/unit/callable_tests.rs index c127dd132..c58098b4a 100644 --- a/crates/ffi/tests/unit/callable_tests.rs +++ b/crates/ffi/tests/unit/callable_tests.rs @@ -152,10 +152,19 @@ unsafe extern "C" fn llm_request_intercept_cb( unsafe extern "C" fn llm_request_null_cb( _user_data: *mut libc::c_void, _request: *const FfiLLMRequest, + _context: NemoRelayLlmSanitizeRequestContext, ) -> *mut FfiLLMRequest { std::ptr::null_mut() } +unsafe extern "C" fn llm_request_alias_cb( + _user_data: *mut libc::c_void, + request: *const FfiLLMRequest, + _context: NemoRelayLlmSanitizeRequestContext, +) -> *mut FfiLLMRequest { + request.cast_mut() +} + unsafe extern "C" fn llm_conditional_cb( _user_data: *mut libc::c_void, request: *const FfiLLMRequest, @@ -167,13 +176,41 @@ unsafe extern "C" fn llm_conditional_cb( } } -unsafe extern "C" fn json_cb(_user_data: *mut libc::c_void, json: *const c_char) -> *mut c_char { +unsafe extern "C" fn json_cb( + _user_data: *mut libc::c_void, + json: *const c_char, + _context: NemoRelayLlmSanitizeResponseContext, +) -> *mut c_char { let mut value: Json = serde_json::from_str(unsafe { CStr::from_ptr(json) }.to_str().unwrap()).unwrap(); value["wrapped"] = json!(true); CString::new(value.to_string()).unwrap().into_raw() } +unsafe extern "C" fn json_alias_cb( + _user_data: *mut libc::c_void, + json: *const c_char, + _context: NemoRelayLlmSanitizeResponseContext, +) -> *mut c_char { + json.cast_mut() +} + +unsafe extern "C" fn invalid_json_cb( + _user_data: *mut libc::c_void, + _json: *const c_char, + _context: NemoRelayLlmSanitizeResponseContext, +) -> *mut c_char { + CString::new("not-json").unwrap().into_raw() +} + +unsafe extern "C" fn invalid_utf8_cb( + _user_data: *mut libc::c_void, + _json: *const c_char, + _context: NemoRelayLlmSanitizeResponseContext, +) -> *mut c_char { + CString::new([0xff]).unwrap().into_raw() +} + unsafe extern "C" fn llm_exec_cb( _user_data: *mut libc::c_void, native_json: *const c_char, @@ -379,9 +416,23 @@ fn test_wrap_llm_request_response_and_conditional_callbacks() { let sanitize_request = wrap_llm_sanitize_request_fn(llm_request_null_cb, std::ptr::null_mut(), None); - let sanitized = sanitize_request(make_request()); - assert_eq!(sanitized.headers.len(), 0); - assert_eq!(sanitized.content, Json::Null); + assert_eq!( + sanitize_request( + make_request(), + nemo_relay::api::runtime::LlmSanitizeRequestContext::default(), + ), + None + ); + + let alias_request = + wrap_llm_sanitize_request_fn(llm_request_alias_cb, std::ptr::null_mut(), None); + assert_eq!( + alias_request( + make_request(), + nemo_relay::api::runtime::LlmSanitizeRequestContext::default(), + ), + Some(make_request()) + ); let conditional = wrap_llm_conditional_fn(llm_conditional_cb, std::ptr::null_mut(), None); assert_eq!( @@ -394,14 +445,74 @@ fn test_wrap_llm_request_response_and_conditional_callbacks() { ); assert_eq!(conditional(&make_request()).unwrap(), None); - let wrapped_json = wrap_json_fn(json_cb, std::ptr::null_mut(), None); - assert_eq!(wrapped_json(json!({"value": 1}))["wrapped"], json!(true)); - - let wrapped_response = wrap_llm_response_fn(json_cb, std::ptr::null_mut(), None); + let wrapped_response = wrap_llm_sanitize_response_fn(json_cb, std::ptr::null_mut(), None); assert_eq!( - wrapped_response(json!({"value": 2}))["wrapped"], + wrapped_response( + json!({"value": 2}), + nemo_relay::api::runtime::LlmSanitizeResponseContext::default(), + ) + .unwrap()["wrapped"], json!(true) ); + + let alias_response = wrap_llm_sanitize_response_fn(json_alias_cb, std::ptr::null_mut(), None); + assert_eq!( + alias_response( + json!({"value": 2}), + nemo_relay::api::runtime::LlmSanitizeResponseContext::default(), + ), + Some(json!({"value": 2})) + ); + + for callback in [invalid_json_cb, invalid_utf8_cb] { + let malformed_response = + wrap_llm_sanitize_response_fn(callback, std::ptr::null_mut(), None); + assert_eq!( + malformed_response( + json!({"secret": "must be omitted"}), + nemo_relay::api::runtime::LlmSanitizeResponseContext::default(), + ), + None + ); + } +} + +#[test] +fn test_llm_sanitizers_fail_closed_for_runtime_codec_ids_with_embedded_nul() { + let runtime_identity = + nemo_relay::api::runtime::LlmCodecIdentity::Runtime("runtime\0codec".to_string()); + + let request_sanitizer = + wrap_llm_sanitize_request_fn(llm_request_alias_cb, std::ptr::null_mut(), None); + assert_eq!( + request_sanitizer( + make_request(), + nemo_relay::api::runtime::LlmSanitizeRequestContext::with_identity( + runtime_identity.clone(), + ), + ), + None + ); + assert!( + last_error_message() + .unwrap() + .contains("runtime codec ID contains an embedded NUL") + ); + + let response_sanitizer = + wrap_llm_sanitize_response_fn(json_alias_cb, std::ptr::null_mut(), None); + assert_eq!( + response_sanitizer( + json!({"secret": "must be omitted"}), + nemo_relay::api::runtime::LlmSanitizeResponseContext::with_identity(runtime_identity), + ), + None + ); + assert!( + last_error_message() + .unwrap() + .contains("runtime codec ID contains an embedded NUL") + ); } #[test] diff --git a/crates/node/package.json b/crates/node/package.json index 9cfa7278c..6f2f371a4 100644 --- a/crates/node/package.json +++ b/crates/node/package.json @@ -62,8 +62,8 @@ "triples": {} }, "scripts": { - "build": "napi build --platform --release", - "build-debug": "napi build --platform", + "build": "napi build --platform --release --dts index.d.ts", + "build-debug": "napi build --platform --dts index.d.ts", "check:docstrings": "node ../../scripts/lint/check_public_docstrings.mjs node", "check:docstrings:all": "node ../../scripts/lint/check_public_docstrings.mjs all", "format": "prettier --write \"*.{js,ts,d.ts}\" \"tests/**/*.{mjs,cjs,js,mts,cts,ts,tsx}\"", diff --git a/crates/node/plugin.d.ts b/crates/node/plugin.d.ts index 608289762..158e25b2f 100644 --- a/crates/node/plugin.d.ts +++ b/crates/node/plugin.d.ts @@ -4,6 +4,28 @@ /// import type { EventSanitizeFields, Json } from './index'; +import type { LlmCodec, LlmResponseCodec } from './typed'; + +/** Codec identity available while a managed LLM event is sanitized. */ +export type LlmCodecIdentity = + | { kind: 'none' } + | { kind: 'builtin'; id: 'openai_chat' | 'openai_responses' | 'anthropic_messages' } + | { kind: 'runtime'; id: string } + | { kind: 'opaque' }; + +/** Codec context available while an LLM request is sanitized. */ +export interface LlmSanitizeRequestContext { + codec: LlmCodecIdentity; + /** Resolve the active codec for this callback. Do not retain the result after the callback returns. */ + resolveCodec(): LlmCodec | null; +} + +/** Codec context available while an LLM response is sanitized. */ +export interface LlmSanitizeResponseContext { + codec: LlmCodecIdentity; + /** Resolve the active codec for this callback. Do not retain the result after the callback returns. */ + resolveCodec(): LlmResponseCodec | null; +} /** Policy behavior for unsupported configuration. */ export type UnsupportedBehavior = 'ignore' | 'warn' | 'error'; @@ -205,10 +227,18 @@ export interface PluginContext { priority: number, callback: (name: string, args: Json) => string | null, ): void; - /** Register an LLM sanitize-request guardrail for this component. */ - registerLlmSanitizeRequestGuardrail(name: string, priority: number, callback: (request: Json) => Json): void; - /** Register an LLM sanitize-response guardrail for this component. */ - registerLlmSanitizeResponseGuardrail(name: string, priority: number, callback: (response: Json) => Json): void; + /** Register an LLM sanitize-request guardrail. The callback receives `(request, context)`. */ + registerLlmSanitizeRequestGuardrail( + name: string, + priority: number, + callback: (request: Json, context: LlmSanitizeRequestContext) => Json | null, + ): void; + /** Register an LLM sanitize-response guardrail. The callback receives `(response, context)`. */ + registerLlmSanitizeResponseGuardrail( + name: string, + priority: number, + callback: (response: Json, context: LlmSanitizeResponseContext) => Json | null, + ): void; /** Register an LLM conditional-execution guardrail for this component. */ registerLlmConditionalExecutionGuardrail( name: string, diff --git a/crates/node/src/api/mod.rs b/crates/node/src/api/mod.rs index 12cc22b96..1f8384fae 100644 --- a/crates/node/src/api/mod.rs +++ b/crates/node/src/api/mod.rs @@ -599,6 +599,61 @@ fn middleware_json_callback_tsfn( Ok(tsfn) } +fn middleware_llm_sanitize_request_callback_tsfn( + env: &Env, + func: &JsFunction, +) -> napi::Result< + ThreadsafeFunction<(Json, callable::JsLlmSanitizeRequestContext), ErrorStrategy::Fatal>, +> { + let callback = callable::safe_middleware_callback(env, func)?; + let mut tsfn = callback.create_threadsafe_function( + 0, + |ctx: napi::threadsafe_function::ThreadSafeCallContext<( + Json, + callable::JsLlmSanitizeRequestContext, + )>| { + let first = unsafe { + JsUnknown::from_raw_unchecked( + ctx.env.raw(), + Json::to_napi_value(ctx.env.raw(), ctx.value.0)?, + ) + }; + let context = callable::js_llm_sanitize_request_context_to_napi(&ctx.env, ctx.value.1)?; + Ok(vec![first, context]) + }, + )?; + tsfn.unref(env)?; + Ok(tsfn) +} + +fn middleware_llm_sanitize_response_callback_tsfn( + env: &Env, + func: &JsFunction, +) -> napi::Result< + ThreadsafeFunction<(Json, callable::JsLlmSanitizeResponseContext), ErrorStrategy::Fatal>, +> { + let callback = callable::safe_middleware_callback(env, func)?; + let mut tsfn = callback.create_threadsafe_function( + 0, + |ctx: napi::threadsafe_function::ThreadSafeCallContext<( + Json, + callable::JsLlmSanitizeResponseContext, + )>| { + let first = unsafe { + JsUnknown::from_raw_unchecked( + ctx.env.raw(), + Json::to_napi_value(ctx.env.raw(), ctx.value.0)?, + ) + }; + let context = + callable::js_llm_sanitize_response_context_to_napi(&ctx.env, ctx.value.1)?; + Ok(vec![first, context]) + }, + )?; + tsfn.unref(env)?; + Ok(tsfn) +} + #[allow(clippy::too_many_arguments)] fn add_plugin_event_sanitizer( env: &Env, @@ -854,9 +909,9 @@ fn build_plugin_context( core_registry_api::register_llm_sanitize_request_guardrail( &name, priority, - callable::wrap_js_llm_sanitize_request_fn(middleware_json_callback_tsfn( - ctx.env, &callback, - )?), + callable::wrap_js_llm_sanitize_request_fn( + middleware_llm_sanitize_request_callback_tsfn(ctx.env, &callback)?, + ), ) .map_err(to_napi_err)?; @@ -900,9 +955,9 @@ fn build_plugin_context( core_registry_api::register_llm_sanitize_response_guardrail( &name, priority, - callable::wrap_js_llm_response_fn(middleware_json_callback_tsfn( - ctx.env, &callback, - )?), + callable::wrap_js_llm_sanitize_response_fn( + middleware_llm_sanitize_response_callback_tsfn(ctx.env, &callback)?, + ), ) .map_err(to_napi_err)?; @@ -1204,19 +1259,42 @@ struct NodePluginRegisterCall { /// struct. `reference` must be a valid N-API reference created for a live /// JavaScript function in `env`, and `env` must not be used after the /// corresponding Node.js environment has been torn down. -struct PersistentJsFunction { +pub(crate) struct PersistentJsFunction { env: napi::sys::napi_env, reference: napi::sys::napi_ref, + cleanup: napi::sys::napi_threadsafe_function, } -// SAFETY: `PersistentJsFunction` only stores raw N-API handles. Callers are -// responsible for constructing it from a live environment and function -// reference, and all access goes back through that same environment. +// SAFETY: Direct function access is restricted by callers to the registration +// thread. Releasing the cleanup TSFN is thread-safe, and its event-loop +// finalizer deletes the N-API reference. unsafe impl Send for PersistentJsFunction {} -// SAFETY: The same invariants as `Send` apply. The struct does not provide -// interior mutation beyond the N-API reference lifecycle managed by Node. +// SAFETY: The same invariants as `Send` apply. The stored handles are immutable, +// and reference deletion is serialized by the cleanup TSFN finalizer. unsafe impl Sync for PersistentJsFunction {} +unsafe extern "C" fn delete_persistent_js_function_reference( + env: napi::sys::napi_env, + finalize_data: *mut std::ffi::c_void, + _finalize_hint: *mut std::ffi::c_void, +) { + if !env.is_null() && !finalize_data.is_null() { + // SAFETY: `finalize_data` is the live N-API reference passed to the + // cleanup TSFN at creation. This finalizer runs once on Node's event + // loop after the final TSFN release. + let _ = + unsafe { napi::sys::napi_delete_reference(env, finalize_data as napi::sys::napi_ref) }; + } +} + +unsafe extern "C" fn persistent_js_function_cleanup_call( + _env: napi::sys::napi_env, + _js_callback: napi::sys::napi_value, + _context: *mut std::ffi::c_void, + _data: *mut std::ffi::c_void, +) { +} + impl PersistentJsFunction { fn new(env: &Env, func: &JsFunction) -> napi::Result { let mut reference = ptr::null_mut(); @@ -1225,17 +1303,74 @@ impl PersistentJsFunction { // writable storage for the created reference. let status = unsafe { napi::sys::napi_create_reference(env.raw(), func.raw(), 1, &mut reference) }; - if status == napi::sys::Status::napi_ok { - Ok(Self { - env: env.raw(), - reference, - }) - } else { - Err(napi::Error::from_reason(format!( + if status != napi::sys::Status::napi_ok { + return Err(napi::Error::from_reason(format!( "failed to create JS function reference: {:?}", napi::Status::from(status) - ))) + ))); + } + + let mut resource_name = ptr::null_mut(); + let resource_name_bytes = b"nemo_relay_persistent_js_function\0"; + let status = unsafe { + napi::sys::napi_create_string_utf8( + env.raw(), + resource_name_bytes.as_ptr().cast(), + resource_name_bytes.len() - 1, + &mut resource_name, + ) + }; + if status != napi::sys::Status::napi_ok { + let _ = unsafe { napi::sys::napi_delete_reference(env.raw(), reference) }; + return Err(napi::Error::from_reason(format!( + "failed to create persistent JS function cleanup resource name: {:?}", + napi::Status::from(status) + ))); + } + + let mut cleanup = ptr::null_mut(); + let status = unsafe { + napi::sys::napi_create_threadsafe_function( + env.raw(), + ptr::null_mut(), + ptr::null_mut(), + resource_name, + 0, + 1, + reference.cast(), + Some(delete_persistent_js_function_reference), + ptr::null_mut(), + Some(persistent_js_function_cleanup_call), + &mut cleanup, + ) + }; + if status != napi::sys::Status::napi_ok { + let _ = unsafe { napi::sys::napi_delete_reference(env.raw(), reference) }; + return Err(napi::Error::from_reason(format!( + "failed to create persistent JS function cleanup handle: {:?}", + napi::Status::from(status) + ))); } + + let status = unsafe { napi::sys::napi_unref_threadsafe_function(env.raw(), cleanup) }; + if status != napi::sys::Status::napi_ok { + let _ = unsafe { + napi::sys::napi_release_threadsafe_function( + cleanup, + napi::sys::ThreadsafeFunctionReleaseMode::release, + ) + }; + return Err(napi::Error::from_reason(format!( + "failed to unref persistent JS function cleanup handle: {:?}", + napi::Status::from(status) + ))); + } + + Ok(Self { + env: env.raw(), + reference, + cleanup, + }) } fn call_validate(&self, plugin_config: &Json) -> napi::Result { @@ -1303,6 +1438,29 @@ impl PersistentJsFunction { // SAFETY: `returned` is the live result of invoking `func` in this environment. unsafe { Option::::from_napi_value(self.env, returned.raw()) }.map(callback_json) } + + fn call_json(&self, argument: Json) -> napi::Result { + let mut value = ptr::null_mut(); + // SAFETY: `self.reference` is a live N-API reference created in + // `self.env`, and `value` is writable storage for the borrowed + // function value. + let status = + unsafe { napi::sys::napi_get_reference_value(self.env, self.reference, &mut value) }; + if status != napi::sys::Status::napi_ok { + return Err(napi::Error::from_reason("failed to borrow codec function")); + } + // SAFETY: `value` was resolved from this struct's function reference, + // so it is a live function value in `self.env` for this call. + let func = unsafe { JsFunction::from_raw_unchecked(self.env, value) }; + // SAFETY: `Json::to_napi_value` created this argument in `self.env`, + // so wrapping it as `JsUnknown` is valid for the immediate callback. + let argument = unsafe { + JsUnknown::from_raw_unchecked(self.env, Json::to_napi_value(self.env, argument)?) + }; + let returned = func.call(None, &[argument])?; + // SAFETY: `returned` is the live result of invoking `func` in this environment. + unsafe { Option::::from_napi_value(self.env, returned.raw()) }.map(callback_json) + } } fn core_event_fields( @@ -1390,11 +1548,91 @@ fn node_event_sanitize_fn(env: &Env, func: &JsFunction) -> napi::Result, + Vec>, +); +type NodeLlmResponseCodec = ( + Arc, + Vec>, +); + +fn node_llm_codec( + env: &Env, + decode: &JsFunction, + encode: &JsFunction, +) -> napi::Result { + let direct_decode = Arc::new(PersistentJsFunction::new(env, decode)?); + let direct_encode = Arc::new(PersistentJsFunction::new(env, encode)?); + let references = vec![direct_decode.clone(), direct_encode.clone()]; + let register_thread = std::thread::current().id(); + + let mut decode_tsfn = decode.create_threadsafe_function( + 0, + |ctx: napi::threadsafe_function::ThreadSafeCallContext| Ok(vec![ctx.value]), + )?; + decode_tsfn.unref(env)?; + let mut encode_tsfn = encode.create_threadsafe_function( + 0, + |ctx: napi::threadsafe_function::ThreadSafeCallContext| Ok(vec![ctx.value]), + )?; + encode_tsfn.unref(env)?; + + Ok(( + callable::wrap_js_codec( + decode_tsfn, + encode_tsfn, + register_thread, + Arc::new(move |argument| { + direct_decode.call_json(argument).map_err(|error| { + FlowError::Internal(format!("JS codec decode callback failed: {error}")) + }) + }), + Arc::new(move |argument| { + direct_encode.call_json(argument).map_err(|error| { + FlowError::Internal(format!("JS codec encode callback failed: {error}")) + }) + }), + ), + references, + )) +} + +fn node_llm_response_codec(env: &Env, decode: &JsFunction) -> napi::Result { + let direct_decode = Arc::new(PersistentJsFunction::new(env, decode)?); + let references = vec![direct_decode.clone()]; + let register_thread = std::thread::current().id(); + let mut decode_tsfn = decode.create_threadsafe_function( + 0, + |ctx: napi::threadsafe_function::ThreadSafeCallContext| Ok(vec![ctx.value]), + )?; + decode_tsfn.unref(env)?; + Ok(( + callable::wrap_js_response_codec( + decode_tsfn, + register_thread, + Arc::new(move |argument| { + direct_decode.call_json(argument).map_err(|error| { + FlowError::Internal(format!( + "JS response codec decode callback failed: {error}" + )) + }) + }), + ), + references, + )) +} + impl Drop for PersistentJsFunction { fn drop(&mut self) { - // SAFETY: `self.reference` was created by `napi_create_reference` for - // `self.env` and is deleted exactly once here during drop. - let _ = unsafe { napi::sys::napi_delete_reference(self.env, self.reference) }; + // SAFETY: N-API permits releasing a TSFN from any thread. Its finalizer + // runs on the event loop and deletes `self.reference` exactly once. + let _ = unsafe { + napi::sys::napi_release_threadsafe_function( + self.cleanup, + napi::sys::ThreadsafeFunctionReleaseMode::release, + ) + }; } } @@ -1523,8 +1761,8 @@ pub fn scope_stack_active() -> bool { /// Returns the most recent callback error that could not be surfaced through a direct exception. /// -/// This is primarily used for sanitize callback paths that fail open and cannot -/// surface their errors directly. +/// This is primarily used for sanitize callback paths that omit observability +/// payloads and cannot surface their errors directly. #[napi] pub fn get_last_callback_error() -> Option { get_recorded_callback_error() @@ -1549,30 +1787,32 @@ pub fn test_closed_tool_callback( wrapped(&name, args) } -/// Internal test helper: invoke a closed JS LLM sanitize-request wrapper and return the fallback request. +/// Internal test helper: model a closed JS LLM request sanitizer. #[napi(js_name = "__testClosedLlmSanitizeRequestCallback")] pub fn test_closed_llm_sanitize_request_callback( callback: ThreadsafeFunction, request: Json, -) -> Result { +) -> Result> { clear_recorded_callback_error(); let _ = callback.clone().abort(); let llm_request: LlmRequest = serde_json::from_value(request) .map_err(|e| napi::Error::from_reason(format!("invalid LlmRequest: {e}")))?; - let wrapped = callable::wrap_js_llm_sanitize_request_fn(callback); - Ok(serde_json::to_value(wrapped(llm_request)).unwrap_or(Json::Null)) + drop(llm_request); + record_callback_error("nemo_relay: failed to queue JS LLM sanitize request callback"); + Ok(None) } -/// Internal test helper: invoke a closed JS LLM sanitize-response wrapper and return the fallback response. +/// Internal test helper: model a closed JS LLM response sanitizer. #[napi(js_name = "__testClosedLlmResponseCallback")] pub fn test_closed_llm_response_callback( callback: ThreadsafeFunction, response: Json, -) -> Json { +) -> Option { clear_recorded_callback_error(); let _ = callback.clone().abort(); - let wrapped = callable::wrap_js_llm_response_fn(callback); - wrapped(response) + drop(response); + record_callback_error("nemo_relay: failed to queue JS LLM sanitize response callback"); + None } /// Internal test helper: invoke a closed JS collector wrapper and surface the queue failure. @@ -2123,9 +2363,9 @@ pub fn llm_call_execute( data: Option, metadata: Option, model_name: Option, - codec_decode: Option>, - codec_encode: Option>, - response_codec_decode: Option>, + #[napi(ts_arg_type = "(arg: Json) => any")] codec_decode: Option, + #[napi(ts_arg_type = "(arg: Json) => any")] codec_encode: Option, + #[napi(ts_arg_type = "(arg: Json) => any")] response_codec_decode: Option, ) -> Result { let attrs = LlmAttributes::from_bits_truncate(attributes.unwrap_or(0)); let parent = handle @@ -2136,11 +2376,28 @@ pub fn llm_call_execute( let callback = callable::safe_execution_callback(&env, &func)?; let exec_fn = callable::wrap_js_llm_exec_fn(json_callback_tsfn(&env, &callback)?); let default_fn: LlmExecutionNextFn = std::sync::Arc::new(move |req| exec_fn(req)); - let codec = match (codec_decode, codec_encode) { - (Some(d), Some(e)) => Some(callable::wrap_js_codec(d, e)), - _ => None, + let mut codec_references = Vec::new(); + let codec = match (codec_decode.as_ref(), codec_encode.as_ref()) { + (Some(d), Some(e)) => { + let (codec, references) = node_llm_codec(&env, d, e)?; + codec_references.extend(references); + Some(codec) + } + (None, None) => None, + _ => { + return Err(napi::Error::from_reason( + "codecDecode and codecEncode must be provided together", + )); + } }; - let response_codec = response_codec_decode.map(callable::wrap_js_response_codec); + let response_codec = response_codec_decode + .as_ref() + .map(|decode| node_llm_response_codec(&env, decode)) + .transpose()? + .map(|(codec, references)| { + codec_references.extend(references); + codec + }); let scope_stack = current_scope_stack_handle(); env.execute_tokio_future( @@ -2165,7 +2422,10 @@ pub fn llm_call_execute( }) .await }, - |_env, result| Ok(result), + move |_env, result| { + drop(codec_references); + Ok(result) + }, ) } @@ -2185,9 +2445,9 @@ pub fn llm_call_execute_async( data: Option, metadata: Option, model_name: Option, - codec_decode: Option>, - codec_encode: Option>, - response_codec_decode: Option>, + #[napi(ts_arg_type = "(arg: Json) => any")] codec_decode: Option, + #[napi(ts_arg_type = "(arg: Json) => any")] codec_encode: Option, + #[napi(ts_arg_type = "(arg: Json) => any")] response_codec_decode: Option, ) -> Result { let attrs = LlmAttributes::from_bits_truncate(attributes.unwrap_or(0)); let parent = handle @@ -2209,11 +2469,28 @@ pub fn llm_call_execute_async( Box::pin(async move { pa_fn.call(req_json).await }) }); - let codec = match (codec_decode, codec_encode) { - (Some(d), Some(e)) => Some(callable::wrap_js_codec(d, e)), - _ => None, + let mut codec_references = Vec::new(); + let codec = match (codec_decode.as_ref(), codec_encode.as_ref()) { + (Some(d), Some(e)) => { + let (codec, references) = node_llm_codec(&env, d, e)?; + codec_references.extend(references); + Some(codec) + } + (None, None) => None, + _ => { + return Err(napi::Error::from_reason( + "codecDecode and codecEncode must be provided together", + )); + } }; - let response_codec = response_codec_decode.map(callable::wrap_js_response_codec); + let response_codec = response_codec_decode + .as_ref() + .map(|decode| node_llm_response_codec(&env, decode)) + .transpose()? + .map(|(codec, references)| { + codec_references.extend(references); + codec + }); env.execute_tokio_future( async move { @@ -2237,7 +2514,10 @@ pub fn llm_call_execute_async( }) .await }, - |_env, result| Ok(result), + move |_env, result| { + drop(codec_references); + Ok(result) + }, ) } @@ -2272,9 +2552,9 @@ pub fn llm_stream_call_execute( data: Option, metadata: Option, model_name: Option, - codec_decode: Option>, - codec_encode: Option>, - response_codec_decode: Option>, + #[napi(ts_arg_type = "(arg: Json) => any")] codec_decode: Option, + #[napi(ts_arg_type = "(arg: Json) => any")] codec_encode: Option, + #[napi(ts_arg_type = "(arg: Json) => any")] response_codec_decode: Option, ) -> Result { let attrs = LlmAttributes::from_bits_truncate(attributes.unwrap_or(0)); let parent = handle @@ -2325,11 +2605,29 @@ pub fn llm_stream_call_execute( }) }); - let codec = match (codec_decode, codec_encode) { - (Some(d), Some(e)) => Some(callable::wrap_js_codec(d, e)), - _ => None, + let mut codec_references = Vec::new(); + let codec = match (codec_decode.as_ref(), codec_encode.as_ref()) { + (Some(d), Some(e)) => { + let (codec, references) = node_llm_codec(&env, d, e)?; + codec_references.extend(references); + Some(codec) + } + (None, None) => None, + _ => { + return Err(napi::Error::from_reason( + "codecDecode and codecEncode must be provided together", + )); + } }; - let response_codec = response_codec_decode.map(callable::wrap_js_response_codec); + let response_codec = response_codec_decode + .as_ref() + .map(|decode| node_llm_response_codec(&env, decode)) + .transpose()? + .map(|(codec, references)| { + codec_references.extend(references); + codec + }); + let completion_codec_references = codec_references.clone(); let scope_stack = current_scope_stack_handle(); env.execute_tokio_future( @@ -2368,11 +2666,15 @@ pub fn llm_stream_call_execute( receiver: tokio::sync::Mutex::new(rx), cancel, closed: closed_rx, + codec_references, }) }) .await }, - |_env, result| Ok(result), + move |_env, result| { + drop(completion_codec_references); + Ok(result) + }, ) } @@ -2610,18 +2912,21 @@ pub fn deregister_tool_execution_intercept(name: String) -> Result { /// Register a guardrail that sanitizes LLM request data before execution. /// -/// The `guardrail` callback receives the LLM request as JSON and must return the sanitized request. -/// Higher `priority` values run first. Throws if a guardrail with the same `name` already exists. -/// If the callback throws, Relay preserves the current emitted payload and records the error for -/// `getLastCallbackError()`. +/// The `guardrail` callback receives `(request, context)` and must return the sanitized request, +/// or `null` to omit the observability payload. Lower `priority` values run first. Throws if a +/// guardrail with the same `name` already exists. If the callback throws, Relay omits the payload +/// and records the error for `getLastCallbackError()`. #[napi] pub fn register_llm_sanitize_request_guardrail( env: Env, name: String, priority: i32, + #[napi( + ts_arg_type = "(request: Json, context: import('./plugin').LlmSanitizeRequestContext) => Json | null" + )] guardrail: JsFunction, ) -> Result<()> { - let callback = middleware_json_callback_tsfn(&env, &guardrail)?; + let callback = middleware_llm_sanitize_request_callback_tsfn(&env, &guardrail)?; core_registry_api::register_llm_sanitize_request_guardrail( &name, priority, @@ -2640,23 +2945,25 @@ pub fn deregister_llm_sanitize_request_guardrail(name: String) -> Result { /// Register a guardrail that sanitizes LLM response data after execution. /// -/// The `guardrail` callback receives the LLM response as a JSON value and must return -/// the sanitized response as JSON. Higher `priority` values run first. Throws if a guardrail -/// with the same `name` already exists. -/// If the callback throws, Relay preserves the current emitted payload and records the error for -/// `getLastCallbackError()`. +/// The `guardrail` callback receives `(response, context)` and must return the sanitized response, +/// or `null` to omit the observability payload. Lower `priority` values run first. Throws if a +/// guardrail with the same `name` already exists. If the callback throws, Relay omits the payload +/// and records the error for `getLastCallbackError()`. #[napi] pub fn register_llm_sanitize_response_guardrail( env: Env, name: String, priority: i32, + #[napi( + ts_arg_type = "(response: Json, context: import('./plugin').LlmSanitizeResponseContext) => Json | null" + )] guardrail: JsFunction, ) -> Result<()> { - let callback = middleware_json_callback_tsfn(&env, &guardrail)?; + let callback = middleware_llm_sanitize_response_callback_tsfn(&env, &guardrail)?; core_registry_api::register_llm_sanitize_response_guardrail( &name, priority, - callable::wrap_js_llm_response_fn(callback), + callable::wrap_js_llm_sanitize_response_fn(callback), ) .map_err(to_napi_err) } @@ -3140,17 +3447,19 @@ pub fn scope_deregister_tool_execution_intercept(scope_uuid: String, name: Strin /// Register a scope-local guardrail that sanitizes LLM request data before execution. /// -/// The `guardrail` callback receives the LLM request as JSON and must return the sanitized request. -/// Higher `priority` values run first. Throws if a guardrail with the same `name` already exists -/// on the specified scope. -/// If the callback throws, Relay preserves the current emitted payload and records the error for -/// `getLastCallbackError()`. +/// The `guardrail` callback receives `(request, context)` and must return the sanitized request, +/// or `null` to omit the observability payload. Lower `priority` values run first. Throws if a +/// guardrail with the same `name` already exists on the specified scope. If the callback throws, +/// Relay omits the payload and records the error for `getLastCallbackError()`. #[napi] pub fn scope_register_llm_sanitize_request_guardrail( env: Env, scope_uuid: String, name: String, priority: i32, + #[napi( + ts_arg_type = "(request: Json, context: import('./plugin').LlmSanitizeRequestContext) => Json | null" + )] guardrail: JsFunction, ) -> Result<()> { let uuid = uuid::Uuid::parse_str(&scope_uuid) @@ -3159,7 +3468,9 @@ pub fn scope_register_llm_sanitize_request_guardrail( &uuid, &name, priority, - callable::wrap_js_llm_sanitize_request_fn(middleware_json_callback_tsfn(&env, &guardrail)?), + callable::wrap_js_llm_sanitize_request_fn(middleware_llm_sanitize_request_callback_tsfn( + &env, &guardrail, + )?), ) .map_err(to_napi_err) } @@ -3180,17 +3491,19 @@ pub fn scope_deregister_llm_sanitize_request_guardrail( /// Register a scope-local guardrail that sanitizes LLM response data after execution. /// -/// The `guardrail` callback receives the LLM response as a JSON value and must return -/// the sanitized response as JSON. Higher `priority` values run first. Throws if a guardrail -/// with the same `name` already exists on the specified scope. -/// If the callback throws, Relay preserves the current emitted payload and records the error for -/// `getLastCallbackError()`. +/// The `guardrail` callback receives `(response, context)` and must return the sanitized response, +/// or `null` to omit the observability payload. Lower `priority` values run first. Throws if a +/// guardrail with the same `name` already exists on the specified scope. If the callback throws, +/// Relay omits the payload and records the error for `getLastCallbackError()`. #[napi] pub fn scope_register_llm_sanitize_response_guardrail( env: Env, scope_uuid: String, name: String, priority: i32, + #[napi( + ts_arg_type = "(response: Json, context: import('./plugin').LlmSanitizeResponseContext) => Json | null" + )] guardrail: JsFunction, ) -> Result<()> { let uuid = uuid::Uuid::parse_str(&scope_uuid) @@ -3199,7 +3512,9 @@ pub fn scope_register_llm_sanitize_response_guardrail( &uuid, &name, priority, - callable::wrap_js_llm_response_fn(middleware_json_callback_tsfn(&env, &guardrail)?), + callable::wrap_js_llm_sanitize_response_fn(middleware_llm_sanitize_response_callback_tsfn( + &env, &guardrail, + )?), ) .map_err(to_napi_err) } diff --git a/crates/node/src/callable.rs b/crates/node/src/callable.rs index 6937e8dd5..ba8abafc2 100644 --- a/crates/node/src/callable.rs +++ b/crates/node/src/callable.rs @@ -13,12 +13,15 @@ use std::future::Future; use std::pin::Pin; use std::sync::Arc; +use napi::bindgen_prelude::ToNapiValue; use napi::threadsafe_function::{ErrorStrategy, ThreadsafeFunction, ThreadsafeFunctionCallMode}; -use napi::{Env, JsFunction, JsUnknown, NapiRaw, NapiValue}; +use napi::{Env, JsFunction, JsObject, JsUnknown, NapiRaw, NapiValue}; +use napi_derive::napi; use nemo_relay::api::runtime::{ - EventSanitizeFn, EventSubscriberFn, LlmConditionalFn, LlmExecutionNextFn, LlmJsonStream, - LlmRequestInterceptFn, LlmSanitizeRequestFn, LlmSanitizeResponseFn, LlmStreamExecutionNextFn, - ToolConditionalFn, ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, + EventSanitizeFn, EventSubscriberFn, LlmCodecIdentity, LlmConditionalFn, LlmExecutionNextFn, + LlmJsonStream, LlmRequestInterceptFn, LlmSanitizeRequestContext, LlmSanitizeRequestFn, + LlmSanitizeResponseContext, LlmSanitizeResponseFn, LlmStreamExecutionNextFn, ToolConditionalFn, + ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, }; use serde::{Deserialize, Serialize}; use serde_json::Value as Json; @@ -37,10 +40,32 @@ use nemo_relay::codec::traits::{LlmCodec, LlmResponseCodec}; use nemo_relay::error::{FlowError, Result}; use crate::callback_factory; -use crate::convert::{callback_json, record_callback_error}; +use crate::convert::{callback_json, record_callback_error, to_napi_err}; use crate::promise_call::{JsonNextFn, JsonStreamNextFn, PromiseAwareFn}; use crate::types::{EventSanitizeFields, JsEvent, event_sanitize_fields_from_json}; +/// Structured codec identity delivered to JavaScript LLM sanitizers. +#[napi(object)] +#[derive(Clone)] +pub(crate) struct JsLlmCodecIdentity { + pub kind: String, + pub id: Option, +} + +/// Structured per-call request context delivered to JavaScript LLM sanitizers. +#[derive(Clone)] +pub(crate) struct JsLlmSanitizeRequestContext { + pub codec: JsLlmCodecIdentity, + resolved: Option>, +} + +/// Structured per-call response context delivered to JavaScript LLM sanitizers. +#[derive(Clone)] +pub(crate) struct JsLlmSanitizeResponseContext { + pub codec: JsLlmCodecIdentity, + resolved: Option>, +} + /// JavaScript-facing pending mark DTO. #[derive(Debug, Deserialize, Serialize)] #[serde(rename_all = "camelCase", deny_unknown_fields)] @@ -398,74 +423,227 @@ pub fn wrap_js_llm_request_intercept_fn( ) } -/// Wrap a JS function for LLM sanitize request: `(request: LlmRequest) => LlmRequest`. -/// Since ThreadsafeFunction requires serde-serializable args, we serialize the request as JSON. +/// Wrap a JS function for LLM request sanitization. The callback receives +/// `(request, context)`. pub fn wrap_js_llm_sanitize_request_fn( - func: ThreadsafeFunction, + func: ThreadsafeFunction<(Json, JsLlmSanitizeRequestContext), ErrorStrategy::Fatal>, ) -> LlmSanitizeRequestFn { let func = Arc::new(func); - Arc::new(move |request: LlmRequest| { - let func = func.clone(); - let req_json = serde_json::to_value(&request).unwrap_or(Json::Null); - let (tx, rx) = std::sync::mpsc::channel(); - let status = func.call_with_return_value( - req_json, - ThreadsafeFunctionCallMode::Blocking, - move |val: Option| { - let _ = tx.send(callback_json(val)); - Ok(()) + Arc::new( + move |request: LlmRequest, context: LlmSanitizeRequestContext| { + let context = js_llm_sanitize_request_context(&context); + let request = serde_json::to_value(request).unwrap_or(Json::Null); + let (tx, rx) = std::sync::mpsc::channel(); + if func.call_with_return_value( + (request.clone(), context), + ThreadsafeFunctionCallMode::Blocking, + move |value: Option| { + let _ = tx.send(callback_json(value)); + Ok(()) + }, + ) != napi::Status::Ok + { + record_callback_error( + "nemo_relay: failed to queue JS LLM sanitize request callback", + ); + return None; + } + let value = recv_middleware_json_or_value( + rx, + "nemo_relay: JS LLM request sanitizer callback failed", + Json::Null, + ); + if value.is_null() { + return None; + } + serde_json::from_value(value).map_or_else( + |error| { + record_callback_error(format!( + "nemo_relay: JS LLM sanitize request callback failed: failed to deserialize LlmRequest: {error}" + )); + None }, - ); - if status != napi::Status::Ok { - record_callback_error(format!( - "nemo_relay: failed to queue JS LLM sanitize request callback: {status:?}" - )); - return request; - } - // TODO: This closure returns LlmRequest (not Result), so we cannot propagate - // errors through the type system. Log the error so failures are not silent. - let result = recv_middleware_json_or_value( - rx, - "nemo_relay: JS LLM sanitize request callback failed", - serde_json::to_value(&request).unwrap_or(Json::Null), - ); - serde_json::from_value(result).unwrap_or_else(|error| { - record_callback_error(format!( - "nemo_relay: JS LLM sanitize request callback failed: failed to deserialize LlmRequest: {error}" - )); - request - }) - }) + Some, + ) + }, + ) } -/// Wrap a JS function for LLM sanitize response: `(response: Json) => Json`. -pub fn wrap_js_llm_response_fn( - func: ThreadsafeFunction, +/// Wrap a JS function for LLM response sanitization. The callback receives +/// `(response, context)`; returning `null` omits the event payload. +pub fn wrap_js_llm_sanitize_response_fn( + func: ThreadsafeFunction<(Json, JsLlmSanitizeResponseContext), ErrorStrategy::Fatal>, ) -> LlmSanitizeResponseFn { let func = Arc::new(func); - Arc::new(move |response: Json| { - let func = func.clone(); + Arc::new(move |response: Json, context: LlmSanitizeResponseContext| { + let context = js_llm_sanitize_response_context(&context); let (tx, rx) = std::sync::mpsc::channel(); - let status = func.call_with_return_value( - response.clone(), + if func.call_with_return_value( + (response, context), ThreadsafeFunctionCallMode::Blocking, - move |val: Option| { - let _ = tx.send(callback_json(val)); + move |value: Option| { + let _ = tx.send(callback_json(value)); Ok(()) }, - ); - if status != napi::Status::Ok { - record_callback_error(format!( - "nemo_relay: failed to queue JS LLM response callback: {status:?}" - )); - return response; + ) != napi::Status::Ok + { + record_callback_error("nemo_relay: failed to queue JS LLM sanitize response callback"); + return None; } - // TODO: This closure returns Json (not Result), so we cannot propagate - // errors through the type system. Log the error and fall back to original response. - recv_middleware_json_or_value(rx, "nemo_relay: JS LLM response callback failed", response) + let value = recv_middleware_json_or_value( + rx, + "nemo_relay: JS LLM response sanitizer callback failed", + Json::Null, + ); + Some(value).and_then(|value| (!value.is_null()).then_some(value)) }) } +fn js_llm_sanitize_request_context( + context: &LlmSanitizeRequestContext, +) -> JsLlmSanitizeRequestContext { + JsLlmSanitizeRequestContext { + codec: js_codec_identity(context.codec()), + resolved: context.resolve_codec(), + } +} + +fn js_codec_identity(identity: &LlmCodecIdentity) -> JsLlmCodecIdentity { + match identity { + LlmCodecIdentity::None => JsLlmCodecIdentity { + kind: "none".into(), + id: None, + }, + LlmCodecIdentity::BuiltIn(codec) => JsLlmCodecIdentity { + kind: "builtin".into(), + id: Some(codec.id().into()), + }, + LlmCodecIdentity::Runtime(id) => JsLlmCodecIdentity { + kind: "runtime".into(), + id: Some(id.clone()), + }, + LlmCodecIdentity::Opaque => JsLlmCodecIdentity { + kind: "opaque".into(), + id: None, + }, + } +} + +fn js_llm_sanitize_response_context( + context: &LlmSanitizeResponseContext, +) -> JsLlmSanitizeResponseContext { + JsLlmSanitizeResponseContext { + codec: js_codec_identity(context.codec()), + resolved: context.resolve_codec(), + } +} + +fn js_object_to_unknown(env: &Env, object: JsObject) -> JsUnknown { + unsafe { JsUnknown::from_raw_unchecked(env.raw(), object.raw()) } +} + +fn request_codec_object(env: &Env, codec: Arc) -> napi::Result { + let mut object = env.create_object()?; + let decode_codec = codec.clone(); + let decode = env.create_function_from_closure("decode", move |ctx| { + let request = ctx.get::(0)?; + let request = serde_json::from_value(request) + .map_err(|error| napi::Error::from_reason(format!("invalid LlmRequest: {error}")))?; + serde_json::to_value(decode_codec.decode(&request).map_err(to_napi_err)?) + .map_err(|error| napi::Error::from_reason(error.to_string())) + })?; + let encode_codec = codec; + let encode = env.create_function_from_closure("encode", move |ctx| { + let annotated = ctx.get::(0)?; + let original = ctx.get::(1)?; + let annotated = serde_json::from_value(annotated).map_err(|error| { + napi::Error::from_reason(format!("invalid AnnotatedLlmRequest: {error}")) + })?; + let original = serde_json::from_value(original) + .map_err(|error| napi::Error::from_reason(format!("invalid LlmRequest: {error}")))?; + serde_json::to_value( + encode_codec + .encode(&annotated, &original) + .map_err(to_napi_err)?, + ) + .map_err(|error| napi::Error::from_reason(error.to_string())) + })?; + object.set_named_property("decode", decode)?; + object.set_named_property("encode", encode)?; + Ok(object) +} + +fn response_codec_object(env: &Env, codec: Arc) -> napi::Result { + let mut object = env.create_object()?; + let decode = env.create_function_from_closure("decodeResponse", move |ctx| { + let response = ctx.get::(0)?; + serde_json::to_value(codec.decode_response(&response).map_err(to_napi_err)?) + .map_err(|error| napi::Error::from_reason(error.to_string())) + })?; + object.set_named_property("decodeResponse", decode)?; + Ok(object) +} + +/// Convert a request sanitizer context into the JavaScript object passed to a callback. +pub(crate) fn js_llm_sanitize_request_context_to_napi( + env: &Env, + context: JsLlmSanitizeRequestContext, +) -> napi::Result { + let mut object = env.create_object()?; + let codec = unsafe { + JsUnknown::from_raw_unchecked( + env.raw(), + JsLlmCodecIdentity::to_napi_value(env.raw(), context.codec)?, + ) + }; + object.set_named_property("codec", codec)?; + + let resolved = context.resolved; + let resolve_codec = + env.create_function_from_closure("resolveCodec", move |ctx| match resolved.clone() { + Some(codec) => Ok(js_object_to_unknown( + ctx.env, + request_codec_object(ctx.env, codec)?, + )), + None => ctx + .env + .get_null() + .map(|value| unsafe { JsUnknown::from_raw_unchecked(ctx.env.raw(), value.raw()) }), + })?; + object.set_named_property("resolveCodec", resolve_codec)?; + Ok(js_object_to_unknown(env, object)) +} + +/// Convert a response sanitizer context into the JavaScript object passed to a callback. +pub(crate) fn js_llm_sanitize_response_context_to_napi( + env: &Env, + context: JsLlmSanitizeResponseContext, +) -> napi::Result { + let mut object = env.create_object()?; + let codec = unsafe { + JsUnknown::from_raw_unchecked( + env.raw(), + JsLlmCodecIdentity::to_napi_value(env.raw(), context.codec)?, + ) + }; + object.set_named_property("codec", codec)?; + + let resolved = context.resolved; + let resolve_codec = + env.create_function_from_closure("resolveCodec", move |ctx| match resolved.clone() { + Some(codec) => Ok(js_object_to_unknown( + ctx.env, + response_codec_object(ctx.env, codec)?, + )), + None => ctx + .env + .get_null() + .map(|value| unsafe { JsUnknown::from_raw_unchecked(ctx.env.raw(), value.raw()) }), + })?; + object.set_named_property("resolveCodec", resolve_codec)?; + Ok(js_object_to_unknown(env, object)) +} + /// Wrap a JS function for LLM conditional guardrails: `(request: object) => string | null`. pub fn wrap_js_llm_conditional_fn( func: ThreadsafeFunction, @@ -683,11 +861,22 @@ pub fn wrap_js_event_sanitize_fn( struct NapiCodec { decode: Arc>, encode: Arc>, + register_thread: std::thread::ThreadId, + direct_decode: Arc Result + Send + Sync>, + direct_encode: Arc Result + Send + Sync>, } impl LlmCodec for NapiCodec { fn decode(&self, request: &LlmRequest) -> Result { let req_json = serde_json::to_value(request).unwrap_or(Json::Null); + if std::thread::current().id() == self.register_thread { + let result = (self.direct_decode)(req_json)?; + return serde_json::from_value(result).map_err(|e| { + FlowError::Internal(format!( + "JS codec decode callback: failed to deserialize AnnotatedLlmRequest: {e}" + )) + }); + } let (tx, rx) = std::sync::mpsc::channel(); let status = self.decode.call_with_return_value( req_json, @@ -714,6 +903,13 @@ impl LlmCodec for NapiCodec { let annotated_json = serde_json::to_value(annotated).unwrap_or(Json::Null); let original_json = serde_json::to_value(original).unwrap_or(Json::Null); let arg = serde_json::json!({"annotated": annotated_json, "original": original_json}); + if std::thread::current().id() == self.register_thread { + return serde_json::from_value((self.direct_encode)(arg)?).map_err(|e| { + FlowError::Internal(format!( + "JS codec encode callback: failed to deserialize LlmRequest: {e}" + )) + }); + } let (tx, rx) = std::sync::mpsc::channel(); let status = self.encode.call_with_return_value( arg, @@ -737,10 +933,16 @@ impl LlmCodec for NapiCodec { pub fn wrap_js_codec( decode: ThreadsafeFunction, encode: ThreadsafeFunction, + register_thread: std::thread::ThreadId, + direct_decode: Arc Result + Send + Sync>, + direct_encode: Arc Result + Send + Sync>, ) -> Arc { Arc::new(NapiCodec { decode: Arc::new(decode), encode: Arc::new(encode), + register_thread, + direct_decode, + direct_encode, }) } @@ -752,10 +954,20 @@ pub fn wrap_js_codec( /// delegating `decode_response` to a JavaScript function via `ThreadsafeFunction`. struct NapiResponseCodec { decode_response: Arc>, + register_thread: std::thread::ThreadId, + direct_decode_response: Arc Result + Send + Sync>, } impl LlmResponseCodec for NapiResponseCodec { fn decode_response(&self, response: &Json) -> Result { + if std::thread::current().id() == self.register_thread { + let result = (self.direct_decode_response)(response.clone())?; + return serde_json::from_value(result).map_err(|e| { + FlowError::Internal(format!( + "decode_response returned invalid AnnotatedLlmResponse: {e}" + )) + }); + } let (tx, rx) = std::sync::mpsc::channel(); let status = self.decode_response.call_with_return_value( response.clone(), @@ -784,9 +996,13 @@ impl LlmResponseCodec for NapiResponseCodec { /// Wrap a JS decode_response function into an `Arc`. pub fn wrap_js_response_codec( decode_response: ThreadsafeFunction, + register_thread: std::thread::ThreadId, + direct_decode_response: Arc Result + Send + Sync>, ) -> Arc { Arc::new(NapiResponseCodec { decode_response: Arc::new(decode_response), + register_thread, + direct_decode_response, }) } diff --git a/crates/node/src/stream.rs b/crates/node/src/stream.rs index c9308d66a..37b1e43a3 100644 --- a/crates/node/src/stream.rs +++ b/crates/node/src/stream.rs @@ -11,6 +11,9 @@ use napi::bindgen_prelude::*; use napi_derive::napi; use nemo_relay::error::Result as FlowResult; use serde_json::Value as Json; +use std::sync::Arc; + +use crate::api::PersistentJsFunction; /// An async iterator over chunks from a streaming LLM response. /// @@ -21,6 +24,7 @@ pub struct LlmStream { pub(crate) receiver: tokio::sync::Mutex>>, pub(crate) cancel: tokio::sync::watch::Sender, pub(crate) closed: tokio::sync::watch::Receiver>>, + pub(crate) codec_references: Vec>, } #[napi] diff --git a/crates/node/tests/callback_error_tests.mjs b/crates/node/tests/callback_error_tests.mjs index 94130142b..83e684a5f 100644 --- a/crates/node/tests/callback_error_tests.mjs +++ b/crates/node/tests/callback_error_tests.mjs @@ -35,7 +35,7 @@ function makeNative() { describe('callback error helpers', () => { it('getLastCallbackError and clearLastCallbackError expose malformed sanitize-request failures', async () => { clearLastCallbackError(); - registerLlmSanitizeRequestGuardrail('node_llm_san_req_public_error', 10, () => null); + registerLlmSanitizeRequestGuardrail('node_llm_san_req_public_error', 10, () => ({ broken: true })); try { const result = await llmCallExecute( 'san_req_public_error_llm', @@ -83,7 +83,7 @@ describe('callback error helpers', () => { clearLastCallbackError(); }); - it('closed llm sanitize-request callbacks fall back to the original request and record the queue failure', () => { + it('closed llm sanitize-request callbacks omit the payload and record the queue failure', () => { const request = makeNative(); const result = __testClosedLlmSanitizeRequestCallback( () => ({ @@ -91,12 +91,12 @@ describe('callback error helpers', () => { }), request, ); - assert.deepEqual(result, request); + assert.equal(result, null); assert.match(getLastCallbackError() ?? '', /failed to queue JS LLM sanitize request callback/i); clearLastCallbackError(); }); - it('closed llm sanitize-response callbacks fall back to the original response and record the queue failure', () => { + it('closed llm sanitize-response callbacks omit the payload and record the queue failure', () => { const response = { ok: true, }; @@ -106,8 +106,8 @@ describe('callback error helpers', () => { }), response, ); - assert.deepEqual(result, response); - assert.match(getLastCallbackError() ?? '', /failed to queue JS LLM response callback/i); + assert.equal(result, null); + assert.match(getLastCallbackError() ?? '', /failed to queue JS LLM sanitize response callback/i); clearLastCallbackError(); }); diff --git a/crates/node/tests/event_sanitizers_tests.mjs b/crates/node/tests/event_sanitizers_tests.mjs index fdd67520d..812a2a0ef 100644 --- a/crates/node/tests/event_sanitizers_tests.mjs +++ b/crates/node/tests/event_sanitizers_tests.mjs @@ -4,6 +4,9 @@ import assert from 'node:assert/strict'; import { describe, it } from 'node:test'; import { createRequire } from 'node:module'; +import { mkdtempSync, rmSync } from 'node:fs'; +import { tmpdir } from 'node:os'; +import path from 'node:path'; const require = createRequire(import.meta.url); const lib = require('../index.js'); @@ -28,6 +31,18 @@ function assertSanitizerFieldsCleared(event) { assert.equal(event.metadata, null); } +async function initializeWithoutDiscoveredPluginConfig(config) { + const previousDirectory = process.cwd(); + const directory = mkdtempSync(path.join(tmpdir(), 'nemo-relay-node-')); + try { + process.chdir(directory); + return await plugin.initialize(config); + } finally { + process.chdir(previousDirectory); + rmSync(directory, { recursive: true, force: true }); + } +} + describe('event sanitizer registries', () => { it('orders mark sanitizers and supports field removal', async () => { const events = capture('node-event-sanitize-order-sub'); @@ -251,7 +266,10 @@ describe('event sanitizer registries', () => { }, }); try { - await plugin.initialize({ version: 1, components: [plugin.ComponentSpec(kind)] }); + await initializeWithoutDiscoveredPluginConfig({ + version: 1, + components: [plugin.ComponentSpec(kind)], + }); lib.event('configured', null, { raw: true }); lib.flushSubscribers(); await waitFor(events, 1); @@ -289,7 +307,10 @@ describe('event sanitizer registries', () => { }); lib.clearLastCallbackError(); try { - await plugin.initialize({ version: 1, components: [plugin.ComponentSpec(kind)] }); + await initializeWithoutDiscoveredPluginConfig({ + version: 1, + components: [plugin.ComponentSpec(kind)], + }); lib.event('plugin-throw', null, { raw: true }, { raw: true }); lib.flushSubscribers(); await waitFor(events, 1); diff --git a/crates/node/tests/llm_tests.mjs b/crates/node/tests/llm_tests.mjs index bb4e9e821..3115b2136 100644 --- a/crates/node/tests/llm_tests.mjs +++ b/crates/node/tests/llm_tests.mjs @@ -3,11 +3,15 @@ import { describe, it } from 'node:test'; import assert from 'node:assert/strict'; +import { execFileSync } from 'node:child_process'; import { createRequire } from 'node:module'; import { readFileSync } from 'node:fs'; +import path from 'node:path'; +import { fileURLToPath } from 'node:url'; const require = createRequire(import.meta.url); const lib = require('../index.js'); +const nodeDir = fileURLToPath(new URL('..', import.meta.url)); const { pushScope, @@ -356,6 +360,275 @@ describe('LLM execute', () => { // =========================================================================== describe('LLM guardrails', () => { + it('contextual sanitizers receive payload first and codec context second', async () => { + const events = []; + let requestContextChecked = false; + let responseContextChecked = false; + registerSubscriber('node_contextual_llm_sanitize_events', (event) => events.push(event)); + registerLlmSanitizeRequestGuardrail('node_contextual_llm_sanitize_request', 10, (request, context) => { + assert.deepEqual(context.codec, { kind: 'none' }); + assert.equal(context.resolveCodec(), null); + requestContextChecked = true; + return { + ...request, + headers: { ...request.headers, 'X-Contextual-Sanitized': 'request' }, + }; + }); + registerLlmSanitizeResponseGuardrail('node_contextual_llm_sanitize_response', 10, (response, context) => { + assert.deepEqual(context.codec, { kind: 'none' }); + assert.equal(context.resolveCodec(), null); + responseContextChecked = true; + return { ...response, contextualSanitized: true }; + }); + + try { + const result = await llmCallExecute('contextual_sanitize_llm', makeNative(), () => ({ ok: true })); + assert.deepEqual(result, { ok: true }); + assert.equal(requestContextChecked, true); + assert.equal(responseContextChecked, true); + await flushSubscriberCallbacks(); + const start = events.find( + (event) => event.name === 'contextual_sanitize_llm' && event.scope_category === 'start', + ); + const end = events.find((event) => event.name === 'contextual_sanitize_llm' && event.scope_category === 'end'); + assert.equal(start.data.headers['X-Contextual-Sanitized'], 'request'); + assert.equal(end.data.contextualSanitized, true); + } finally { + deregisterLlmSanitizeRequestGuardrail('node_contextual_llm_sanitize_request'); + deregisterLlmSanitizeResponseGuardrail('node_contextual_llm_sanitize_response'); + deregisterSubscriber('node_contextual_llm_sanitize_events'); + } + }); + + it('rejects incomplete request codec callback pairs', () => { + const decode = (request) => request; + const encode = ({ annotated }) => annotated; + + assert.throws( + () => llmCallExecute('partial_codec_execute', makeNative(), () => ({}), null, null, null, null, null, decode), + /codecDecode and codecEncode must be provided together/, + ); + assert.throws( + () => + llmCallExecuteAsync( + 'partial_codec_execute_async', + makeNative(), + async () => ({}), + null, + null, + null, + null, + null, + null, + encode, + ), + /codecDecode and codecEncode must be provided together/, + ); + assert.throws( + () => + llmStreamCallExecute( + 'partial_codec_stream', + makeNative(), + () => {}, + null, + null, + null, + null, + null, + null, + null, + decode, + ), + /codecDecode and codecEncode must be provided together/, + ); + }); + + it('resolved sanitizer codecs expose directional operations', async () => { + const codec = new lib.OpenAIChatCodec(); + let requestDecoded = false; + let responseDecoded = false; + registerLlmSanitizeRequestGuardrail('node_resolved_llm_request_codec', 10, (request, context) => { + assert.deepEqual(context.codec, { kind: 'opaque' }); + const resolved = context.resolveCodec(); + assert.notEqual(resolved, null); + const annotated = resolved.decode(request); + requestDecoded = annotated.model === 'test-model'; + return resolved.encode(annotated, request); + }); + registerLlmSanitizeResponseGuardrail('node_resolved_llm_response_codec', 10, (response, context) => { + assert.deepEqual(context.codec, { kind: 'opaque' }); + const resolved = context.resolveCodec(); + assert.notEqual(resolved, null); + const annotated = resolved.decodeResponse(response); + responseDecoded = annotated.model === 'test-model'; + return response; + }); + + try { + const response = { + id: 'chatcmpl-test', + model: 'test-model', + choices: [ + { + index: 0, + message: { role: 'assistant', content: 'ok' }, + finish_reason: 'stop', + }, + ], + }; + const result = await llmCallExecute( + 'resolved_sanitizer_codec_llm', + makeNative(), + () => response, + null, + null, + null, + null, + null, + codec.decode.bind(codec), + ({ annotated, original }) => codec.encode(annotated, original), + codec.decodeResponse.bind(codec), + ); + assert.deepEqual(result, response); + assert.equal(requestDecoded, true); + assert.equal(responseDecoded, true); + } finally { + deregisterLlmSanitizeRequestGuardrail('node_resolved_llm_request_codec'); + deregisterLlmSanitizeResponseGuardrail('node_resolved_llm_response_codec'); + } + }); + + it('streaming sanitizers resolve codecs only during callbacks and can omit the final payload', async () => { + const codec = new lib.OpenAIChatCodec(); + const events = []; + let requestCodecUsed = false; + let responseCodecUsed = false; + registerSubscriber('node_streaming_codec_sanitize_events', (event) => events.push(event)); + registerLlmSanitizeRequestGuardrail('node_streaming_request_codec', 10, (request, context) => { + const resolved = context.resolveCodec(); + assert.notEqual(resolved, null); + const annotated = resolved.decode(request); + requestCodecUsed = true; + return resolved.encode(annotated, request); + }); + registerLlmSanitizeResponseGuardrail('node_streaming_response_codec', 10, (response, context) => { + const resolved = context.resolveCodec(); + assert.notEqual(resolved, null); + resolved.decodeResponse(response); + responseCodecUsed = true; + return null; + }); + + const response = { + id: 'chatcmpl-stream', + model: 'test-model', + choices: [{ index: 0, message: { role: 'assistant', content: 'secret' }, finish_reason: 'stop' }], + }; + try { + const stream = await llmStreamCallExecute( + 'streaming_resolved_sanitizer_codec', + makeNative(), + (wrapper) => { + lib.pushStreamChunk(wrapper.__nemo_relay_stream_id, { delta: 'secret' }); + lib.endStream(wrapper.__nemo_relay_stream_id); + }, + null, + () => response, + null, + null, + null, + null, + null, + codec.decode.bind(codec), + ({ annotated, original }) => codec.encode(annotated, original), + codec.decodeResponse.bind(codec), + ); + assert.deepEqual(await stream.next(), { delta: 'secret' }); + assert.equal(await stream.next(), null); + await flushSubscriberCallbacks(); + assert.equal(requestCodecUsed, true); + assert.equal(responseCodecUsed, true); + const end = events.find( + (event) => event.name === 'streaming_resolved_sanitizer_codec' && event.scope_category === 'end', + ); + assert.equal(end.data, null); + assert.equal(end.category_profile.annotated_response, undefined); + } finally { + deregisterLlmSanitizeRequestGuardrail('node_streaming_request_codec'); + deregisterLlmSanitizeResponseGuardrail('node_streaming_response_codec'); + deregisterSubscriber('node_streaming_codec_sanitize_events'); + } + }); + + it('releases custom stream codec references safely after early garbage collection', () => { + const modulePath = path.join(nodeDir, 'index.js'); + const script = ` + import { createRequire } from 'node:module'; + const require = createRequire(${JSON.stringify(path.join(nodeDir, 'package.json'))}); + const lib = require(${JSON.stringify(modulePath)}); + const codec = new lib.OpenAIChatCodec(); + let ended = false; + lib.registerSubscriber('early_drop_custom_codec_stream_events', (event) => { + if ( + event.name === 'early_drop_custom_codec_stream' && + event.scope_category === 'end' + ) { + ended = true; + } + }); + let stream = await lib.llmStreamCallExecute( + 'early_drop_custom_codec_stream', + { + headers: {}, + content: { model: 'test-model', messages: [] }, + }, + (wrapper) => { + setTimeout(() => { + lib.endStream(wrapper.__nemo_relay_stream_id); + }, 25); + }, + null, + () => ({ model: 'test-model', choices: [] }), + null, + null, + null, + null, + null, + codec.decode.bind(codec), + ({ annotated, original }) => codec.encode(annotated, original), + codec.decodeResponse.bind(codec), + ); + const weak = new WeakRef(stream); + stream = null; + let collected = false; + for (let index = 0; index < 100; index += 1) { + global.gc(); + await new Promise((resolve) => setImmediate(resolve)); + if (weak.deref() === undefined) { + collected = true; + break; + } + await new Promise((resolve) => setImmediate(resolve)); + } + if (!collected) { + throw new Error('unfinished custom-codec stream was not garbage collected'); + } + for (let index = 0; index < 100 && !ended; index += 1) { + global.gc(); + await new Promise((resolve) => setImmediate(resolve)); + } + if (!ended) { + throw new Error('early-dropped custom-codec stream did not finish cleanup'); + } + lib.deregisterSubscriber('early_drop_custom_codec_stream_events'); + await new Promise((resolve) => setImmediate(resolve)); + `; + execFileSync(process.execPath, ['--expose-gc', '--input-type=module', '--eval', script], { + stdio: 'inherit', + timeout: 30_000, + }); + }); + it('sanitize request guardrail', () => { registerLlmSanitizeRequestGuardrail('node_llm_san_req', 10, (request) => { request.extra = 'sanitized'; @@ -420,7 +693,7 @@ describe('LLM guardrails', () => { } }); - it('sanitize request guardrail falls back on malformed return', async () => { + it('sanitize request guardrail can omit the observability payload', async () => { registerLlmSanitizeRequestGuardrail('node_llm_san_req_bad', 10, () => null); try { const result = await llmCallExecute( @@ -445,7 +718,7 @@ describe('LLM guardrails', () => { } }); - it('sanitize request guardrail failures preserve the payload and remain usable', async () => { + it('sanitize request guardrail failures omit the payload and remain usable', async () => { const events = []; clearLastCallbackError(); registerSubscriber('node_llm_san_req_throw_sub', (event) => events.push(event)); @@ -463,7 +736,7 @@ describe('LLM guardrails', () => { event.category === 'llm' && event.scope_category === 'start', ); - assert.deepEqual(start.data, request); + assert.equal(start.data, null); assert.match(getLastCallbackError() ?? '', /JavaScript callback threw/i); deregisterLlmSanitizeRequestGuardrail('node_llm_san_req_throw'); @@ -567,7 +840,7 @@ describe('LLM guardrails', () => { } }); - it('sanitize response guardrail failures preserve the payload and remain usable', async () => { + it('sanitize response guardrail failures omit the payload and remain usable', async () => { const events = []; clearLastCallbackError(); registerSubscriber('node_llm_san_resp_throw_sub', (event) => events.push(event)); @@ -585,7 +858,7 @@ describe('LLM guardrails', () => { event.category === 'llm' && event.scope_category === 'end', ); - assert.deepEqual(end.data, response); + assert.equal(end.data, null); assert.match(getLastCallbackError() ?? '', /response sanitizer boom/i); deregisterLlmSanitizeResponseGuardrail('node_llm_san_resp_throw'); @@ -1074,6 +1347,14 @@ describe('LLM intercepts', () => { assert.equal(declarations.split(openKind).length - 1, 3); }); + it('generated LLM sanitizer declarations expose directional codec contexts', () => { + const declarations = readFileSync(new URL('../index.d.ts', import.meta.url), 'utf8'); + + assert.equal(declarations.split("context: import('./plugin').LlmSanitizeRequestContext").length - 1, 2); + assert.equal(declarations.split("context: import('./plugin').LlmSanitizeResponseContext").length - 1, 2); + assert.doesNotMatch(declarations, /registerLlmSanitizeRequestGuardrail\([^\n]*\.\.\.args: any\[\]/); + }); + it('standalone conditional execution helper throws on rejection', async () => { registerLlmConditionalExecutionGuardrail('node_llm_cond_helper', 10, () => 'llm blocked by helper'); try { diff --git a/crates/pii-redaction/README.md b/crates/pii-redaction/README.md index 1b0f820da..4c8528582 100644 --- a/crates/pii-redaction/README.md +++ b/crates/pii-redaction/README.md @@ -85,7 +85,6 @@ sanitization surface: kind = "pii_redaction" [components.config] -codec = "openai_chat" [[components.config.profiles]] mode = "builtin" @@ -124,7 +123,6 @@ enabled = true [components.config] version = 1 -codec = "anthropic_messages" [[components.config.profiles]] mode = "builtin" diff --git a/crates/pii-redaction/src/builtin.rs b/crates/pii-redaction/src/builtin.rs index 43cc9cdb8..6e6d1dc5c 100644 --- a/crates/pii-redaction/src/builtin.rs +++ b/crates/pii-redaction/src/builtin.rs @@ -6,16 +6,19 @@ use std::sync::Arc; use regex::Regex; use serde::Serialize; use serde::de::DeserializeOwned; -use serde_json::Value as Json; +use serde_json::{Map, Value as Json}; use sha2::{Digest, Sha256}; use nemo_relay::api::event::{CategoryProfile, Event}; use nemo_relay::api::llm::LlmRequest; use nemo_relay::api::runtime::{ - EventSanitizeFn, LlmSanitizeRequestFn, LlmSanitizeResponseFn, ToolSanitizeFn, + BuiltinLlmCodec, EventSanitizeFn, LlmCodecIdentity, LlmSanitizeRequestFn, + LlmSanitizeResponseFn, ToolSanitizeFn, }; +use nemo_relay::codec::request::AnnotatedLlmRequest; use nemo_relay::codec::resolve::{ - ProviderSurface, request_codec as build_request_codec, response_codec as build_response_codec, + ProviderSurface, detect_response_surface, request_codec as build_request_codec, + response_codec as build_response_codec, }; use nemo_relay::codec::traits::{LlmCodec, LlmResponseCodec}; use nemo_relay::plugin::{PluginError, Result as PluginResult}; @@ -29,9 +32,7 @@ use super::trajectory::{CustomMarkPayloadPolicy, TrajectorySanitizer}; pub(super) struct CompiledBuiltinBackend { action: BuiltinAction, target_paths: Arc>, - request_codec: Option>, - response_codec: Option>, - codec_name: Option, + legacy_surface: Option, trajectory: Option, } @@ -168,9 +169,7 @@ impl CompiledBuiltinBackend { Ok(Self { action, target_paths: Arc::new(config.target_paths), - request_codec: surface.map(build_request_codec), - response_codec: surface.map(build_response_codec), - codec_name: surface.map(BuiltinCodecName::from_provider_surface), + legacy_surface: surface, trajectory, }) } @@ -286,19 +285,171 @@ impl CompiledBuiltinBackend { } } - fn sanitize_request_with_codec(&self, request: &LlmRequest) -> Option { - let codec = self.request_codec.as_ref()?; + fn selected_surface(&self, codec: &LlmCodecIdentity) -> Option { + match codec { + LlmCodecIdentity::None => self.legacy_surface, + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiChat) => { + Some(ProviderSurface::OpenAIChat) + } + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiResponses) => { + Some(ProviderSurface::OpenAIResponses) + } + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::AnthropicMessages) => { + Some(ProviderSurface::AnthropicMessages) + } + LlmCodecIdentity::Runtime(_) | LlmCodecIdentity::Opaque => None, + } + } + + fn uses_compatible_legacy_response_codec(&self, payload: &Json) -> bool { + self.legacy_surface + .is_some_and(|surface| detect_response_surface(payload) == Some(surface)) + } + + fn sanitize_request_with_codec( + &self, + codec: &dyn LlmCodec, + request: &LlmRequest, + ) -> Option { let annotated = codec.decode(request).ok()?; let sanitized_annotated = sanitize_serializable_with_backend(self, annotated).ok()?; - codec.encode(&sanitized_annotated, request).ok() + codec + .encode(&sanitized_annotated, request) + .ok() + .or_else(|| { + self.sanitize_request_target_paths_incrementally( + codec, + request, + sanitized_annotated, + ) + }) + } + + fn sanitize_request_target_paths_incrementally( + &self, + codec: &dyn LlmCodec, + request: &LlmRequest, + sanitized_annotated: AnnotatedLlmRequest, + ) -> Option { + let sanitized = serde_json::to_value(sanitized_annotated).ok()?; + let mut sanitized_request = request.clone(); + + for target_path in self.target_paths.iter() { + let target_segments = json_pointer_segments(target_path)?; + let current_annotated = codec.decode(&sanitized_request).ok()?; + let mut current = serde_json::to_value(¤t_annotated).ok()?; + match ( + sanitized_json_pointer_value(&sanitized, &target_segments), + sanitized_json_pointer_value(¤t, &target_segments), + ) { + (None, None) => continue, + (Some(target_value), Some(current_value)) if current_value == target_value => { + continue; + } + (Some(target_value), Some(_)) => { + replace_sanitized_json_pointer_value( + &mut current, + &target_segments, + target_value.clone(), + )?; + } + (None, Some(_)) if matches!(self.action, BuiltinAction::Remove) => { + remove_sanitized_json_pointer_value(&mut current, &target_segments)?; + } + _ => return None, + } + let updated = serde_json::from_value(current).ok()?; + sanitized_request = codec.encode(&updated, &sanitized_request).ok()?; + } + + Some(sanitized_request) + } + + fn sanitize_request_headers(&self, headers: Map) -> Map { + let sanitized = self.sanitize_json_preorder_dfs(Json::Object( + [("headers".to_string(), Json::Object(headers))] + .into_iter() + .collect(), + )); + sanitized + .get("headers") + .and_then(Json::as_object) + .cloned() + .unwrap_or_default() } - fn sanitize_response_with_codec(&self, payload: Json) -> Option { - let codec = self.response_codec.as_ref()?; - let codec_name = self.codec_name?; + fn sanitize_response_with_codec( + &self, + codec: &dyn LlmResponseCodec, + surface: ProviderSurface, + payload: Json, + ) -> Option { + if surface == ProviderSurface::OpenAIChat + && payload + .get("choices") + .and_then(Json::as_array) + .is_some_and(|choices| choices.len() > 1) + && self.targets_normalized_openai_chat_choice() + { + return None; + } + let codec_name = BuiltinCodecName::from_provider_surface(surface); let annotated = codec.decode_response(&payload).ok()?; + let annotated_json = serde_json::to_value(&annotated).ok()?; let sanitized_annotated = sanitize_serializable_with_backend(self, annotated).ok()?; - Some(codec_name.overlay_response_payload(payload, &sanitized_annotated)) + let sanitized_json = serde_json::to_value(&sanitized_annotated).ok()?; + let payload = codec_name.overlay_response_payload(payload, &sanitized_annotated); + let payload = self.sanitize_json_preorder_dfs(payload); + let target_segments = self + .target_paths + .iter() + .map(|target_path| json_pointer_segments(target_path)) + .collect::>>()?; + let has_normalized_target = target_segments.iter().any(|target_segments| { + sanitized_json_pointer_value(&annotated_json, target_segments).is_some() + || sanitized_json_pointer_value(&sanitized_json, target_segments).is_some() + }); + if !has_normalized_target { + return Some(payload); + } + let projected = codec.decode_response(&payload).ok()?; + let projected = serde_json::to_value(projected).ok()?; + Self::normalized_response_targets_match( + &target_segments, + &annotated_json, + &sanitized_json, + &projected, + ) + .then_some(payload) + } + + fn normalized_response_targets_match( + target_paths: &[Vec], + annotated: &Json, + sanitized: &Json, + projected: &Json, + ) -> bool { + target_paths.iter().all(|target_segments| { + let original = sanitized_json_pointer_value(annotated, target_segments); + let expected = sanitized_json_pointer_value(sanitized, target_segments); + if original.is_none() && expected.is_none() { + return true; + } + sanitized_json_pointer_value(projected, target_segments) == expected + }) + } + + fn targets_normalized_openai_chat_choice(&self) -> bool { + self.target_paths.iter().any(|path| { + json_pointer_segments(path) + .and_then(|segments| segments.into_iter().next()) + .is_some_and(|segment| { + matches!( + segment.as_str(), + "message" | "tool_calls" | "finish_reason" | "api_specific" + ) + }) + }) } } @@ -368,37 +519,183 @@ fn event_sanitize_callback_with_scope_categories( pub(super) fn llm_sanitize_request_callback( backend: CompiledBuiltinBackend, ) -> LlmSanitizeRequestFn { - Arc::new(move |mut request: LlmRequest| { + Arc::new(move |mut request: LlmRequest, context| { if let Some(trajectory) = backend.trajectory.as_ref() { + request.headers = trajectory + .sanitize_tool_payload(Json::Object(request.headers)) + .as_object() + .cloned() + .unwrap_or_default(); request.content = trajectory.sanitize_provider_payload(request.content); - return request; + return Some(request); } - if let Some(encoded) = backend.sanitize_request_with_codec(&request) { - return encoded; + request.headers = backend.sanitize_request_headers(request.headers); + if backend.target_paths.is_empty() { + request.content = backend.sanitize_json_preorder_dfs(request.content); + return Some(request); + } + let resolved = context.resolve_codec(); + let fallback = if resolved.is_none() { + backend + .selected_surface(context.codec()) + .map(build_request_codec) + } else { + None + }; + let Some(codec) = resolved.as_deref().or(fallback.as_deref()) else { + log_llm_payload_omitted("request", context.codec(), "no usable request codec"); + return None; + }; + let sanitized = backend.sanitize_request_with_codec(codec, &request); + if sanitized.is_none() { + log_llm_payload_omitted( + "request", + context.codec(), + "codec decode, sanitize, or encode failure", + ); } - request.content = backend.sanitize_json_preorder_dfs(request.content); - request + sanitized }) } pub(super) fn llm_sanitize_response_callback( backend: CompiledBuiltinBackend, ) -> LlmSanitizeResponseFn { - Arc::new(move |payload: Json| { + Arc::new(move |payload: Json, context| { if let Some(trajectory) = backend.trajectory.as_ref() { - return trajectory.sanitize_provider_payload(payload); + return Some(trajectory.sanitize_provider_payload(payload)); } if backend.target_paths.is_empty() { - return backend.sanitize_json_preorder_dfs(payload); + return Some(backend.sanitize_json_preorder_dfs(payload)); } - - let payload = backend - .sanitize_response_with_codec(payload.clone()) - .unwrap_or(payload); - backend.sanitize_json_preorder_dfs(payload) + if matches!(context.codec(), LlmCodecIdentity::None) + && !backend.uses_compatible_legacy_response_codec(&payload) + { + log_llm_payload_omitted( + "response", + context.codec(), + "no active response codec or compatible legacy codec", + ); + return None; + } + let Some(surface) = backend.selected_surface(context.codec()) else { + log_llm_payload_omitted( + "response", + context.codec(), + "no recognized response codec surface", + ); + return None; + }; + let resolved = context.resolve_codec(); + let fallback = if resolved.is_none() { + Some(build_response_codec(surface)) + } else { + None + }; + let Some(codec) = resolved.as_deref().or(fallback.as_deref()) else { + log_llm_payload_omitted("response", context.codec(), "no usable response codec"); + return None; + }; + let sanitized = backend.sanitize_response_with_codec(codec, surface, payload); + if sanitized.is_none() { + log_llm_payload_omitted( + "response", + context.codec(), + "codec decode, sanitize, or encode failure", + ); + } + sanitized }) } +fn log_llm_payload_omitted(direction: &str, codec: &LlmCodecIdentity, reason: &str) { + let codec_kind = match codec { + LlmCodecIdentity::None => "none", + LlmCodecIdentity::BuiltIn(_) => "builtin", + LlmCodecIdentity::Runtime(_) => "runtime", + LlmCodecIdentity::Opaque => "opaque", + }; + log::warn!( + target: "nemo_relay.plugin", + event = "pii_llm_payload_omitted", + codec_kind, + reason; + "PII redaction omitted an LLM {direction} payload" + ); +} + +fn json_pointer_segments(pointer: &str) -> Option> { + pointer + .strip_prefix('/') + .map(|path| path.split('/').map(unescape_json_pointer_segment).collect()) +} + +fn unescape_json_pointer_segment(segment: &str) -> String { + segment.replace("~1", "/").replace("~0", "~") +} + +fn sanitized_json_pointer_value<'a>(value: &'a Json, segments: &[String]) -> Option<&'a Json> { + segments + .iter() + .try_fold(value, |value, segment| match value { + Json::Object(values) => values.get(segment), + Json::Array(values) => segment + .parse::() + .ok() + .and_then(|index| values.get(index)), + _ => None, + }) +} + +fn replace_sanitized_json_pointer_value( + value: &mut Json, + segments: &[String], + replacement: Json, +) -> Option<()> { + let (last, parents) = segments.split_last()?; + let parent = parents + .iter() + .try_fold(value, |value, segment| match value { + Json::Object(values) => values.get_mut(segment), + Json::Array(values) => segment + .parse::() + .ok() + .and_then(|index| values.get_mut(index)), + _ => None, + })?; + match parent { + Json::Object(values) => { + values.insert(last.clone(), replacement); + Some(()) + } + Json::Array(values) => { + let index = last.parse::().ok()?; + let value = values.get_mut(index)?; + *value = replacement; + Some(()) + } + _ => None, + } +} + +fn remove_sanitized_json_pointer_value(value: &mut Json, segments: &[String]) -> Option<()> { + let (last, parents) = segments.split_last()?; + let parent = parents + .iter() + .try_fold(value, |value, segment| match value { + Json::Object(values) => values.get_mut(segment), + Json::Array(values) => segment + .parse::() + .ok() + .and_then(|index| values.get_mut(index)), + _ => None, + })?; + match parent { + Json::Object(values) => values.remove(last).map(|_| ()), + Json::Array(_) | Json::Null | Json::Bool(_) | Json::Number(_) | Json::String(_) => None, + } +} + fn render_json_pointer_path(path_segments: &[String]) -> String { if path_segments.is_empty() { return String::new(); diff --git a/crates/pii-redaction/src/component.rs b/crates/pii-redaction/src/component.rs index 0c4c3b415..2f6b9bfd3 100644 --- a/crates/pii-redaction/src/component.rs +++ b/crates/pii-redaction/src/component.rs @@ -97,7 +97,8 @@ pub struct PiiRedactionConfig { skip_serializing_if = "is_default_priority" )] pub priority: i32, - /// Provider request/response codec for LLM-managed surfaces. + /// Compatibility fallback codec for LLM-managed surfaces without an active + /// per-call codec. #[serde(default, skip_serializing_if = "Option::is_none")] #[cfg_attr(feature = "schema", schemars(schema_with = "codec_schema"))] pub codec: Option, @@ -1113,14 +1114,6 @@ fn validate_codec_requirements( } let Some(codec) = config.codec.as_deref() else { - push_policy_diag( - diagnostics, - policy.unsupported_value, - "pii_redaction.unsupported_value", - Some(PII_REDACTION_PLUGIN_KIND.to_string()), - Some("codec".to_string()), - "codec is required when any LLM surface is enabled".to_string(), - ); return; }; diff --git a/crates/pii-redaction/src/overlay.rs b/crates/pii-redaction/src/overlay.rs index 2cab7c35f..e74da6e26 100644 --- a/crates/pii-redaction/src/overlay.rs +++ b/crates/pii-redaction/src/overlay.rs @@ -88,8 +88,12 @@ fn overlay_openai_responses_response(mut payload: Json, annotated: &AnnotatedLlm .map(openai_responses_status), ); + let message_text = annotated_message_text(annotated.message.as_ref()); + if root.contains_key("output_text") { + set_optional_string_field(root, "output_text", message_text.as_deref()); + } if let Some(items) = root.get_mut("output").and_then(Json::as_array_mut) { - overlay_output_text_blocks(items, annotated_message_text(annotated.message.as_ref())); + overlay_output_text_blocks(items, message_text); overlay_openai_responses_tool_calls(items, annotated.tool_calls.as_deref()); } payload diff --git a/crates/pii-redaction/tests/unit/component_tests.rs b/crates/pii-redaction/tests/unit/component_tests.rs index 3ac54253e..cd7ad4026 100644 --- a/crates/pii-redaction/tests/unit/component_tests.rs +++ b/crates/pii-redaction/tests/unit/component_tests.rs @@ -14,8 +14,9 @@ use crate::api::llm::{ llm_call_execute, llm_stream_call_execute, }; use crate::api::runtime::{ - LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn, NemoRelayContextState, - create_scope_stack, global_context, set_thread_scope_stack, + BuiltinLlmCodec, LlmCodecIdentity, LlmExecutionNextFn, LlmJsonStream, + LlmSanitizeRequestContext, LlmSanitizeResponseContext, LlmStreamExecutionNextFn, + NemoRelayContextState, create_scope_stack, global_context, set_thread_scope_stack, }; use crate::api::scope::{ EmitMarkEventParams, PopScopeParams, PushScopeParams, ScopeType, event, pop_scope, push_scope, @@ -24,12 +25,13 @@ use crate::api::subscriber::{deregister_subscriber, register_subscriber}; use crate::api::tool::{ToolCallEndParams, ToolCallParams, tool_call, tool_call_end}; use crate::codec::openai_chat::OpenAIChatCodec; use crate::codec::openai_responses::OpenAIResponsesCodec; +use crate::codec::request::AnnotatedLlmRequest; use crate::codec::traits::{LlmCodec, LlmResponseCodec}; use crate::plugin::{ ConfigPolicy, DiagnosticLevel, PluginComponentSpec, PluginConfig, PluginError, PluginRegistrationContext, UnsupportedBehavior, clear_plugin_configuration, - ensure_builtin_plugins_registered, initialize_plugins, list_plugin_kinds, - rollback_registrations, validate_plugin_config, + ensure_builtin_plugins_registered, initialize_plugins_exact as initialize_plugins, + list_plugin_kinds, rollback_registrations, validate_plugin_config, }; use futures::StreamExt; use nemo_relay::observability::atif::{AtifAgentInfo, AtifExporter}; @@ -263,6 +265,253 @@ fn trajectory_backend(codec: Option<&str>, policy: &str) -> crate::builtin::Comp .unwrap() } +fn no_codec_context() -> LlmSanitizeResponseContext { + LlmSanitizeResponseContext::default() +} + +fn no_codec_request_context() -> LlmSanitizeRequestContext { + LlmSanitizeRequestContext::default() +} + +struct IdentifiedRequestCodec { + identity: LlmCodecIdentity, + inner: OpenAIResponsesCodec, +} + +impl LlmCodec for IdentifiedRequestCodec { + fn codec_identity(&self) -> LlmCodecIdentity { + self.identity.clone() + } + + fn decode(&self, request: &LlmRequest) -> nemo_relay::error::Result { + self.inner.decode(request) + } + + fn encode( + &self, + annotated: &AnnotatedLlmRequest, + original: &LlmRequest, + ) -> nemo_relay::error::Result { + self.inner.encode(annotated, original) + } +} + +#[test] +fn normalized_llm_paths_use_the_active_codec_and_fail_closed_for_unknown_codecs() { + let backend = crate::builtin::CompiledBuiltinBackend::new( + BuiltinBackendConfig { + action: "regex_replace".to_string(), + pattern: Some("sk-[A-Za-z0-9_-]+".to_string()), + replacement: Some("[REDACTED]".to_string()), + target_paths: vec![ + "/messages/0/content/0/text".to_string(), + "/message".to_string(), + ], + ..BuiltinBackendConfig::default() + }, + Some("openai_chat".to_string()), + ) + .unwrap(); + let sanitize_request = crate::builtin::llm_sanitize_request_callback(backend.clone()); + let sanitize_response = crate::builtin::llm_sanitize_response_callback(backend); + let active_request = sanitize_request( + LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "gpt-4.1-mini", + "input": [{"role": "user", "content": [{"type": "input_text", "text": "sk-request-secret"}]}] + }), + }, + LlmSanitizeRequestContext::for_request_codec(Some(Arc::new(OpenAIResponsesCodec))), + ) + .expect("the active OpenAI Responses codec must override the legacy fallback"); + assert_eq!( + active_request.content["input"][0]["content"][0]["text"], + json!("[REDACTED]") + ); + + for identity in [ + LlmCodecIdentity::Runtime("com.example.responses.v1".to_owned()), + LlmCodecIdentity::Opaque, + ] { + let active_request = sanitize_request( + LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "gpt-4.1-mini", + "input": [{ + "role": "user", + "content": [{"type": "input_text", "text": "sk-request-secret"}] + }] + }), + }, + LlmSanitizeRequestContext::for_request_codec(Some(Arc::new(IdentifiedRequestCodec { + identity, + inner: OpenAIResponsesCodec, + }))), + ) + .expect("an active runtime or opaque request codec must remain usable"); + assert_eq!( + active_request.content["input"][0]["content"][0]["text"], + json!("[REDACTED]") + ); + } + + let responses_payload = json!({ + "id": "resp_123", + "model": "gpt-4.1-mini", + "status": "completed", + "output": [{ + "type": "message", + "content": [{"type": "output_text", "text": "sk-responses-secret"}] + }] + }); + + let active_responses = sanitize_response( + responses_payload.clone(), + LlmSanitizeResponseContext::with_identity(LlmCodecIdentity::BuiltIn( + BuiltinLlmCodec::OpenAiResponses, + )), + ) + .expect("the active OpenAI Responses codec must override the legacy fallback"); + assert_eq!( + active_responses["output"][0]["content"][0]["text"], + json!("[REDACTED]") + ); + + assert!( + sanitize_response(responses_payload.clone(), no_codec_context()).is_none(), + "an incompatible configured fallback codec must omit a normalized payload" + ); + + assert!( + sanitize_response( + responses_payload, + LlmSanitizeResponseContext::with_identity(LlmCodecIdentity::Opaque), + ) + .is_none(), + "a normalized-path policy must omit an unknown active provider payload" + ); + + assert!( + sanitize_response( + json!({ + "id": "resp_123", + "output": [{"content": [{"type": "output_text", "text": "sk-runtime-secret"}]}] + }), + LlmSanitizeResponseContext::with_identity(LlmCodecIdentity::Runtime( + "com.example.chat.v1".to_owned(), + )), + ) + .is_none(), + "a normalized-path policy must omit a runtime codec until it has a compatible projection" + ); +} + +#[test] +fn normalized_llm_paths_omit_payloads_when_legacy_codec_decode_fails() { + let backend = crate::builtin::CompiledBuiltinBackend::new( + BuiltinBackendConfig { + action: "regex_replace".to_string(), + pattern: Some("sk-[A-Za-z0-9_-]+".to_string()), + replacement: Some("[REDACTED]".to_string()), + target_paths: vec!["/messages/0/content".to_string()], + ..BuiltinBackendConfig::default() + }, + Some("openai_chat".to_string()), + ) + .unwrap(); + let sanitize_request = crate::builtin::llm_sanitize_request_callback(backend.clone()); + let sanitize_response = crate::builtin::llm_sanitize_response_callback(backend); + + assert!( + sanitize_request( + LlmRequest { + headers: serde_json::Map::new(), + content: json!({"messages": "sk-request-secret"}), + }, + no_codec_request_context(), + ) + .is_none(), + "a shallow legacy surface match must not enable a raw-payload fallback" + ); + assert!( + sanitize_response(json!({"choices": "sk-response-secret"}), no_codec_context()).is_none(), + "a legacy response codec failure must omit the payload instead of emitting raw content" + ); +} + +#[test] +fn normalized_openai_chat_api_specific_policy_omits_multiple_choices() { + let backend = crate::builtin::CompiledBuiltinBackend::new( + BuiltinBackendConfig { + action: "remove".to_string(), + target_paths: vec!["/api_specific".to_string()], + ..BuiltinBackendConfig::default() + }, + None, + ) + .unwrap(); + let sanitize_response = crate::builtin::llm_sanitize_response_callback(backend); + + assert!( + sanitize_response( + json!({ + "id": "chatcmpl-multi", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "first"}, + "logprobs": {"content": [{"token": "SECRET-FIRST"}]} + }, + { + "index": 1, + "message": {"role": "assistant", "content": "second"}, + "logprobs": {"content": [{"token": "SECRET-SECOND"}]} + } + ] + }), + LlmSanitizeResponseContext::with_identity(LlmCodecIdentity::BuiltIn( + BuiltinLlmCodec::OpenAiChat, + )), + ) + .is_none() + ); +} + +#[test] +fn normalized_llm_paths_use_configured_anthropic_codec_without_a_system_message() { + let backend = crate::builtin::CompiledBuiltinBackend::new( + BuiltinBackendConfig { + action: "regex_replace".to_string(), + pattern: Some("sk-[A-Za-z0-9_-]+".to_string()), + replacement: Some("[REDACTED]".to_string()), + target_paths: vec!["/messages/0/content".to_string()], + ..BuiltinBackendConfig::default() + }, + Some("anthropic_messages".to_string()), + ) + .unwrap(); + let sanitize_request = crate::builtin::llm_sanitize_request_callback(backend); + + let sanitized = sanitize_request( + LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "claude-sonnet-4-6", + "messages": [{"role": "user", "content": "sk-anthropic-secret"}], + }), + }, + no_codec_request_context(), + ) + .expect("the configured Anthropic codec must sanitize a valid message-only request"); + + assert_eq!( + sanitized.content["messages"][0]["content"], + json!("[REDACTED]") + ); +} + #[test] fn trajectory_preset_redacts_chat_content_without_erasing_request_structure() { let callback = crate::builtin::llm_sanitize_request_callback(trajectory_backend( @@ -298,7 +547,8 @@ fn trajectory_preset_redacts_chat_content_without_erasing_request_structure() { "participant": {"name": "Alice Example", "username": "alice"}, "person_name": "Alice Example" }), - }); + }, no_codec_request_context()) + .unwrap(); assert_eq!(request.content["model"], "claude-sonnet-4-6"); assert_eq!(request.content["temperature"], 0.2); @@ -351,19 +601,23 @@ fn trajectory_preset_preserves_response_analytics_and_redacts_response_content() Some("openai_chat"), "preserve", )); - let sanitized = callback(json!({ - "id": "chatcmpl_1", - "model": "claude-opus-4-6", - "choices": [{"index": 0, "finish_reason": "tool_calls", "message": { - "role": "assistant", - "content": "private answer", - "tool_calls": [{"id": "call_1", "type": "function", "function": { - "name": "terminal", "arguments": "{\"command\":\"cat secret.txt\"}" - }}] - }, "logprobs": {"content": [{"token": "secret", "logprob": -0.5}]}}], - "usage": {"prompt_tokens": 20, "completion_tokens": 5, "total_tokens": 25}, - "cost": {"total": 1.25} - })); + let sanitized = callback( + json!({ + "id": "chatcmpl_1", + "model": "claude-opus-4-6", + "choices": [{"index": 0, "finish_reason": "tool_calls", "message": { + "role": "assistant", + "content": "private answer", + "tool_calls": [{"id": "call_1", "type": "function", "function": { + "name": "terminal", "arguments": "{\"command\":\"cat secret.txt\"}" + }}] + }, "logprobs": {"content": [{"token": "secret", "logprob": -0.5}]}}], + "usage": {"prompt_tokens": 20, "completion_tokens": 5, "total_tokens": 25}, + "cost": {"total": 1.25} + }), + no_codec_context(), + ) + .unwrap(); assert_eq!(sanitized["id"], "chatcmpl_1"); assert_eq!(sanitized["model"], "claude-opus-4-6"); @@ -402,7 +656,8 @@ fn trajectory_preset_covers_responses_and_anthropic_provider_shapes() { "reasoning": {"effort": "high", "summary": "private reasoning"}, "max_output_tokens": 100 }), - }); + }, no_codec_request_context()) + .unwrap(); assert_eq!(responses_request.content["model"], "gpt-5"); assert_eq!(responses_request.content["input"][0]["role"], "user"); assert_eq!( @@ -414,15 +669,19 @@ fn trajectory_preset_covers_responses_and_anthropic_provider_shapes() { let responses_response = crate::builtin::llm_sanitize_response_callback(trajectory_backend( Some("openai_responses"), "preserve", - ))(json!({ - "id": "resp_1", - "model": "gpt-5", - "status": "completed", - "output": [{"id": "msg_1", "type": "message", "role": "assistant", "content": [ - {"type": "output_text", "text": "private output"} - ]}], - "usage": {"input_tokens": 10, "output_tokens": 4, "total_tokens": 14} - })); + ))( + json!({ + "id": "resp_1", + "model": "gpt-5", + "status": "completed", + "output": [{"id": "msg_1", "type": "message", "role": "assistant", "content": [ + {"type": "output_text", "text": "private output"} + ]}], + "usage": {"input_tokens": 10, "output_tokens": 4, "total_tokens": 14} + }), + no_codec_context(), + ) + .unwrap(); assert_eq!(responses_response["id"], "resp_1"); assert_eq!(responses_response["status"], "completed"); assert_eq!(responses_response["output"][0]["id"], "msg_1"); @@ -446,7 +705,8 @@ fn trajectory_preset_covers_responses_and_anthropic_provider_shapes() { ]}], "max_tokens": 128 }), - }); + }, no_codec_request_context()) + .unwrap(); assert_eq!(anthropic_request.content["model"], "claude-sonnet-4-6"); assert_eq!(anthropic_request.content["system"], "[REDACTED]"); assert_eq!( @@ -474,7 +734,8 @@ fn trajectory_preset_covers_responses_and_anthropic_provider_shapes() { ], "stop_reason": "end_turn", "usage": {"input_tokens": 12, "output_tokens": 6, "cache_read_input_tokens": 8} - })); + }), no_codec_context()) + .unwrap(); assert_eq!(anthropic_response["id"], "msg_1"); assert_eq!(anthropic_response["role"], "assistant"); assert_eq!(anthropic_response["content"][0]["type"], "thinking"); @@ -1451,7 +1712,7 @@ fn validate_rejects_builtin_mode_without_builtin_section() { } #[test] -fn validate_rejects_llm_surfaces_without_codec() { +fn validate_allows_llm_surfaces_without_codec() { let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); reset_runtime(); @@ -1464,12 +1725,7 @@ fn validate_rejects_llm_surfaces_without_codec() { "output": false, }))); - assert!(report.diagnostics.iter().any(|diag| { - diag.field.as_deref() == Some("codec") - && diag - .message - .contains("codec is required when any LLM surface is enabled") - })); + assert!(report.diagnostics.is_empty(), "{report:?}"); } #[test] @@ -2029,6 +2285,93 @@ async fn trajectory_preset_sanitizes_stream_finalization_without_changing_client clear_plugin_configuration().unwrap(); } +#[tokio::test] +async fn normalized_paths_use_the_active_codec_for_stream_finalization() { + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + setup_isolated_thread(); + + initialize_plugins(plugin_config(json!({ + "mode": "builtin", + "codec": "openai_responses", + "input": false, + "output": true, + "tool_input": false, + "tool_output": false, + "builtin": { + "action": "regex_replace", + "pattern": "sk-[A-Za-z0-9_-]+", + "replacement": "[REDACTED]", + "target_paths": ["/message"] + } + }))) + .await + .unwrap(); + + let events = capture_events("pii-active-codec-stream-finalization"); + let provider: LlmStreamExecutionNextFn = Arc::new(move |_request| { + Box::pin(async move { + Ok(LlmJsonStream::new(futures::stream::iter(vec![Ok(json!({ + "id": "chatcmpl-stream", + "object": "chat.completion.chunk", + "choices": [{"index": 0, "delta": {"content": "sk-client-visible"}}] + }))]))) + }) + }); + let request_codec: Arc = Arc::new(OpenAIChatCodec); + let response_codec: Arc = Arc::new(OpenAIChatCodec); + let mut stream = llm_stream_call_execute( + LlmStreamCallExecuteParams::builder() + .name("openai") + .request(LlmRequest { + headers: serde_json::Map::new(), + content: json!({"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hello"}]}), + }) + .func(provider) + .collector(Box::new(|_| Ok(()))) + .finalizer(Box::new(|| { + json!({ + "id": "chatcmpl-stream", + "model": "gpt-4o-mini", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "sk-stream-secret"}, + "finish_reason": "stop" + }] + }) + })) + .codec(request_codec) + .response_codec(response_codec) + .build(), + ) + .await + .unwrap(); + + assert_eq!( + stream.next().await.unwrap().unwrap()["choices"][0]["delta"]["content"], + json!("sk-client-visible") + ); + assert!(stream.next().await.is_none()); + + let captured = captured_events_snapshot(&events); + let end = captured + .iter() + .find(|event| event.scope_category() == Some(ScopeCategory::End)) + .unwrap(); + assert_eq!( + end.output().unwrap()["choices"][0]["message"]["content"], + json!("[REDACTED]") + ); + assert_eq!( + end.annotated_response() + .and_then(|response| response.response_text()), + Some("[REDACTED]") + ); + + deregister_subscriber("pii-active-codec-stream-finalization").unwrap(); + clear_plugin_configuration().unwrap(); +} + #[test] fn builtin_backend_sanitizes_tool_start_and_end_payloads_with_preorder_targets() { let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); @@ -3743,7 +4086,7 @@ fn builtin_backend_sanitizes_llm_start_payload_via_codec_and_reencodes_provider_ } #[tokio::test] -async fn builtin_backend_sanitizes_llm_end_payload_and_response_codec_decodes_sanitized_output() { +async fn builtin_backend_removes_targeted_message_names_and_ignores_missing_normalized_paths() { let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); reset_runtime(); setup_isolated_thread(); @@ -3751,50 +4094,286 @@ async fn builtin_backend_sanitizes_llm_end_payload_and_response_codec_decodes_sa initialize_plugins(plugin_config(json!({ "mode": "builtin", "codec": "openai_chat", - "input": false, - "output": true, + "input": true, + "output": false, "tool_input": false, "tool_output": false, "builtin": { - "action": "regex_replace", - "pattern": "sk-[A-Za-z0-9_-]+", - "replacement": "[REDACTED]", - "target_paths": ["/choices/0/message/content", "/audit_owner"] + "action": "remove", + "target_paths": [ + "/messages/0/missing", + "/messages/0/name", + "/messages/1/name" + ] } }))) .await .unwrap(); - let events = capture_events("pii-redaction-llm-end-events"); + let events = capture_events("pii-redaction-remove-message-names"); let request = LlmRequest { headers: serde_json::Map::new(), content: json!({ "model": "gpt-4o-mini", "messages": [ - {"role": "user", "content": "hello"} + {"role": "system", "name": "SECRET-SYSTEM", "content": "instructions"}, + {"role": "user", "name": "SECRET-USER", "content": "hello"} ] }), }; - let response = json!({ - "id": "chatcmpl-123", - "model": "gpt-4o-mini", - "choices": [ - { - "index": 0, - "message": { - "role": "assistant", - "content": "sk-response-secret" - }, - "finish_reason": "stop" - } - ], - "usage": { - "prompt_tokens": 3, - "completion_tokens": 2, - "total_tokens": 5 - } - }); - let response_codec: Arc = Arc::new(OpenAIChatCodec); + + llm_call( + LlmCallParams::builder() + .name("openai") + .request(&request) + .build(), + ) + .unwrap(); + + let captured_events = captured_events_snapshot(&events); + assert_eq!(captured_events.len(), 1); + assert_eq!( + captured_events[0].input(), + Some(&json!({ + "headers": {}, + "content": { + "model": "gpt-4o-mini", + "messages": [ + {"role": "system", "content": "instructions"}, + {"role": "user", "content": "hello"} + ] + } + })) + ); + assert!( + !serde_json::to_string(&captured_events[0]) + .unwrap() + .contains("SECRET-") + ); + + deregister_subscriber("pii-redaction-remove-message-names").unwrap(); + clear_plugin_configuration().unwrap(); +} + +#[tokio::test] +async fn builtin_backend_omits_request_and_annotation_for_unsafe_normalized_array_removal() { + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + setup_isolated_thread(); + + initialize_plugins(plugin_config(json!({ + "mode": "builtin", + "codec": "openai_responses", + "input": true, + "output": false, + "tool_input": false, + "tool_output": false, + "builtin": { + "action": "remove", + "target_paths": ["/messages/0/content/0"] + } + }))) + .await + .unwrap(); + + let events = capture_events("pii-redaction-unsafe-normalized-array-removal"); + let request = LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "gpt-4.1-mini", + "input": [{ + "role": "user", + "content": [{"type": "input_text", "text": "SECRET-REQUEST"}] + }] + }), + }; + let annotated_request = Arc::new(OpenAIResponsesCodec.decode(&request).unwrap()); + + llm_call( + LlmCallParams::builder() + .name("openai") + .request(&request) + .annotated_request(annotated_request) + .build(), + ) + .unwrap(); + + let captured_events = captured_events_snapshot(&events); + assert_eq!(captured_events.len(), 1); + assert!(captured_events[0].input().is_none()); + assert!(captured_events[0].annotated_request().is_none()); + assert!( + !serde_json::to_string(&captured_events[0]) + .unwrap() + .contains("SECRET-REQUEST") + ); + + deregister_subscriber("pii-redaction-unsafe-normalized-array-removal").unwrap(); + clear_plugin_configuration().unwrap(); +} + +#[tokio::test] +async fn builtin_backend_sanitizes_observable_llm_request_headers() { + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + setup_isolated_thread(); + + initialize_plugins(plugin_config(json!({ + "mode": "builtin", + "input": true, + "output": false, + "tool_input": false, + "tool_output": false, + "builtin": { + "action": "redact", + "detector": "email" + } + }))) + .await + .unwrap(); + + let events = capture_events("pii-redaction-llm-request-headers"); + let request = LlmRequest { + headers: [("x-user-email".to_string(), json!("alice@example.com"))] + .into_iter() + .collect(), + content: json!({ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello"}] + }), + }; + + llm_call( + LlmCallParams::builder() + .name("openai") + .request(&request) + .build(), + ) + .unwrap(); + + let captured_events = captured_events_snapshot(&events); + assert_eq!(captured_events.len(), 1); + assert_eq!( + captured_events[0].input().unwrap()["headers"]["x-user-email"], + json!("[REDACTED]") + ); + assert!( + !serde_json::to_string(&captured_events[0]) + .unwrap() + .contains("alice@example.com") + ); + + deregister_subscriber("pii-redaction-llm-request-headers").unwrap(); + clear_plugin_configuration().unwrap(); +} + +#[tokio::test] +async fn trajectory_preset_redacts_opaque_llm_request_header_values() { + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + setup_isolated_thread(); + + initialize_plugins(plugin_config(json!({ + "mode": "builtin", + "input": true, + "output": false, + "tool_input": false, + "tool_output": false, + "builtin": { + "preset": "trajectory_context" + } + }))) + .await + .unwrap(); + + let events = capture_events("pii-trajectory-llm-request-headers"); + let request = LlmRequest { + headers: [("model".to_string(), json!("SECRET-HEADER"))] + .into_iter() + .collect(), + content: json!({ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello"}] + }), + }; + + llm_call( + LlmCallParams::builder() + .name("openai") + .request(&request) + .build(), + ) + .unwrap(); + + let captured_events = captured_events_snapshot(&events); + assert_eq!(captured_events.len(), 1); + assert_eq!( + captured_events[0].input().unwrap()["headers"]["model"], + json!("[REDACTED]") + ); + assert!( + !serde_json::to_string(&captured_events[0]) + .unwrap() + .contains("SECRET-HEADER") + ); + + deregister_subscriber("pii-trajectory-llm-request-headers").unwrap(); + clear_plugin_configuration().unwrap(); +} + +#[tokio::test] +async fn builtin_backend_sanitizes_llm_end_payload_and_response_codec_decodes_sanitized_output() { + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + setup_isolated_thread(); + + initialize_plugins(plugin_config(json!({ + "mode": "builtin", + "codec": "openai_chat", + "input": false, + "output": true, + "tool_input": false, + "tool_output": false, + "builtin": { + "action": "regex_replace", + "pattern": "sk-[A-Za-z0-9_-]+", + "replacement": "[REDACTED]", + "target_paths": ["/choices/0/message/content", "/audit_owner"] + } + }))) + .await + .unwrap(); + + let events = capture_events("pii-redaction-llm-end-events"); + let request = LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "gpt-4o-mini", + "messages": [ + {"role": "user", "content": "hello"} + ] + }), + }; + let response = json!({ + "id": "chatcmpl-123", + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "sk-response-secret" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 3, + "completion_tokens": 2, + "total_tokens": 5 + } + }); + let response_codec: Arc = Arc::new(OpenAIChatCodec); let result = llm_call_execute( LlmCallExecuteParams::builder() @@ -3915,6 +4494,142 @@ async fn builtin_backend_sanitizes_openai_chat_response_from_normalized_message_ clear_plugin_configuration().unwrap(); } +#[tokio::test] +async fn builtin_backend_omits_unprojectable_openai_chat_api_specific_response() { + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + setup_isolated_thread(); + + initialize_plugins(plugin_config(json!({ + "mode": "builtin", + "codec": "openai_chat", + "input": false, + "output": true, + "tool_input": false, + "tool_output": false, + "builtin": { + "action": "remove", + "target_paths": ["/api_specific/system_fingerprint"] + } + }))) + .await + .unwrap(); + + let events = capture_events("pii-redaction-openai-chat-api-specific-response"); + let response = json!({ + "id": "chatcmpl-api-specific", + "model": "gpt-4o-mini", + "system_fingerprint": "SECRET-FINGERPRINT", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "hello"}, + "finish_reason": "stop" + }] + }); + + let result = llm_call_execute( + LlmCallExecuteParams::builder() + .name("openai") + .request(LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello"}] + }), + }) + .func(noop_openai_chat_exec_fn(response.clone())) + .response_codec(Arc::new(OpenAIChatCodec)) + .build(), + ) + .await + .unwrap(); + + assert_eq!(result, response); + let captured_events = captured_events_snapshot(&events); + assert_eq!(captured_events.len(), 2); + assert!(captured_events[1].output().is_none()); + assert!(captured_events[1].annotated_response().is_none()); + assert!( + !serde_json::to_string(&captured_events[1]) + .unwrap() + .contains("SECRET-FINGERPRINT") + ); + + deregister_subscriber("pii-redaction-openai-chat-api-specific-response").unwrap(); + clear_plugin_configuration().unwrap(); +} + +#[tokio::test] +async fn builtin_backend_omits_multi_choice_openai_chat_normalized_response() { + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + setup_isolated_thread(); + + initialize_plugins(plugin_config(json!({ + "mode": "builtin", + "codec": "openai_chat", + "input": false, + "output": true, + "tool_input": false, + "tool_output": false, + "builtin": { + "action": "regex_replace", + "pattern": "sk-[A-Za-z0-9_-]+", + "replacement": "[REDACTED]", + "target_paths": ["/message"] + } + }))) + .await + .unwrap(); + + let events = capture_events("pii-redaction-openai-chat-multi-choice-response"); + let response = json!({ + "id": "chatcmpl-multi", + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "sk-first-secret"}, + "finish_reason": "stop" + }, + { + "index": 1, + "message": {"role": "assistant", "content": "sk-second-secret"}, + "finish_reason": "stop" + } + ] + }); + + let result = llm_call_execute( + LlmCallExecuteParams::builder() + .name("openai") + .request(LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello"}] + }), + }) + .func(noop_openai_chat_exec_fn(response.clone())) + .response_codec(Arc::new(OpenAIChatCodec)) + .build(), + ) + .await + .unwrap(); + + assert_eq!(result, response); + let captured_events = captured_events_snapshot(&events); + assert_eq!(captured_events.len(), 2); + assert!(captured_events[1].output().is_none()); + assert!(captured_events[1].annotated_response().is_none()); + let serialized_end = serde_json::to_string(&captured_events[1]).unwrap(); + assert!(!serialized_end.contains("sk-first-secret")); + assert!(!serialized_end.contains("sk-second-secret")); + + deregister_subscriber("pii-redaction-openai-chat-multi-choice-response").unwrap(); + clear_plugin_configuration().unwrap(); +} + #[tokio::test] async fn builtin_redact_sanitizes_openai_chat_response_from_detector_path() { let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); @@ -4044,6 +4759,65 @@ async fn builtin_backend_sanitizes_anthropic_response_from_normalized_message_pa clear_plugin_configuration().unwrap(); } +#[tokio::test] +async fn builtin_backend_omits_unprojectable_anthropic_normalized_usage_response() { + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + setup_isolated_thread(); + + initialize_plugins(plugin_config(json!({ + "mode": "builtin", + "codec": "anthropic_messages", + "input": false, + "output": true, + "tool_input": false, + "tool_output": false, + "builtin": { + "action": "remove", + "target_paths": ["/usage/prompt_tokens"] + } + }))) + .await + .unwrap(); + + let events = capture_events("pii-redaction-anthropic-normalized-usage-response"); + let response = json!({ + "id": "msg_usage", + "model": "claude-sonnet-4-20250514", + "role": "assistant", + "type": "message", + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 123, "output_tokens": 4} + }); + + let result = llm_call_execute( + LlmCallExecuteParams::builder() + .name("anthropic") + .request(LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "claude-sonnet-4-20250514", + "messages": [{"role": "user", "content": "hello"}] + }), + }) + .func(noop_openai_chat_exec_fn(response.clone())) + .response_codec(Arc::new(crate::codec::anthropic::AnthropicMessagesCodec)) + .build(), + ) + .await + .unwrap(); + + assert_eq!(result, response); + let captured_events = captured_events_snapshot(&events); + assert_eq!(captured_events.len(), 2); + assert!(captured_events[1].output().is_none()); + assert!(captured_events[1].annotated_response().is_none()); + + deregister_subscriber("pii-redaction-anthropic-normalized-usage-response").unwrap(); + clear_plugin_configuration().unwrap(); +} + #[tokio::test] async fn builtin_backend_sanitizes_openai_responses_response_from_normalized_message_path() { let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); @@ -4052,7 +4826,7 @@ async fn builtin_backend_sanitizes_openai_responses_response_from_normalized_mes initialize_plugins(plugin_config(json!({ "mode": "builtin", - "codec": "openai_responses", + "codec": "openai_chat", "input": false, "output": true, "tool_input": false, @@ -4081,6 +4855,7 @@ async fn builtin_backend_sanitizes_openai_responses_response_from_normalized_mes "id": "resp_123", "model": "gpt-4.1-mini", "status": "completed", + "output_text": "sk-responses-secret", "output": [ { "type": "message", @@ -4101,13 +4876,117 @@ async fn builtin_backend_sanitizes_openai_responses_response_from_normalized_mes captured_events[1].output().unwrap()["output"][0]["content"][0]["text"], json!("[REDACTED]") ); + assert_eq!( + captured_events[1].output().unwrap()["output_text"], + json!("[REDACTED]") + ); assert_eq!( captured_events[1] .annotated_response() .and_then(|response| response.response_text()), Some("[REDACTED]") ); + assert!( + !captured_events[1] + .to_json_string() + .unwrap() + .contains("sk-responses-secret") + ); deregister_subscriber("pii-redaction-openai-responses-normalized-response").unwrap(); clear_plugin_configuration().unwrap(); } + +#[tokio::test] +async fn builtin_backend_sanitizes_openai_responses_output_text_alias_on_stream_finalization() { + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + setup_isolated_thread(); + + initialize_plugins(plugin_config(json!({ + "mode": "builtin", + "input": false, + "output": true, + "tool_input": false, + "tool_output": false, + "builtin": { + "action": "regex_replace", + "pattern": "sk-[A-Za-z0-9_-]+", + "replacement": "[REDACTED]", + "target_paths": ["/message"] + } + }))) + .await + .unwrap(); + + let events = capture_events("pii-redaction-openai-responses-output-text-stream"); + let provider: LlmStreamExecutionNextFn = Arc::new(move |_request| { + Box::pin(async move { + Ok(LlmJsonStream::new(futures::stream::iter(vec![Ok(json!({ + "type": "response.output_text.delta", + "delta": "sk-client-visible" + }))]))) + }) + }); + let request_codec: Arc = Arc::new(OpenAIResponsesCodec); + let response_codec: Arc = Arc::new(OpenAIResponsesCodec); + let mut stream = llm_stream_call_execute( + LlmStreamCallExecuteParams::builder() + .name("openai") + .request(LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "gpt-4.1-mini", + "input": "hello" + }), + }) + .func(provider) + .collector(Box::new(|_| Ok(()))) + .finalizer(Box::new(|| { + json!({ + "id": "resp_stream", + "model": "gpt-4.1-mini", + "status": "completed", + "output_text": "sk-stream-secret", + "output": [{ + "type": "message", + "content": [{ + "type": "output_text", + "text": "sk-stream-secret" + }] + }] + }) + })) + .codec(request_codec) + .response_codec(response_codec) + .build(), + ) + .await + .unwrap(); + + assert_eq!( + stream.next().await.unwrap().unwrap()["delta"], + json!("sk-client-visible") + ); + assert!(stream.next().await.is_none()); + + let captured_events = captured_events_snapshot(&events); + let end = captured_events + .iter() + .find(|event| event.scope_category() == Some(ScopeCategory::End)) + .unwrap(); + assert_eq!( + end.output().unwrap()["output"][0]["content"][0]["text"], + json!("[REDACTED]") + ); + assert_eq!(end.output().unwrap()["output_text"], json!("[REDACTED]")); + assert_eq!( + end.annotated_response() + .and_then(|response| response.response_text()), + Some("[REDACTED]") + ); + assert!(!end.to_json_string().unwrap().contains("sk-stream-secret")); + + deregister_subscriber("pii-redaction-openai-responses-output-text-stream").unwrap(); + clear_plugin_configuration().unwrap(); +} diff --git a/crates/plugin/src/lib.rs b/crates/plugin/src/lib.rs index 056661db0..727e09aa2 100644 --- a/crates/plugin/src/lib.rs +++ b/crates/plugin/src/lib.rs @@ -36,7 +36,52 @@ use serde::{Serialize, de::DeserializeOwned}; use serde_json::Map; /// Native plugin ABI version supported by this crate. -pub const NEMO_RELAY_NATIVE_ABI_VERSION: u32 = 1; +pub const NEMO_RELAY_NATIVE_ABI_VERSION: u32 = 2; + +/// Built-in LLM codec identities available to native plugins. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, serde::Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum BuiltinLlmCodec { + /// OpenAI Chat Completions. + #[serde(rename = "openai_chat")] + OpenAiChat, + /// OpenAI Responses. + #[serde(rename = "openai_responses")] + OpenAiResponses, + /// Anthropic Messages. + #[serde(rename = "anthropic_messages")] + AnthropicMessages, +} + +/// Per-call LLM codec identity delivered to native plugins. +#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, serde::Deserialize)] +#[serde(tag = "kind", content = "id", rename_all = "snake_case")] +pub enum LlmCodecIdentity { + /// No codec was active. + #[default] + None, + /// A Relay built-in codec was active. + #[serde(rename = "builtin")] + BuiltIn(BuiltinLlmCodec), + /// A runtime-registered codec was active, identified by its stable ID. + Runtime(String), + /// A codec was active but has no registered identity. + Opaque, +} + +/// Per-call request codec context delivered to an LLM sanitizer. +pub struct LlmSanitizeRequestContext<'a> { + /// Identity of the active codec. + pub codec: LlmCodecIdentity, + resolved: Option>, +} + +/// Per-call response codec context delivered to an LLM sanitizer. +pub struct LlmSanitizeResponseContext<'a> { + /// Identity of the active codec. + pub codec: LlmCodecIdentity, + resolved: Option>, +} /// Status codes returned by stable native ABI functions. #[repr(i32)] @@ -73,6 +118,144 @@ pub struct NemoRelayNativeString { _marker: PhantomData<(*mut u8, PhantomPinned)>, } +/// Opaque callback-scoped request codec capability owned by the host. +#[repr(C)] +pub struct NemoRelayNativeLlmRequestCodec { + _private: [u8; 0], + _marker: PhantomData<(*mut u8, PhantomPinned)>, +} + +/// Opaque callback-scoped response codec capability owned by the host. +#[repr(C)] +pub struct NemoRelayNativeLlmResponseCodec { + _private: [u8; 0], + _marker: PhantomData<(*mut u8, PhantomPinned)>, +} + +/// Discriminator for the codec supplied to an LLM sanitizer over the native ABI. +#[repr(u32)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum NemoRelayNativeLlmCodecKind { + /// No codec was active for this call. + None = 0, + /// A Relay built-in codec was active. + BuiltIn = 1, + /// A runtime-registered codec was active. + Runtime = 2, + /// A codec was active but has no registered identity. + Opaque = 3, +} + +/// Per-call LLM sanitizer context passed over the native ABI. +/// +/// `codec_id` is borrowed for the duration of the callback. It is null for +/// [`NemoRelayNativeLlmCodecKind::None`] and +/// [`NemoRelayNativeLlmCodecKind::Opaque`]. For `BuiltIn`, it is one of the +/// stable built-in codec IDs; for `Runtime`, it is the registered codec ID. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct NemoRelayNativeLlmSanitizeRequestContext { + /// Discriminator for the active codec. + pub codec_kind: NemoRelayNativeLlmCodecKind, + /// Optional borrowed codec identifier. + pub codec_id: *const NemoRelayNativeString, + /// Borrowed request codec capability, or null when no codec is active. + pub codec: *const NemoRelayNativeLlmRequestCodec, +} + +/// Per-call response sanitizer context passed over the native ABI. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct NemoRelayNativeLlmSanitizeResponseContext { + /// Discriminator for the active codec. + pub codec_kind: NemoRelayNativeLlmCodecKind, + /// Optional borrowed codec identifier. + pub codec_id: *const NemoRelayNativeString, + /// Borrowed response codec capability, or null when no codec is active. + pub codec: *const NemoRelayNativeLlmResponseCodec, +} + +/// Safe callback-scoped request codec facade for typed native plugins. +pub struct LlmSanitizeRequestCodec<'a> { + host: NemoRelayNativeHostApiV1, + handle: *const NemoRelayNativeLlmRequestCodec, + _lifetime: PhantomData<&'a NemoRelayNativeLlmRequestCodec>, +} + +impl LlmSanitizeRequestCodec<'_> { + /// Decode an opaque request into Relay's normalized request model. + pub fn decode(&self, request: &LlmRequest) -> Result { + native_codec_call(&self.host, |out| unsafe { + let request = HostString::from_json(&self.host, request) + .ok_or_else(|| "failed to serialize LLM request".to_string())?; + codec_status( + &self.host, + (self.host.llm_request_codec_decode)(self.handle, request.as_ptr(), out), + ) + }) + } + + /// Encode normalized changes onto the original opaque request. + pub fn encode( + &self, + annotated: &AnnotatedLlmRequest, + original: &LlmRequest, + ) -> Result { + native_codec_call(&self.host, |out| unsafe { + let annotated = HostString::from_json(&self.host, annotated) + .ok_or_else(|| "failed to serialize annotated request".to_string())?; + let original = HostString::from_json(&self.host, original) + .ok_or_else(|| "failed to serialize original request".to_string())?; + codec_status( + &self.host, + (self.host.llm_request_codec_encode)( + self.handle, + annotated.as_ptr(), + original.as_ptr(), + out, + ), + ) + }) + } +} + +/// Safe callback-scoped response codec facade for typed native plugins. +pub struct LlmSanitizeResponseCodec<'a> { + host: NemoRelayNativeHostApiV1, + handle: *const NemoRelayNativeLlmResponseCodec, + _lifetime: PhantomData<&'a NemoRelayNativeLlmResponseCodec>, +} + +impl LlmSanitizeResponseCodec<'_> { + /// Decode an opaque response into Relay's normalized response model. + pub fn decode(&self, response: &Json) -> Result { + native_codec_call(&self.host, |out| unsafe { + let response = HostString::from_json(&self.host, response) + .ok_or_else(|| "failed to serialize LLM response".to_string())?; + codec_status( + &self.host, + (self.host.llm_response_codec_decode)(self.handle, response.as_ptr(), out), + ) + }) + } +} + +impl<'a> LlmSanitizeRequestContext<'a> { + /// Resolve the active request codec capability. + #[must_use] + pub fn resolve_codec(&self) -> Option<&LlmSanitizeRequestCodec<'a>> { + self.resolved.as_ref() + } +} + +impl<'a> LlmSanitizeResponseContext<'a> { + /// Resolve the active response codec capability. + #[must_use] + pub fn resolve_codec(&self) -> Option<&LlmSanitizeResponseCodec<'a>> { + self.resolved.as_ref() + } +} + /// Opaque plugin registration context borrowed from the host during registration. #[repr(C)] pub struct NemoRelayNativePluginContext { @@ -241,17 +424,27 @@ pub type NemoRelayNativeToolExecutionCb = unsafe extern "C" fn( out_outcome_json: *mut *mut NemoRelayNativeString, ) -> NemoRelayStatus; -/// Native LLM request transform callback for request sanitizers. -pub type NemoRelayNativeLlmRequestCb = unsafe extern "C" fn( +/// Native LLM request sanitizer callback. Return a successful null output to +/// omit the observability payload and annotation. `request_json` is borrowed, +/// but may be written directly to `out_request_json` as a pass-through; the +/// host releases an aliased input/output once. Any other non-null output must +/// be host-allocated and transfers ownership to the host. +pub type NemoRelayNativeLlmSanitizeRequestCb = unsafe extern "C" fn( user_data: *mut c_void, request_json: *const NemoRelayNativeString, + context: NemoRelayNativeLlmSanitizeRequestContext, out_request_json: *mut *mut NemoRelayNativeString, ) -> NemoRelayStatus; -/// Native JSON transform callback for LLM response sanitizers. -pub type NemoRelayNativeJsonCb = unsafe extern "C" fn( +/// Native LLM response sanitizer callback. Return a successful null output to +/// omit the observability payload and annotation. `payload_json` is borrowed, +/// but may be written directly to `out_json` as a pass-through; the host +/// releases an aliased input/output once. Any other non-null output must be +/// host-allocated and transfers ownership to the host. +pub type NemoRelayNativeLlmSanitizeResponseCb = unsafe extern "C" fn( user_data: *mut c_void, payload_json: *const NemoRelayNativeString, + context: NemoRelayNativeLlmSanitizeResponseContext, out_json: *mut *mut NemoRelayNativeString, ) -> NemoRelayStatus; @@ -334,6 +527,25 @@ pub struct NemoRelayNativeHostApiV1 { pub last_error_clear: unsafe extern "C" fn(), /// Sets the host thread-local native ABI error message. pub last_error_set: unsafe extern "C" fn(message: *const NemoRelayNativeString), + /// Decodes an LLM request through a callback-scoped codec capability. + pub llm_request_codec_decode: unsafe extern "C" fn( + codec: *const NemoRelayNativeLlmRequestCodec, + request_json: *const NemoRelayNativeString, + out: *mut *mut NemoRelayNativeString, + ) -> NemoRelayStatus, + /// Encodes normalized request changes through a callback-scoped codec capability. + pub llm_request_codec_encode: unsafe extern "C" fn( + codec: *const NemoRelayNativeLlmRequestCodec, + annotated_json: *const NemoRelayNativeString, + original_json: *const NemoRelayNativeString, + out: *mut *mut NemoRelayNativeString, + ) -> NemoRelayStatus, + /// Decodes an LLM response through a callback-scoped codec capability. + pub llm_response_codec_decode: unsafe extern "C" fn( + codec: *const NemoRelayNativeLlmResponseCodec, + response_json: *const NemoRelayNativeString, + out: *mut *mut NemoRelayNativeString, + ) -> NemoRelayStatus, /// Registers an event subscriber through the plugin context. pub plugin_context_register_subscriber: unsafe extern "C" fn( ctx: *mut NemoRelayNativePluginContext, @@ -399,7 +611,7 @@ pub struct NemoRelayNativeHostApiV1 { ctx: *mut NemoRelayNativePluginContext, name: *const NemoRelayNativeString, priority: i32, - cb: NemoRelayNativeLlmRequestCb, + cb: NemoRelayNativeLlmSanitizeRequestCb, user_data: *mut c_void, free_fn: NemoRelayNativeFreeFn, ) -> NemoRelayStatus, @@ -409,7 +621,7 @@ pub struct NemoRelayNativeHostApiV1 { ctx: *mut NemoRelayNativePluginContext, name: *const NemoRelayNativeString, priority: i32, - cb: NemoRelayNativeJsonCb, + cb: NemoRelayNativeLlmSanitizeResponseCb, user_data: *mut c_void, free_fn: NemoRelayNativeFreeFn, ) -> NemoRelayStatus, @@ -1526,7 +1738,10 @@ impl<'a> PluginContext<'a> { callback: F, ) -> Result<()> where - F: Fn(LlmRequest) -> LlmRequest + Send + Sync + 'static, + F: for<'ctx> Fn(LlmRequest, LlmSanitizeRequestContext<'ctx>) -> Option + + Send + + Sync + + 'static, { let user_data = typed_callback_user_data(self.host, callback); let status = unsafe { @@ -1554,7 +1769,10 @@ impl<'a> PluginContext<'a> { callback: F, ) -> Result<()> where - F: Fn(Json) -> Json + Send + Sync + 'static, + F: for<'ctx> Fn(Json, LlmSanitizeResponseContext<'ctx>) -> Option + + Send + + Sync + + 'static, { let user_data = typed_callback_user_data(self.host, callback); let status = unsafe { @@ -1655,7 +1873,7 @@ impl<'a> PluginContext<'a> { /// Registers a typed LLM stream execution intercept. /// - /// Native ABI v1 represents stream execution as one JSON result. The host + /// Native ABI v2 represents stream execution as one JSON result. The host /// wraps that result as a one-chunk stream. pub fn register_llm_stream_execution_intercept( &mut self, @@ -1890,7 +2108,7 @@ impl<'a> PluginContext<'a> { &mut self, name: &str, priority: i32, - cb: NemoRelayNativeLlmRequestCb, + cb: NemoRelayNativeLlmSanitizeRequestCb, user_data: *mut c_void, free_fn: NemoRelayNativeFreeFn, ) -> NemoRelayStatus { @@ -1911,7 +2129,7 @@ impl<'a> PluginContext<'a> { &mut self, name: &str, priority: i32, - cb: NemoRelayNativeJsonCb, + cb: NemoRelayNativeLlmSanitizeResponseCb, user_data: *mut c_void, free_fn: NemoRelayNativeFreeFn, ) -> NemoRelayStatus { @@ -2274,10 +2492,14 @@ where unsafe extern "C" fn typed_llm_sanitize_request_trampoline( user_data: *mut c_void, request_json: *const NemoRelayNativeString, + context: NemoRelayNativeLlmSanitizeRequestContext, out_request_json: *mut *mut NemoRelayNativeString, ) -> NemoRelayStatus where - F: Fn(LlmRequest) -> LlmRequest + Send + Sync + 'static, + F: for<'a> Fn(LlmRequest, LlmSanitizeRequestContext<'a>) -> Option + + Send + + Sync + + 'static, { if user_data.is_null() || out_request_json.is_null() { return NemoRelayStatus::NullPointer; @@ -2285,9 +2507,14 @@ where unsafe { *out_request_json = ptr::null_mut() }; let state = unsafe { &*(user_data as *const TypedCallback) }; let result = catch_unwind(AssertUnwindSafe(|| { + let context = llm_sanitize_request_context_from_native(&state.host, context)?; let request: LlmRequest = read_json_value(&state.host, request_json, "LLM request")?; - let output = (state.callback)(request); - Ok::<_, NemoRelayStatus>(write_json(&state.host, &output, out_request_json)) + match (state.callback)(request, context) { + Some(output) => { + Ok::<_, NemoRelayStatus>(write_json(&state.host, &output, out_request_json)) + } + None => Ok(NemoRelayStatus::Ok), + } })); match result { Ok(Ok(status)) => status, @@ -2299,10 +2526,11 @@ where unsafe extern "C" fn typed_llm_sanitize_response_trampoline( user_data: *mut c_void, payload_json: *const NemoRelayNativeString, + context: NemoRelayNativeLlmSanitizeResponseContext, out_json: *mut *mut NemoRelayNativeString, ) -> NemoRelayStatus where - F: Fn(Json) -> Json + Send + Sync + 'static, + F: for<'a> Fn(Json, LlmSanitizeResponseContext<'a>) -> Option + Send + Sync + 'static, { if user_data.is_null() || out_json.is_null() { return NemoRelayStatus::NullPointer; @@ -2310,9 +2538,12 @@ where unsafe { *out_json = ptr::null_mut() }; let state = unsafe { &*(user_data as *const TypedCallback) }; let result = catch_unwind(AssertUnwindSafe(|| { + let context = llm_sanitize_response_context_from_native(&state.host, context)?; let payload: Json = read_json_value(&state.host, payload_json, "LLM response")?; - let output = (state.callback)(payload); - Ok::<_, NemoRelayStatus>(write_json(&state.host, &output, out_json)) + match (state.callback)(payload, context) { + Some(output) => Ok::<_, NemoRelayStatus>(write_json(&state.host, &output, out_json)), + None => Ok(NemoRelayStatus::Ok), + } })); match result { Ok(Ok(status)) => status, @@ -2625,6 +2856,29 @@ impl Drop for HostString<'_> { } } +fn codec_status(host: &NemoRelayNativeHostApiV1, status: NemoRelayStatus) -> Result<()> { + if status == NemoRelayStatus::Ok { + Ok(()) + } else { + Err(status_error(host, status, "LLM codec operation")) + } +} + +fn native_codec_call( + host: &NemoRelayNativeHostApiV1, + call: impl FnOnce(*mut *mut NemoRelayNativeString) -> Result<()>, +) -> Result { + let mut out = ptr::null_mut(); + call(&mut out)?; + if out.is_null() { + return Err("LLM codec operation returned null".into()); + } + let out = HostString { host, ptr: out }; + let text = read_host_string(host, out.as_ptr()) + .map_err(|_| "LLM codec operation returned invalid UTF-8".to_string())?; + serde_json::from_str(&text).map_err(|error| format!("invalid LLM codec result: {error}")) +} + struct OptionalHostJson<'a>(Option>); impl<'a> OptionalHostJson<'a> { @@ -2765,6 +3019,60 @@ fn read_optional_json_value( } } +fn llm_codec_identity_from_native( + host: &NemoRelayNativeHostApiV1, + codec_kind: NemoRelayNativeLlmCodecKind, + codec_id: *const NemoRelayNativeString, +) -> std::result::Result { + let codec = match codec_kind { + NemoRelayNativeLlmCodecKind::None => LlmCodecIdentity::None, + NemoRelayNativeLlmCodecKind::Opaque => LlmCodecIdentity::Opaque, + NemoRelayNativeLlmCodecKind::BuiltIn => { + let id = read_required_host_string(host, codec_id, "LLM built-in codec ID")?; + let builtin = match id.as_str() { + "openai_chat" => BuiltinLlmCodec::OpenAiChat, + "openai_responses" => BuiltinLlmCodec::OpenAiResponses, + "anthropic_messages" => BuiltinLlmCodec::AnthropicMessages, + _ => { + set_last_error(host, &format!("unknown built-in LLM codec ID: {id}")); + return Err(NemoRelayStatus::InvalidArg); + } + }; + LlmCodecIdentity::BuiltIn(builtin) + } + NemoRelayNativeLlmCodecKind::Runtime => LlmCodecIdentity::Runtime( + read_required_host_string(host, codec_id, "LLM runtime codec ID")?, + ), + }; + Ok(codec) +} + +fn llm_sanitize_request_context_from_native<'a>( + host: &NemoRelayNativeHostApiV1, + context: NemoRelayNativeLlmSanitizeRequestContext, +) -> std::result::Result, NemoRelayStatus> { + let codec = llm_codec_identity_from_native(host, context.codec_kind, context.codec_id)?; + let resolved = (!context.codec.is_null()).then_some(LlmSanitizeRequestCodec { + host: *host, + handle: context.codec, + _lifetime: PhantomData, + }); + Ok(LlmSanitizeRequestContext { codec, resolved }) +} + +fn llm_sanitize_response_context_from_native<'a>( + host: &NemoRelayNativeHostApiV1, + context: NemoRelayNativeLlmSanitizeResponseContext, +) -> std::result::Result, NemoRelayStatus> { + let codec = llm_codec_identity_from_native(host, context.codec_kind, context.codec_id)?; + let resolved = (!context.codec.is_null()).then_some(LlmSanitizeResponseCodec { + host: *host, + handle: context.codec, + _lifetime: PhantomData, + }); + Ok(LlmSanitizeResponseContext { codec, resolved }) +} + enum HostStringReadError { Null, InvalidUtf8, diff --git a/crates/plugin/tests/typed_callbacks.rs b/crates/plugin/tests/typed_callbacks.rs index 4a1b20041..fb3562bd3 100644 --- a/crates/plugin/tests/typed_callbacks.rs +++ b/crates/plugin/tests/typed_callbacks.rs @@ -13,13 +13,16 @@ use std::sync::{ }; use nemo_relay_plugin::{ - AnnotatedLlmRequest, CategoryProfile, ConfigDiagnostic, DiagnosticLevel, Event, EventCategory, - EventSanitizeFields, Json, LlmJsonStream, LlmNext, LlmRequest, LlmRequestInterceptOutcome, - LlmStream, LlmStreamNext, NEMO_RELAY_NATIVE_ABI_VERSION, NativePlugin, - NemoRelayNativeEventSanitizeCb, NemoRelayNativeEventSubscriberCb, NemoRelayNativeFreeFn, - NemoRelayNativeHostApiV1, NemoRelayNativeJsonCb, NemoRelayNativeLlmConditionalCb, - NemoRelayNativeLlmExecutionCb, NemoRelayNativeLlmRequestCb, - NemoRelayNativeLlmRequestInterceptCb, NemoRelayNativeLlmStreamExecutionCb, + AnnotatedLlmRequest, BuiltinLlmCodec, CategoryProfile, ConfigDiagnostic, DiagnosticLevel, + Event, EventCategory, EventSanitizeFields, Json, LlmCodecIdentity, LlmJsonStream, LlmNext, + LlmRequest, LlmRequestInterceptOutcome, LlmStream, LlmStreamNext, + NEMO_RELAY_NATIVE_ABI_VERSION, NativePlugin, NemoRelayNativeEventSanitizeCb, + NemoRelayNativeEventSubscriberCb, NemoRelayNativeFreeFn, NemoRelayNativeHostApiV1, + NemoRelayNativeLlmCodecKind, NemoRelayNativeLlmConditionalCb, NemoRelayNativeLlmExecutionCb, + NemoRelayNativeLlmRequestCodec, NemoRelayNativeLlmRequestInterceptCb, + NemoRelayNativeLlmResponseCodec, NemoRelayNativeLlmSanitizeRequestCb, + NemoRelayNativeLlmSanitizeRequestContext, NemoRelayNativeLlmSanitizeResponseCb, + NemoRelayNativeLlmSanitizeResponseContext, NemoRelayNativeLlmStreamExecutionCb, NemoRelayNativeLlmStreamV1, NemoRelayNativePluginContext, NemoRelayNativePluginV1, NemoRelayNativeScopeHandle, NemoRelayNativeScopeStack, NemoRelayNativeScopeStackBinding, NemoRelayNativeScopeType, NemoRelayNativeString, NemoRelayNativeToolConditionalCb, @@ -114,7 +117,7 @@ impl RegisteredToolExecution { struct RegisteredLlmRequest { name: String, priority: i32, - cb: NemoRelayNativeLlmRequestCb, + cb: NemoRelayNativeLlmSanitizeRequestCb, user_data: usize, free_fn: NemoRelayNativeFreeFn, } @@ -130,7 +133,7 @@ impl RegisteredLlmRequest { struct RegisteredLlmJson { name: String, priority: i32, - cb: NemoRelayNativeJsonCb, + cb: NemoRelayNativeLlmSanitizeResponseCb, user_data: usize, free_fn: NemoRelayNativeFreeFn, } @@ -297,8 +300,8 @@ static LLM_REQUEST_INTERCEPT_REGISTRATION: Mutex(), test_host().struct_size @@ -316,13 +319,13 @@ fn native_abi_v1_struct_sizes_are_self_describing() { #[cfg(target_pointer_width = "64")] { assert_eq!(align_of::(), 8); - assert_eq!(size_of::(), 296); + assert_eq!(size_of::(), 320); assert_eq!( host_api_offsets(), [ 0, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248, 256, 264, 272, - 280, 288, + 280, 288, 296, 304, 312, ] ); assert_eq!(align_of::(), 8); @@ -336,12 +339,13 @@ fn native_abi_v1_struct_sizes_are_self_describing() { #[cfg(target_pointer_width = "32")] { assert_eq!(align_of::(), 4); - assert_eq!(size_of::(), 148); + assert_eq!(size_of::(), 160); assert_eq!( host_api_offsets(), [ 0, 4, 8, 12, 16, 20, 24, 28, 32, 36, 40, 44, 48, 52, 56, 60, 64, 68, 72, 76, 80, - 84, 88, 92, 96, 100, 104, 108, 112, 116, 120, 124, 128, 132, 136, 140, 144, + 84, 88, 92, 96, 100, 104, 108, 112, 116, 120, 124, 128, 132, 136, 140, 144, 148, + 152, 156, ] ); assert_eq!(align_of::(), 4); @@ -353,7 +357,7 @@ fn native_abi_v1_struct_sizes_are_self_describing() { } } -fn host_api_offsets() -> [usize; 37] { +fn host_api_offsets() -> [usize; 40] { [ offset_of!(NemoRelayNativeHostApiV1, abi_version), offset_of!(NemoRelayNativeHostApiV1, struct_size), @@ -364,6 +368,9 @@ fn host_api_offsets() -> [usize; 37] { offset_of!(NemoRelayNativeHostApiV1, string_free), offset_of!(NemoRelayNativeHostApiV1, last_error_clear), offset_of!(NemoRelayNativeHostApiV1, last_error_set), + offset_of!(NemoRelayNativeHostApiV1, llm_request_codec_decode), + offset_of!(NemoRelayNativeHostApiV1, llm_request_codec_encode), + offset_of!(NemoRelayNativeHostApiV1, llm_response_codec_decode), offset_of!(NemoRelayNativeHostApiV1, plugin_context_register_subscriber), offset_of!( NemoRelayNativeHostApiV1, @@ -688,7 +695,7 @@ unsafe extern "C" fn capture_llm_request( _ctx: *mut NemoRelayNativePluginContext, name: *const NemoRelayNativeString, priority: i32, - cb: NemoRelayNativeLlmRequestCb, + cb: NemoRelayNativeLlmSanitizeRequestCb, user_data: *mut c_void, free_fn: NemoRelayNativeFreeFn, ) -> NemoRelayStatus { @@ -717,7 +724,7 @@ unsafe extern "C" fn capture_llm_json( _ctx: *mut NemoRelayNativePluginContext, name: *const NemoRelayNativeString, priority: i32, - cb: NemoRelayNativeJsonCb, + cb: NemoRelayNativeLlmSanitizeResponseCb, user_data: *mut c_void, free_fn: NemoRelayNativeFreeFn, ) -> NemoRelayStatus { @@ -1126,6 +1133,57 @@ unsafe extern "C" fn true_scope_stack_active() -> bool { true } +unsafe extern "C" fn unavailable_request_codec_decode( + _codec: *const NemoRelayNativeLlmRequestCodec, + _request: *const NemoRelayNativeString, + _out: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + NemoRelayStatus::Internal +} + +unsafe extern "C" fn unavailable_request_codec_encode( + _codec: *const NemoRelayNativeLlmRequestCodec, + _annotated: *const NemoRelayNativeString, + _original: *const NemoRelayNativeString, + _out: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + NemoRelayStatus::Internal +} + +unsafe extern "C" fn unavailable_response_codec_decode( + _codec: *const NemoRelayNativeLlmResponseCodec, + _response: *const NemoRelayNativeString, + _out: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + NemoRelayStatus::Internal +} + +unsafe extern "C" fn successful_request_codec_decode( + _codec: *const NemoRelayNativeLlmRequestCodec, + _request: *const NemoRelayNativeString, + out: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + unsafe { test_string_new(c"{}".as_ptr().cast(), 2, out) } +} + +unsafe extern "C" fn successful_request_codec_encode( + _codec: *const NemoRelayNativeLlmRequestCodec, + _annotated: *const NemoRelayNativeString, + original: *const NemoRelayNativeString, + out: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + let bytes = unsafe { &*(original.cast::()) }.0.as_slice(); + unsafe { test_string_new(bytes.as_ptr(), bytes.len(), out) } +} + +unsafe extern "C" fn successful_response_codec_decode( + _codec: *const NemoRelayNativeLlmResponseCodec, + _response: *const NemoRelayNativeString, + out: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + unsafe { test_string_new(c"{}".as_ptr().cast(), 2, out) } +} + fn test_host() -> NemoRelayNativeHostApiV1 { NemoRelayNativeHostApiV1 { abi_version: NEMO_RELAY_NATIVE_ABI_VERSION, @@ -1137,6 +1195,9 @@ fn test_host() -> NemoRelayNativeHostApiV1 { string_free: test_string_free, last_error_clear: test_last_error_clear, last_error_set: test_last_error_set, + llm_request_codec_decode: unavailable_request_codec_decode, + llm_request_codec_encode: unavailable_request_codec_encode, + llm_response_codec_decode: unavailable_response_codec_decode, plugin_context_register_subscriber: capture_register_subscriber, plugin_context_register_tool_sanitize_request_guardrail: capture_tool_json, plugin_context_register_tool_sanitize_response_guardrail: capture_tool_json, @@ -1284,6 +1345,22 @@ fn json_host_string(host: &NemoRelayNativeHostApiV1, value: Json) -> *mut NemoRe host_string(host, &serde_json::to_string(&value).unwrap()) } +fn native_no_codec_context() -> NemoRelayNativeLlmSanitizeRequestContext { + NemoRelayNativeLlmSanitizeRequestContext { + codec_kind: NemoRelayNativeLlmCodecKind::None, + codec_id: ptr::null(), + codec: ptr::null(), + } +} + +fn native_no_response_codec_context() -> NemoRelayNativeLlmSanitizeResponseContext { + NemoRelayNativeLlmSanitizeResponseContext { + codec_kind: NemoRelayNativeLlmCodecKind::None, + codec_id: ptr::null(), + codec: ptr::null(), + } +} + fn read_json_and_free(host: &NemoRelayNativeHostApiV1, value: *mut NemoRelayNativeString) -> Json { let result: Json = serde_json::from_str(&read_host_string(host, value).unwrap()).unwrap(); unsafe { (host.string_free)(value) }; @@ -2802,12 +2879,21 @@ fn typed_callbacks_reject_null_abi_pointers_before_decoding_inputs() { } let mut ctx = test_context(&host); - ctx.register_llm_sanitize_request_guardrail("llm-request", 0, |request| request) - .unwrap(); + ctx.register_llm_sanitize_request_guardrail("llm-request", 0, |request, _context| { + Some(request) + }) + .unwrap(); let registration = take_llm_request_registration(); let request = json_host_string(&host, serde_json::to_value(test_llm_request()).unwrap()); assert_eq!( - unsafe { (registration.cb)(ptr::null_mut(), request, &mut out) }, + unsafe { + (registration.cb)( + ptr::null_mut(), + request, + native_no_codec_context(), + &mut out, + ) + }, NemoRelayStatus::NullPointer ); assert_eq!( @@ -2815,6 +2901,7 @@ fn typed_callbacks_reject_null_abi_pointers_before_decoding_inputs() { (registration.cb)( registration.user_data as *mut c_void, request, + native_no_codec_context(), ptr::null_mut(), ) }, @@ -2826,12 +2913,19 @@ fn typed_callbacks_reject_null_abi_pointers_before_decoding_inputs() { } let mut ctx = test_context(&host); - ctx.register_llm_sanitize_response_guardrail("llm-response", 0, |value| value) + ctx.register_llm_sanitize_response_guardrail("llm-response", 0, |value, _context| Some(value)) .unwrap(); let registration = take_llm_json_registration(); let response = json_host_string(&host, json!({})); assert_eq!( - unsafe { (registration.cb)(ptr::null_mut(), response, &mut out) }, + unsafe { + (registration.cb)( + ptr::null_mut(), + response, + native_no_response_codec_context(), + &mut out, + ) + }, NemoRelayStatus::NullPointer ); assert_eq!( @@ -2839,6 +2933,7 @@ fn typed_callbacks_reject_null_abi_pointers_before_decoding_inputs() { (registration.cb)( registration.user_data as *mut c_void, response, + native_no_response_codec_context(), ptr::null_mut(), ) }, @@ -3166,14 +3261,24 @@ fn typed_callbacks_report_invalid_json_for_each_decoder_family() { } let mut ctx = test_context(&host); - ctx.register_llm_sanitize_request_guardrail("llm-request", 0, |request| request) - .unwrap(); + ctx.register_llm_sanitize_request_guardrail("llm-request", 0, |request, _context| { + Some(request) + }) + .unwrap(); let registration = take_llm_request_registration(); let request = host_string(&host, "{not json"); + let context = native_no_codec_context(); let stale_out = host_string(&host, r#"{"stale":true}"#); let mut out = stale_out; assert_eq!( - unsafe { (registration.cb)(registration.user_data as *mut c_void, request, &mut out) }, + unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + request, + context, + &mut out, + ) + }, NemoRelayStatus::InvalidJson ); assert!(out.is_null()); @@ -3184,14 +3289,22 @@ fn typed_callbacks_report_invalid_json_for_each_decoder_family() { } let mut ctx = test_context(&host); - ctx.register_llm_sanitize_response_guardrail("llm-response", 0, |value| value) + ctx.register_llm_sanitize_response_guardrail("llm-response", 0, |value, _context| Some(value)) .unwrap(); let registration = take_llm_json_registration(); let response = host_string(&host, "{not json"); + let context = native_no_response_codec_context(); let stale_out = host_string(&host, r#"{"stale":true}"#); let mut out = stale_out; assert_eq!( - unsafe { (registration.cb)(registration.user_data as *mut c_void, response, &mut out) }, + unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + response, + context, + &mut out, + ) + }, NemoRelayStatus::InvalidJson ); assert!(out.is_null()); @@ -4073,20 +4186,31 @@ fn typed_llm_sanitize_guardrails_transform_request_and_response() { let _guard = begin_test(); let host = test_host(); let mut ctx = test_context(&host); - ctx.register_llm_sanitize_request_guardrail("llm-sanitize-request", 12, |mut request| { - request.headers.insert("x-policy".into(), json!("sdk")); - request.content["sanitized"] = json!(true); - request - }) + ctx.register_llm_sanitize_request_guardrail( + "llm-sanitize-request", + 12, + |mut request, _context| { + request.headers.insert("x-policy".into(), json!("sdk")); + request.content["sanitized"] = json!(true); + Some(request) + }, + ) .unwrap(); let registration = take_llm_request_registration(); assert_eq!(registration.name, "llm-sanitize-request"); assert_eq!(registration.priority, 12); let request = json_host_string(&host, serde_json::to_value(test_llm_request()).unwrap()); + let context = native_no_codec_context(); let mut out = ptr::null_mut(); - let status = - unsafe { (registration.cb)(registration.user_data as *mut c_void, request, &mut out) }; + let status = unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + request, + context, + &mut out, + ) + }; assert_eq!(status, NemoRelayStatus::Ok); let output = read_json_and_free(&host, out); assert_eq!(output["headers"]["x-policy"], json!("sdk")); @@ -4097,19 +4221,30 @@ fn typed_llm_sanitize_guardrails_transform_request_and_response() { } let mut ctx = test_context(&host); - ctx.register_llm_sanitize_response_guardrail("llm-sanitize-response", 13, |mut payload| { - payload["sanitized"] = json!(true); - payload - }) + ctx.register_llm_sanitize_response_guardrail( + "llm-sanitize-response", + 13, + |mut payload, _context| { + payload["sanitized"] = json!(true); + Some(payload) + }, + ) .unwrap(); let registration = take_llm_json_registration(); assert_eq!(registration.name, "llm-sanitize-response"); assert_eq!(registration.priority, 13); let response = json_host_string(&host, json!({ "output": true })); + let context = native_no_response_codec_context(); let mut out = ptr::null_mut(); - let status = - unsafe { (registration.cb)(registration.user_data as *mut c_void, response, &mut out) }; + let status = unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + response, + context, + &mut out, + ) + }; assert_eq!(status, NemoRelayStatus::Ok); assert_eq!(read_json_and_free(&host, out)["sanitized"], json!(true)); unsafe { @@ -4118,6 +4253,176 @@ fn typed_llm_sanitize_guardrails_transform_request_and_response() { } } +#[test] +fn typed_contextual_llm_sanitize_guardrails_receive_payload_before_context() { + let _guard = begin_test(); + let mut host = test_host(); + host.llm_request_codec_decode = successful_request_codec_decode; + host.llm_request_codec_encode = successful_request_codec_encode; + host.llm_response_codec_decode = successful_response_codec_decode; + let mut ctx = test_context(&host); + ctx.register_llm_sanitize_request_guardrail( + "contextual-request", + 14, + |mut request, callback_context| { + assert_eq!( + callback_context.codec, + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiChat) + ); + let codec = callback_context + .resolve_codec() + .expect("active request codec must resolve"); + let annotated = codec.decode(&request).expect("request decode succeeds"); + request = codec + .encode(&annotated, &request) + .expect("request encode succeeds"); + request.headers.insert("x-contextual".into(), json!(true)); + Some(request) + }, + ) + .unwrap(); + + let registration = take_llm_request_registration(); + assert_eq!(registration.name, "contextual-request"); + assert_eq!(registration.priority, 14); + let request = json_host_string(&host, serde_json::to_value(test_llm_request()).unwrap()); + let context_id = host_string(&host, "openai_chat"); + let request_codec_placeholder = Box::new(0_usize); + let native_context = NemoRelayNativeLlmSanitizeRequestContext { + codec_kind: NemoRelayNativeLlmCodecKind::BuiltIn, + codec_id: context_id, + codec: std::ptr::from_ref(request_codec_placeholder.as_ref()) + .cast::(), + }; + let mut out = ptr::null_mut(); + let status = unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + request, + native_context, + &mut out, + ) + }; + assert_eq!(status, NemoRelayStatus::Ok); + assert_eq!( + read_json_and_free(&host, out)["headers"]["x-contextual"], + json!(true) + ); + unsafe { + (host.string_free)(request); + (host.string_free)(context_id); + registration.free(); + } + + let mut ctx = test_context(&host); + ctx.register_llm_sanitize_response_guardrail( + "contextual-response", + 15, + |mut payload, callback_context| { + assert_eq!( + callback_context.codec, + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiChat) + ); + callback_context + .resolve_codec() + .expect("active response codec must resolve") + .decode(&payload) + .expect("response decode succeeds"); + payload["contextual"] = json!(true); + Some(payload) + }, + ) + .unwrap(); + + let registration = take_llm_json_registration(); + assert_eq!(registration.name, "contextual-response"); + assert_eq!(registration.priority, 15); + let response = json_host_string(&host, json!({ "output": true })); + let context_id = host_string(&host, "openai_chat"); + let response_codec_placeholder = Box::new(0_usize); + let native_context = NemoRelayNativeLlmSanitizeResponseContext { + codec_kind: NemoRelayNativeLlmCodecKind::BuiltIn, + codec_id: context_id, + codec: std::ptr::from_ref(response_codec_placeholder.as_ref()) + .cast::(), + }; + let mut out = ptr::null_mut(); + let status = unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + response, + native_context, + &mut out, + ) + }; + assert_eq!(status, NemoRelayStatus::Ok); + assert_eq!(read_json_and_free(&host, out)["contextual"], json!(true)); + unsafe { + (host.string_free)(response); + (host.string_free)(context_id); + registration.free(); + } +} + +#[test] +fn typed_contextual_llm_sanitizer_uses_null_output_to_omit_payload() { + let _guard = begin_test(); + let host = test_host(); + let mut ctx = test_context(&host); + ctx.register_llm_sanitize_request_guardrail( + "contextual-omit-request", + 16, + |_request, _context| None, + ) + .unwrap(); + + let registration = take_llm_request_registration(); + let request = json_host_string(&host, serde_json::to_value(test_llm_request()).unwrap()); + let context = native_no_codec_context(); + let mut out = ptr::null_mut(); + let status = unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + request, + context, + &mut out, + ) + }; + assert_eq!(status, NemoRelayStatus::Ok); + assert!(out.is_null(), "null native output must represent omission"); + unsafe { + (host.string_free)(request); + registration.free(); + } + + let mut ctx = test_context(&host); + ctx.register_llm_sanitize_response_guardrail("contextual-omit", 16, |_payload, _context| None) + .unwrap(); + + let registration = take_llm_json_registration(); + let response = json_host_string(&host, json!({"secret": "value"})); + let context = NemoRelayNativeLlmSanitizeResponseContext { + codec_kind: NemoRelayNativeLlmCodecKind::None, + codec_id: ptr::null(), + codec: ptr::null(), + }; + let mut out = ptr::null_mut(); + let status = unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + response, + context, + &mut out, + ) + }; + assert_eq!(status, NemoRelayStatus::Ok); + assert!(out.is_null(), "null native output must represent omission"); + unsafe { + (host.string_free)(response); + registration.free(); + } +} + #[test] fn typed_llm_conditional_guardrail_returns_optional_reason() { let _guard = begin_test(); diff --git a/crates/python/src/py_api/mod.rs b/crates/python/src/py_api/mod.rs index bc026d76e..64e2341e3 100644 --- a/crates/python/src/py_api/mod.rs +++ b/crates/python/src/py_api/mod.rs @@ -73,6 +73,26 @@ fn py_llm_response_codec( }) } +fn py_llm_codec(codec: Option<&Bound<'_, PyAny>>) -> Option> { + codec.and_then(|codec| -> Option> { + if codec.is_none() { + return None; + } + if let Ok(builtin) = codec.extract::>() { + return Some(builtin.inner_codec.clone()); + } + if let Ok(builtin) = codec.extract::>() { + return Some(builtin.inner_codec.clone()); + } + if let Ok(builtin) = codec.extract::>() { + return Some(builtin.inner_codec.clone()); + } + Some(Arc::new(py_callable::PyLlmCodecWrapper { + py_codec: codec.clone().unbind(), + })) + }) +} + fn py_annotated_llm_response( annotated_response: Option<&Bound<'_, PyAny>>, ) -> PyResult>> { @@ -750,11 +770,7 @@ fn llm_call_execute<'py>( let exec_fn = py_callable::wrap_py_llm_exec_fn(func); let default_fn: LlmExecutionNextFn = Arc::new(move |req| exec_fn(req)); let parent_handle = handle.map(|h| h.inner).unwrap_or_else(task_scope_top); - let codec_arc: Option> = codec.map(|c| { - Arc::new(py_callable::PyLlmCodecWrapper { - py_codec: c.clone().unbind(), - }) as Arc - }); + let codec_arc = py_llm_codec(codec); let response_codec_arc = py_llm_response_codec(response_codec); let scope_stack = current_scope_stack_handle(); @@ -855,11 +871,7 @@ fn llm_stream_call_execute<'py>( let collector_fn = py_callable::wrap_py_collector_fn(collector); let finalizer_fn = py_callable::wrap_py_finalizer_fn(finalizer); let parent_handle = handle.map(|h| h.inner).unwrap_or_else(task_scope_top); - let codec_arc: Option> = codec.map(|c| { - Arc::new(py_callable::PyLlmCodecWrapper { - py_codec: c.clone().unbind(), - }) as Arc - }); + let codec_arc = py_llm_codec(codec); let response_codec_arc = py_llm_response_codec(response_codec); let scope_stack = current_scope_stack_handle(); @@ -1081,7 +1093,8 @@ fn deregister_tool_execution_intercept(name: &str) -> PyResult { /// Register an LLM sanitize-request guardrail. /// -/// Callback: ``(request: LlmRequest) -> LlmRequest`` — returns a sanitized request. +/// Callback: ``(request: LlmRequest, context: LlmSanitizeRequestContext) -> +/// Optional[LlmRequest]``. Return ``None`` to omit the observability payload and annotation. #[pyfunction] fn register_llm_sanitize_request_guardrail( name: &str, @@ -1091,7 +1104,7 @@ fn register_llm_sanitize_request_guardrail( core_registry_api::register_llm_sanitize_request_guardrail( name, priority, - py_callable::wrap_py_llm_sanitize_request_fn(guardrail), + py_callable::wrap_py_llm_sanitize_request_fn(guardrail)?, ) .map_err(to_py_err) } @@ -1104,7 +1117,8 @@ fn deregister_llm_sanitize_request_guardrail(name: &str) -> PyResult { /// Register an LLM sanitize-response guardrail. /// -/// Callback: ``(response: dict) -> dict`` — returns a sanitized response. +/// Callback: ``(response: Json, context: LlmSanitizeResponseContext) -> Optional[Json]``. +/// Return ``None`` to omit the observability payload and annotation. #[pyfunction] fn register_llm_sanitize_response_guardrail( name: &str, @@ -1114,7 +1128,7 @@ fn register_llm_sanitize_response_guardrail( core_registry_api::register_llm_sanitize_response_guardrail( name, priority, - py_callable::wrap_py_llm_sanitize_response_fn(guardrail), + py_callable::wrap_py_llm_sanitize_response_fn(guardrail)?, ) .map_err(to_py_err) } @@ -1542,7 +1556,7 @@ fn scope_register_llm_sanitize_request_guardrail( &uuid, name, priority, - py_callable::wrap_py_llm_sanitize_request_fn(guardrail), + py_callable::wrap_py_llm_sanitize_request_fn(guardrail)?, ) .map_err(to_py_err) } @@ -1568,7 +1582,7 @@ fn scope_register_llm_sanitize_response_guardrail( &uuid, name, priority, - py_callable::wrap_py_llm_sanitize_response_fn(guardrail), + py_callable::wrap_py_llm_sanitize_response_fn(guardrail)?, ) .map_err(to_py_err) } diff --git a/crates/python/src/py_callable.rs b/crates/python/src/py_callable.rs index b08fb2f2d..57fffd5d5 100644 --- a/crates/python/src/py_callable.rs +++ b/crates/python/src/py_callable.rs @@ -27,8 +27,9 @@ use std::task::{Context, Poll}; use nemo_relay::api::runtime::{ EventSanitizeFn, EventSubscriberFn, LlmConditionalFn, LlmExecutionNextFn, LlmJsonStream, - LlmRequestInterceptFn, LlmSanitizeRequestFn, LlmSanitizeResponseFn, LlmStreamExecutionNextFn, - LlmStreamInner, ToolConditionalFn, ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, + LlmRequestInterceptFn, LlmSanitizeRequestContext, LlmSanitizeRequestFn, + LlmSanitizeResponseContext, LlmSanitizeResponseFn, LlmStreamExecutionNextFn, LlmStreamInner, + ToolConditionalFn, ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, }; use nemo_relay::error::{FlowError, Result as FlowResult}; use pyo3::prelude::*; @@ -47,11 +48,31 @@ use nemo_relay::codec::traits::{LlmCodec, LlmResponseCodec}; use crate::convert::{json_to_py, py_to_json}; use crate::py_types::{ PyAnnotatedLLMRequest, PyAnnotatedLLMResponse, PyLLMRequest, PyLLMRequestInterceptOutcome, - PyToolExecutionInterceptOutcome, + PyLlmSanitizeRequestContext, PyLlmSanitizeResponseContext, PyToolExecutionInterceptOutcome, }; type PyValueFuture = Pin>> + Send>>; +fn validate_python_llm_sanitizer_signature(py_fn: &Py) -> PyResult<()> { + Python::attach(|py| { + let inspect = py.import("inspect")?; + let signature = inspect.call_method1("signature", (py_fn.bind(py),)).map_err(|_| { + pyo3::exceptions::PyTypeError::new_err( + "LLM sanitizer callback signature cannot be inspected; use a callable that accepts `(payload, context)`", + ) + })?; + if signature + .call_method1("bind", (py.None(), py.None())) + .is_ok() + { + return Ok(()); + } + Err(pyo3::exceptions::PyTypeError::new_err( + "LLM sanitizer callback must accept `(payload, context)`", + )) + }) +} + fn split_json_or_future( py: Python<'_>, result: Py, @@ -792,33 +813,35 @@ pub fn wrap_py_llm_stream_exec_intercept_fn( ) } -/// Wrap a Python callable `(LlmRequest) -> LlmRequest` for LLM sanitize request guardrails. -pub fn wrap_py_llm_sanitize_request_fn(py_fn: Py) -> LlmSanitizeRequestFn { - Arc::new(move |request: LlmRequest| { - Python::attach(|py| { - let py_req = PyLLMRequest { - inner: request.clone(), - }; - let result = match py_fn.call1(py, (py_req,)) { - Ok(v) => v, - Err(e) => { - eprintln!("nemo_relay: LLM sanitize request guardrail callable failed: {e}"); - return request; +/// Wrap a Python callable `(LlmRequest, LlmSanitizeRequestContext) -> Optional`. +fn wrap_py_llm_sanitize_request_callback(py_fn: Py) -> LlmSanitizeRequestFn { + Arc::new( + move |request: LlmRequest, context: LlmSanitizeRequestContext| { + Python::attach(|py| { + let py_context = PyLlmSanitizeRequestContext { inner: context }; + let py_request = PyLLMRequest { inner: request }; + let result = match py_fn.call1(py, (py_request, py_context)) { + Ok(value) => value, + Err(error) => { + eprintln!("nemo_relay: LLM sanitize request callable failed: {error}"); + return None; + } + }; + if result.is_none(py) { + return None; } - }; - let extracted = result.extract::(py); - match extracted { - Ok(r) => r.inner, - Err(e) => { - eprintln!( - "nemo_relay: LLM sanitize request guardrail returned unexpected type \ - (expected LlmRequest): {e}" - ); - request + match result.extract::(py) { + Ok(request) => Some(request.inner), + Err(error) => { + eprintln!( + "nemo_relay: LLM sanitize request callable returned unexpected type: {error}" + ); + None + } } - } - }) - }) + }) + }, + ) } /// Wrap a Python callable `(LlmRequest) -> Optional[str]` for LLM conditional guardrails. @@ -989,34 +1012,51 @@ pub fn wrap_py_finalizer_fn(py_fn: Py) -> Box Json + Send }) } -/// Wrap a Python callable `(dict) -> dict` for LLM sanitize response guardrails. -pub fn wrap_py_llm_sanitize_response_fn(py_fn: Py) -> LlmSanitizeResponseFn { - Arc::new(move |response: Json| { +/// Wrap a Python callable `(Json, LlmSanitizeResponseContext) -> Optional[Json]`. +fn wrap_py_llm_sanitize_response_callback(py_fn: Py) -> LlmSanitizeResponseFn { + Arc::new(move |response: Json, context: LlmSanitizeResponseContext| { Python::attach(|py| { - let py_resp = match json_to_py(py, &response) { - Ok(v) => v, - Err(e) => { - eprintln!( - "nemo_relay: json_to_py failed in LLM sanitize response guardrail: {e}" - ); - return response.clone(); + let py_context = PyLlmSanitizeResponseContext { inner: context }; + let py_response = match json_to_py(py, &response) { + Ok(value) => value, + Err(error) => { + eprintln!("nemo_relay: json_to_py failed in LLM sanitize response: {error}"); + return None; } }; - let result = match py_fn.call1(py, (py_resp,)) { - Ok(v) => v, - Err(e) => { - eprintln!("nemo_relay: LLM sanitize response guardrail callable failed: {e}"); - return response.clone(); + let result = match py_fn.call1(py, (py_response, py_context)) { + Ok(value) => value, + Err(error) => { + eprintln!("nemo_relay: LLM sanitize response callable failed: {error}"); + return None; } }; - py_to_json(result.bind(py)).unwrap_or_else(|e| { - eprintln!("nemo_relay: py_to_json failed in LLM sanitize response guardrail: {e}"); - response.clone() - }) + if result.is_none(py) { + return None; + } + match py_to_json(result.bind(py)) { + Ok(response) => Some(response), + Err(error) => { + eprintln!("nemo_relay: py_to_json failed in LLM sanitize response: {error}"); + None + } + } }) }) } +/// Wrap a Python LLM sanitize-request callback. +pub fn wrap_py_llm_sanitize_request_fn(py_fn: Py) -> PyResult { + validate_python_llm_sanitizer_signature(&py_fn)?; + Ok(wrap_py_llm_sanitize_request_callback(py_fn)) +} + +/// Wrap a Python LLM sanitize-response callback. +pub fn wrap_py_llm_sanitize_response_fn(py_fn: Py) -> PyResult { + validate_python_llm_sanitizer_signature(&py_fn)?; + Ok(wrap_py_llm_sanitize_response_callback(py_fn)) +} + /// Wrap a Python callable `(Event) -> None` for event subscribers. pub fn wrap_py_event_subscriber(py_fn: Py) -> EventSubscriberFn { Arc::new(move |event: &Event| { diff --git a/crates/python/src/py_plugin.rs b/crates/python/src/py_plugin.rs index e709fc769..2cf76a4a8 100644 --- a/crates/python/src/py_plugin.rs +++ b/crates/python/src/py_plugin.rs @@ -388,14 +388,11 @@ impl PyPluginContext { priority: i32, callback: Py, ) -> PyResult<()> { + let callback = wrap_py_llm_sanitize_request_fn(callback)?; self.register_callback( name, |qualified_name| { - register_llm_sanitize_request_guardrail( - qualified_name, - priority, - wrap_py_llm_sanitize_request_fn(callback), - ) + register_llm_sanitize_request_guardrail(qualified_name, priority, callback) }, deregister_llm_sanitize_request_guardrail, "llm sanitize request guardrail", @@ -409,14 +406,11 @@ impl PyPluginContext { priority: i32, callback: Py, ) -> PyResult<()> { + let callback = wrap_py_llm_sanitize_response_fn(callback)?; self.register_callback( name, |qualified_name| { - register_llm_sanitize_response_guardrail( - qualified_name, - priority, - wrap_py_llm_sanitize_response_fn(callback), - ) + register_llm_sanitize_response_guardrail(qualified_name, priority, callback) }, deregister_llm_sanitize_response_guardrail, "llm sanitize response guardrail", diff --git a/crates/python/src/py_types/codecs.rs b/crates/python/src/py_types/codecs.rs index 73efdf5a8..1b3a22875 100644 --- a/crates/python/src/py_types/codecs.rs +++ b/crates/python/src/py_types/codecs.rs @@ -26,6 +26,53 @@ use super::{ }; use nemo_relay::codec::response::FinishReason; +/// A resolved request codec available while an LLM request sanitizer runs. +#[pyclass(name = "LlmSanitizeRequestCodec")] +pub struct PyLlmSanitizeRequestCodec { + pub(crate) inner: Arc, +} + +#[pymethods] +impl PyLlmSanitizeRequestCodec { + /// Parse an opaque request into its normalized representation. + fn decode(&self, request: &PyLLMRequest) -> PyResult { + self.inner + .decode(&request.inner) + .map(|inner| PyAnnotatedLLMRequest { inner }) + .map_err(|error| pyo3::exceptions::PyRuntimeError::new_err(error.to_string())) + } + + /// Merge a normalized request back into its provider representation. + fn encode( + &self, + annotated: &PyAnnotatedLLMRequest, + original: &PyLLMRequest, + ) -> PyResult { + self.inner + .encode(&annotated.inner, &original.inner) + .map(|inner| PyLLMRequest { inner }) + .map_err(|error| pyo3::exceptions::PyRuntimeError::new_err(error.to_string())) + } +} + +/// A resolved response codec available while an LLM response sanitizer runs. +#[pyclass(name = "LlmSanitizeResponseCodec")] +pub struct PyLlmSanitizeResponseCodec { + pub(crate) inner: Arc, +} + +#[pymethods] +impl PyLlmSanitizeResponseCodec { + /// Parse an opaque response into its normalized representation. + fn decode_response(&self, response: &Bound<'_, PyAny>) -> PyResult { + let response = py_to_json(response)?; + self.inner + .decode_response(&response) + .map(|inner| PyAnnotatedLLMResponse { inner }) + .map_err(|error| pyo3::exceptions::PyRuntimeError::new_err(error.to_string())) + } +} + // --------------------------------------------------------------------------- // AnnotatedLLMRequest // --------------------------------------------------------------------------- diff --git a/crates/python/src/py_types/core.rs b/crates/python/src/py_types/core.rs index aeca45305..d0f122311 100644 --- a/crates/python/src/py_types/core.rs +++ b/crates/python/src/py_types/core.rs @@ -5,15 +5,96 @@ use std::sync::Arc; use pyo3::prelude::*; +use super::codecs::{PyLlmSanitizeRequestCodec, PyLlmSanitizeResponseCodec}; use super::{ - AnnotatedLLMRequest, Bound, CoreScopeType, FlowResult, LlmAttributes, LlmHandle, LlmRequest, - PyAnnotatedLLMRequest, PyAny, PyErr, PyRef, PyResult, Python, ScopeAttributes, ScopeHandle, - ScopeStackHandle, ToolAttributes, ToolHandle, json_to_py, opt_json_to_py, py_to_json, + AnnotatedLLMRequest, Bound, CoreScopeType, FlowResult, LlmAttributes, LlmCodecIdentity, + LlmHandle, LlmRequest, PyAnnotatedLLMRequest, PyAny, PyErr, PyRef, PyResult, Python, + ScopeAttributes, ScopeHandle, ScopeStackHandle, ToolAttributes, ToolHandle, json_to_py, + opt_json_to_py, py_to_json, }; use nemo_relay::api::event::{CategoryProfile, EventCategory, PendingMarkSpec}; use nemo_relay::api::llm::LlmRequestInterceptOutcome; +use nemo_relay::api::runtime::{LlmSanitizeRequestContext, LlmSanitizeResponseContext}; use nemo_relay::api::tool::ToolExecutionInterceptOutcome; +/// Structured identity of the codec active during LLM sanitization. +#[pyclass(name = "LlmCodecIdentity", frozen)] +pub struct PyLlmCodecIdentity { + pub(crate) inner: LlmCodecIdentity, +} + +#[pymethods] +impl PyLlmCodecIdentity { + /// Identity variant: ``none``, ``builtin``, ``runtime``, or ``opaque``. + #[getter] + fn kind(&self) -> &'static str { + match self.inner { + LlmCodecIdentity::None => "none", + LlmCodecIdentity::BuiltIn(_) => "builtin", + LlmCodecIdentity::Runtime(_) => "runtime", + LlmCodecIdentity::Opaque => "opaque", + } + } + + /// Stable built-in or runtime codec ID, when this identity has one. + #[getter] + fn id(&self) -> Option { + match &self.inner { + LlmCodecIdentity::BuiltIn(codec) => Some(codec.id().to_owned()), + LlmCodecIdentity::Runtime(id) => Some(id.clone()), + LlmCodecIdentity::None | LlmCodecIdentity::Opaque => None, + } + } +} + +/// Structured per-call context delivered to LLM request sanitizer callbacks. +#[pyclass(name = "LlmSanitizeRequestContext", frozen)] +pub struct PyLlmSanitizeRequestContext { + pub(crate) inner: LlmSanitizeRequestContext, +} + +#[pymethods] +impl PyLlmSanitizeRequestContext { + /// The active codec identity for this request or response payload. + #[getter] + fn codec(&self) -> PyLlmCodecIdentity { + PyLlmCodecIdentity { + inner: self.inner.codec().clone(), + } + } + + /// Resolve the active request codec. + fn resolve_codec(&self) -> Option { + self.inner + .resolve_codec() + .map(|inner| PyLlmSanitizeRequestCodec { inner }) + } +} + +/// Structured per-call context delivered to LLM response sanitizer callbacks. +#[pyclass(name = "LlmSanitizeResponseContext", frozen)] +pub struct PyLlmSanitizeResponseContext { + pub(crate) inner: LlmSanitizeResponseContext, +} + +#[pymethods] +impl PyLlmSanitizeResponseContext { + /// The active codec identity for this request or response payload. + #[getter] + fn codec(&self) -> PyLlmCodecIdentity { + PyLlmCodecIdentity { + inner: self.inner.codec().clone(), + } + } + + /// Resolve the active response codec. + fn resolve_codec(&self) -> Option { + self.inner + .resolve_codec() + .map(|inner| PyLlmSanitizeResponseCodec { inner }) + } +} + // --------------------------------------------------------------------------- // LlmStream (async iterator) // --------------------------------------------------------------------------- diff --git a/crates/python/src/py_types/mod.rs b/crates/python/src/py_types/mod.rs index 046b580ff..eb60374d9 100644 --- a/crates/python/src/py_types/mod.rs +++ b/crates/python/src/py_types/mod.rs @@ -15,7 +15,7 @@ use std::time::Duration; use nemo_relay::api::event::{MarkEvent, ScopeEvent}; use nemo_relay::api::llm::{LlmAttributes, LlmHandle, LlmRequest}; -use nemo_relay::api::runtime::ScopeStackHandle; +use nemo_relay::api::runtime::{LlmCodecIdentity, ScopeStackHandle}; use nemo_relay::api::scope::{ScopeAttributes, ScopeHandle, ScopeType as CoreScopeType}; use nemo_relay::api::tool::{ToolAttributes, ToolHandle}; use nemo_relay::codec::request::{ @@ -128,6 +128,14 @@ pub use events::*; pub use observability::*; pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> { + register_runtime_types(m)?; + register_llm_types(m)?; + register_event_types(m)?; + register_observability_types(m)?; + register_codec_types(m) +} + +fn register_runtime_types(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; @@ -138,13 +146,30 @@ pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + Ok(()) +} + +fn register_llm_types(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; + Ok(()) +} + +fn register_event_types(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; + Ok(()) +} + +fn register_observability_types(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; @@ -154,6 +179,10 @@ pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + Ok(()) +} + +fn register_codec_types(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/crates/python/tests/coverage/coverage_tests.rs b/crates/python/tests/coverage/coverage_tests.rs index 765c3850a..67cebd3ba 100644 --- a/crates/python/tests/coverage/coverage_tests.rs +++ b/crates/python/tests/coverage/coverage_tests.rs @@ -4,6 +4,7 @@ //! Coverage tests for coverage in the NeMo Relay Python crate. use std::ffi::CString; +use std::path::PathBuf; use std::sync::Arc; use pyo3::prelude::*; @@ -51,6 +52,29 @@ fn make_request() -> LlmRequest { } } +/// Restores the process working directory after a plugin-layering test. +/// +/// Plugin discovery intentionally walks from the working directory, so this +/// guard keeps coverage tests independent of a developer's parent project +/// configuration. The Python test guard serializes these process-global test +/// changes. +struct CurrentDirectoryGuard(PathBuf); + +impl CurrentDirectoryGuard { + fn move_to_temporary_directory() -> Self { + let original = std::env::current_dir().expect("test process has a working directory"); + std::env::set_current_dir(std::env::temp_dir()) + .expect("temporary directory is available for plugin config discovery"); + Self(original) + } +} + +impl Drop for CurrentDirectoryGuard { + fn drop(&mut self) { + std::env::set_current_dir(&self.0).expect("restore test working directory"); + } +} + fn with_event_loop(py: Python<'_>, f: impl FnOnce(Bound<'_, PyAny>) -> T) -> T { let asyncio = py.import("asyncio").unwrap(); #[cfg(windows)] @@ -346,6 +370,7 @@ fn test_py_adaptive_binding_rejects_zero_sensitivity() { fn test_plugin_bindings_validate_configure_and_clear() { let _python = crate::test_support::init_python_test(); let _plugin_test_state = crate::py_plugin::lock_plugin_test_state_for_tests(); + let _working_directory = CurrentDirectoryGuard::move_to_temporary_directory(); Python::attach(|py| { nemo_relay_adaptive::plugin_component::register_adaptive_component().unwrap(); @@ -405,10 +430,10 @@ def tool_passthrough(name, value): def tool_conditional(name, value): return None -def llm_sanitize_request(request): +def llm_sanitize_request(request, context): return request -def llm_sanitize_response(response): +def llm_sanitize_response(response, context): return response def llm_conditional(request): @@ -610,7 +635,7 @@ def tool_fail(name, args): def tool_cond_bad(name, args): return 123 -def llm_sanitize_bad(request): +def llm_sanitize_bad(request, context): return {"bad": True} def llm_cond_bad(request): @@ -622,7 +647,7 @@ def llm_cond_none(request): def llm_req_bad(name, request): return {"bad": True} -def llm_resp_fail(response): +def llm_resp_fail(response, context): raise RuntimeError("resp boom") def collector_fail(chunk): @@ -656,8 +681,15 @@ def event_fail(event): let request = make_request(); let llm_sanitize = - wrap_py_llm_sanitize_request_fn(module.getattr("llm_sanitize_bad").unwrap().unbind()); - assert_eq!(llm_sanitize(request.clone()).content, request.content); + wrap_py_llm_sanitize_request_fn(module.getattr("llm_sanitize_bad").unwrap().unbind()) + .unwrap(); + assert_eq!( + llm_sanitize( + request.clone(), + nemo_relay::api::runtime::LlmSanitizeRequestContext::default(), + ), + None + ); let llm_cond = wrap_py_llm_conditional_fn(module.getattr("llm_cond_bad").unwrap().unbind()); assert!( @@ -689,8 +721,15 @@ def event_fail(event): ); let llm_resp = - wrap_py_llm_sanitize_response_fn(module.getattr("llm_resp_fail").unwrap().unbind()); - assert_eq!(llm_resp(json!({"ok": true})), json!({"ok": true})); + wrap_py_llm_sanitize_response_fn(module.getattr("llm_resp_fail").unwrap().unbind()) + .unwrap(); + assert_eq!( + llm_resp( + json!({"ok": true}), + nemo_relay::api::runtime::LlmSanitizeResponseContext::default(), + ), + None + ); let mut collector = wrap_py_collector_fn(module.getattr("collector_fail").unwrap().unbind()); diff --git a/crates/python/tests/coverage/py_api_coverage_tests.rs b/crates/python/tests/coverage/py_api_coverage_tests.rs index 0dee9e556..ec4c74fec 100644 --- a/crates/python/tests/coverage/py_api_coverage_tests.rs +++ b/crates/python/tests/coverage/py_api_coverage_tests.rs @@ -179,10 +179,10 @@ async def tool_exec_intercept(name, args, next): result["tool_intercepted"] = True return ToolExecutionInterceptOutcome(result) -def llm_sanitize_request(request): +def llm_sanitize_request(request, context): return request -def llm_sanitize_response(response): +def llm_sanitize_response(response, context): updated = dict(response) updated["llm_sanitized_response"] = True return updated diff --git a/crates/python/tests/coverage/py_callable_coverage_tests.rs b/crates/python/tests/coverage/py_callable_coverage_tests.rs index e6d8e381a..fe815cc1b 100644 --- a/crates/python/tests/coverage/py_callable_coverage_tests.rs +++ b/crates/python/tests/coverage/py_callable_coverage_tests.rs @@ -73,7 +73,7 @@ def collector_ok(chunk): def finalizer_bad_json(): return object() -def llm_resp_bad_json(response): +def llm_resp_bad_json(response, context): return object() class BadCodec: @@ -187,8 +187,15 @@ class RaisingResponseCodec: assert_eq!(finalizer(), serde_json::Value::Null); let llm_response = - wrap_py_llm_sanitize_response_fn(module.getattr("llm_resp_bad_json").unwrap().unbind()); - assert_eq!(llm_response(json!({"ok": true})), json!({"ok": true})); + wrap_py_llm_sanitize_response_fn(module.getattr("llm_resp_bad_json").unwrap().unbind()) + .unwrap(); + assert_eq!( + llm_response( + json!({"ok": true}), + nemo_relay::api::runtime::LlmSanitizeResponseContext::default() + ), + None + ); let bad_codec = PyLlmCodecWrapper { py_codec: module diff --git a/crates/python/tests/coverage/py_plugin_coverage_tests.rs b/crates/python/tests/coverage/py_plugin_coverage_tests.rs index c17b7e469..83bc2c9b5 100644 --- a/crates/python/tests/coverage/py_plugin_coverage_tests.rs +++ b/crates/python/tests/coverage/py_plugin_coverage_tests.rs @@ -67,6 +67,36 @@ fn plugin_context_helpers_and_error_conversion_work() { assert!(err.to_string().contains("boom")); } +#[test] +fn plugin_context_rejects_legacy_and_uninspectable_llm_sanitizers() { + let _python = crate::test_support::init_python_test(); + let context = PyPluginContext { + registrations: Arc::new(Mutex::new(vec![])), + namespace_prefix: "invalid.".to_string(), + }; + + Python::attach(|py| { + let helpers = load_module( + py, + r#" +def one_argument(payload): + return payload +"#, + ); + for callback in [helpers.getattr("one_argument").unwrap().unbind(), py.None()] { + let request_error = context + .register_llm_sanitize_request_guardrail("request", 1, callback.clone_ref(py)) + .unwrap_err(); + assert!(request_error.to_string().contains("payload, context")); + + let response_error = context + .register_llm_sanitize_response_guardrail("response", 1, callback) + .unwrap_err(); + assert!(response_error.to_string().contains("payload, context")); + } + }); +} + #[test] fn register_adds_plugin_management_bindings() { let _python = crate::test_support::init_python_test(); @@ -148,10 +178,10 @@ def tool_fn(name, value): def tool_conditional(name, value): return None -def llm_sanitize_request(request): +def llm_sanitize_request(request, context): return request -def llm_sanitize_response(response): +def llm_sanitize_response(response, context): return response def llm_conditional(request): @@ -628,10 +658,10 @@ def tool_fn(name, value): def tool_conditional(name, value): return None -def llm_sanitize_request(request): +def llm_sanitize_request(request, context): return request -def llm_sanitize_response(response): +def llm_sanitize_response(response, context): return response def llm_conditional(request): @@ -912,10 +942,10 @@ def tool_fn(name, value): def tool_conditional(name, value): return None -def llm_sanitize_request(request): +def llm_sanitize_request(request, context): return request -def llm_sanitize_response(response): +def llm_sanitize_response(response, context): return response def llm_conditional(request): diff --git a/crates/worker-proto/proto/nemo/relay/worker/v1/plugin_worker.proto b/crates/worker-proto/proto/nemo/relay/worker/v1/plugin_worker.proto index 0d17d8b6b..75307e56d 100644 --- a/crates/worker-proto/proto/nemo/relay/worker/v1/plugin_worker.proto +++ b/crates/worker-proto/proto/nemo/relay/worker/v1/plugin_worker.proto @@ -25,6 +25,9 @@ service RelayHostRuntime { rpc ToolNext(ToolNextRequest) returns (JsonResult); rpc LlmNext(LlmNextRequest) returns (JsonResult); rpc LlmStreamNext(LlmStreamNextRequest) returns (stream StreamChunk); + rpc DecodeLlmCodecRequest(LlmCodecDecodeRequest) returns (JsonResult); + rpc EncodeLlmCodecRequest(LlmCodecEncodeRequest) returns (JsonResult); + rpc DecodeLlmCodecResponse(LlmCodecDecodeResponse) returns (JsonResult); } message JsonEnvelope { @@ -51,6 +54,13 @@ enum RegistrationSurface { SCOPE_SANITIZE_END_GUARDRAIL = 32; } +enum LlmCodecKind { + LLM_CODEC_KIND_UNSPECIFIED = 0; + LLM_CODEC_KIND_BUILTIN = 1; + LLM_CODEC_KIND_RUNTIME = 2; + LLM_CODEC_KIND_OPAQUE = 3; +} + enum ScopeType { SCOPE_TYPE_UNSPECIFIED = 0; AGENT = 1; @@ -132,6 +142,7 @@ message Registration { RegistrationSurface surface = 2; int32 priority = 3; bool break_chain = 4; + reserved 5; } message InvokeRequest { @@ -160,6 +171,28 @@ message LlmInvocation { JsonEnvelope request = 2; JsonEnvelope annotated_request = 3; JsonEnvelope response = 4; + reserved 5, 6, 7, 8; + oneof sanitize_context { + LlmSanitizeRequestContext request_sanitize_context = 9; + LlmSanitizeResponseContext response_sanitize_context = 10; + } +} + +message LlmCodecIdentity { + LlmCodecKind kind = 1; + optional string id = 2; +} + +message LlmSanitizeRequestContext { + LlmCodecIdentity codec = 1; + // Opaque, invocation-scoped host capability. SDKs must not expose this value. + optional string codec_capability_id = 2; +} + +message LlmSanitizeResponseContext { + LlmCodecIdentity codec = 1; + // Opaque, invocation-scoped host capability. SDKs must not expose this value. + optional string codec_capability_id = 2; } message InvokeResponse { @@ -302,3 +335,28 @@ message LlmStreamNextRequest { string continuation_id = 3; JsonEnvelope request = 4; } + +message LlmCodecDecodeRequest { + string activation_id = 1; + string auth_token = 2; + string codec_capability_id = 3; + JsonEnvelope request = 4; + string invocation_id = 5; +} + +message LlmCodecEncodeRequest { + string activation_id = 1; + string auth_token = 2; + string codec_capability_id = 3; + JsonEnvelope annotated_request = 4; + JsonEnvelope original_request = 5; + string invocation_id = 6; +} + +message LlmCodecDecodeResponse { + string activation_id = 1; + string auth_token = 2; + string codec_capability_id = 3; + JsonEnvelope response = 4; + string invocation_id = 5; +} diff --git a/crates/worker/src/lib.rs b/crates/worker/src/lib.rs index 982501b64..f2e523f4a 100644 --- a/crates/worker/src/lib.rs +++ b/crates/worker/src/lib.rs @@ -52,7 +52,8 @@ use nemo_relay_worker_proto::v1::relay_host_runtime_client::RelayHostRuntimeClie use nemo_relay_worker_proto::v1::{ CancelInvocationRequest, CreateScopeStackRequest, DropScopeStackRequest, EmitMarkRequest, EmptyResult, GuardrailResult, HandshakeRequest, HandshakeResponse, HealthRequest, - HealthResponse, InvokeRequest, InvokeResponse, JsonEnvelope, JsonResult, LlmNextRequest, + HealthResponse, InvokeRequest, InvokeResponse, JsonEnvelope, JsonResult, LlmCodecDecodeRequest, + LlmCodecDecodeResponse, LlmCodecEncodeRequest, LlmCodecKind, LlmNextRequest, LlmRequestInterceptResult, LlmStreamNextRequest, PopScopeRequest, PushScopeRequest, RegisterRequest, RegisterResponse, Registration, RegistrationSurface, ScopeContext, ShutdownRequest, StreamChunk, ToolExecutionInterceptResult, ToolNextRequest, ValidateRequest, @@ -80,6 +81,9 @@ pub type BoxFutureResult = Pin> + Send>>; /// Boxed JSON stream returned by streaming worker callbacks. pub type JsonStream = Pin> + Send>>; +const JSON_SCHEMA: &str = "nemo.relay.Json@1"; +const LLM_REQUEST_SCHEMA: &str = "nemo.relay.LlmRequest@1"; + tokio::task_local! { static TASK_SCOPE_CONTEXT: Option; } @@ -126,15 +130,137 @@ pub trait WorkerPlugin: Send + Sync + 'static { type SubscriberFn = Arc; type EventSanitizeFn = - Arc EventSanitizeFields + Send + Sync>; -type ToolSanitizeFn = Arc Json + Send + Sync>; + Arc BoxFutureResult + Send + Sync>; +type ToolSanitizeFn = Arc BoxFutureResult + Send + Sync>; type ToolConditionalFn = Arc Result> + Send + Sync>; type ToolRequestFn = Arc Result + Send + Sync>; type ToolExecutionFn = Arc< dyn Fn(&str, Json, ToolNext) -> BoxFutureResult + Send + Sync, >; -type LlmSanitizeRequestFn = Arc LlmRequest + Send + Sync>; -type LlmSanitizeResponseFn = Arc Json + Send + Sync>; +type LlmSanitizeRequestFn = Arc< + dyn Fn(LlmRequest, LlmSanitizeRequestContext) -> BoxFutureResult> + + Send + + Sync, +>; +type LlmSanitizeResponseFn = + Arc BoxFutureResult> + Send + Sync>; + +/// Relay built-in codec identities supplied to worker sanitizers. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BuiltinLlmCodec { + /// OpenAI Chat Completions. + OpenAiChat, + /// OpenAI Responses. + OpenAiResponses, + /// Anthropic Messages. + AnthropicMessages, +} + +/// Per-call LLM codec identity supplied to worker sanitizers. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum LlmCodecIdentity { + /// No codec was active. + None, + /// A Relay built-in codec was active. + BuiltIn(BuiltinLlmCodec), + /// A runtime-registered codec was active, identified by its stable ID. + Runtime(String), + /// A codec was active but has no registered identity. + Opaque, +} + +/// Active codec context supplied to an LLM request sanitizer. +#[derive(Clone)] +pub struct LlmSanitizeRequestContext { + /// Identity of the active codec. + pub codec: LlmCodecIdentity, + runtime: Option, + codec_capability_id: Option, + invocation_id: Option, +} + +/// Active codec context supplied to an LLM response sanitizer. +#[derive(Clone)] +pub struct LlmSanitizeResponseContext { + /// Identity of the active codec. + pub codec: LlmCodecIdentity, + runtime: Option, + codec_capability_id: Option, + invocation_id: Option, +} + +impl LlmSanitizeRequestContext { + /// Resolves the active request codec for this callback. + #[must_use] + pub fn resolve_codec(&self) -> Option { + Some(WorkerRequestCodec { + runtime: self.runtime.clone()?, + capability_id: self.codec_capability_id.clone()?, + invocation_id: self.invocation_id.clone()?, + }) + } +} + +impl LlmSanitizeResponseContext { + /// Resolves the active response codec for this callback. + #[must_use] + pub fn resolve_codec(&self) -> Option { + Some(WorkerResponseCodec { + runtime: self.runtime.clone()?, + capability_id: self.codec_capability_id.clone()?, + invocation_id: self.invocation_id.clone()?, + }) + } +} + +/// Invocation-scoped proxy for the active LLM request codec. +#[derive(Clone)] +pub struct WorkerRequestCodec { + runtime: PluginRuntime, + capability_id: String, + invocation_id: String, +} + +impl WorkerRequestCodec { + /// Decodes an opaque request into its normalized representation. + pub async fn decode(&self, request: &LlmRequest) -> Result { + self.runtime + .decode_llm_codec_request(&self.capability_id, &self.invocation_id, request) + .await + } + /// Encodes normalized request changes onto the original opaque request. + pub async fn encode( + &self, + annotated: &AnnotatedLlmRequest, + original: &LlmRequest, + ) -> Result { + self.runtime + .encode_llm_codec_request( + &self.capability_id, + &self.invocation_id, + annotated, + original, + ) + .await + } +} + +/// Invocation-scoped proxy for the active LLM response codec. +#[derive(Clone)] +pub struct WorkerResponseCodec { + runtime: PluginRuntime, + capability_id: String, + invocation_id: String, +} + +impl WorkerResponseCodec { + /// Decodes an opaque response into its normalized representation. + pub async fn decode(&self, response: &Json) -> Result { + self.runtime + .decode_llm_codec_response(&self.capability_id, &self.invocation_id, response) + .await + } +} type LlmConditionalFn = Arc Result> + Send + Sync>; type LlmRequestFn = Arc< dyn Fn(&str, LlmRequest, Option) -> Result @@ -204,14 +330,15 @@ impl PluginContext { .insert(name.into(), Arc::new(callback)); } - fn register_event_sanitizer( + fn register_event_sanitizer( &mut self, name: &str, priority: i32, surface: RegistrationSurface, callback: F, ) where - F: Fn(&Event, EventSanitizeFields) -> EventSanitizeFields + Send + Sync + 'static, + F: Fn(&Event, EventSanitizeFields) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, { self.push_registration(name, surface, priority, false); let sanitizers = match surface { @@ -224,13 +351,21 @@ impl PluginContext { } _ => unreachable!("event sanitizer registration requires an event sanitizer surface"), }; - sanitizers.insert(name.into(), Arc::new(callback)); + sanitizers.insert( + name.into(), + Arc::new(move |event, fields| Box::pin(callback(event, fields))), + ); } /// Registers a mark event sanitizer. - pub fn register_mark_sanitize_guardrail(&mut self, name: &str, priority: i32, callback: F) - where - F: Fn(&Event, EventSanitizeFields) -> EventSanitizeFields + Send + Sync + 'static, + pub fn register_mark_sanitize_guardrail( + &mut self, + name: &str, + priority: i32, + callback: F, + ) where + F: Fn(&Event, EventSanitizeFields) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, { self.register_event_sanitizer( name, @@ -241,13 +376,14 @@ impl PluginContext { } /// Registers a scope-start event sanitizer. - pub fn register_scope_sanitize_start_guardrail( + pub fn register_scope_sanitize_start_guardrail( &mut self, name: &str, priority: i32, callback: F, ) where - F: Fn(&Event, EventSanitizeFields) -> EventSanitizeFields + Send + Sync + 'static, + F: Fn(&Event, EventSanitizeFields) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, { self.register_event_sanitizer( name, @@ -258,13 +394,14 @@ impl PluginContext { } /// Registers a scope-end event sanitizer. - pub fn register_scope_sanitize_end_guardrail( + pub fn register_scope_sanitize_end_guardrail( &mut self, name: &str, priority: i32, callback: F, ) where - F: Fn(&Event, EventSanitizeFields) -> EventSanitizeFields + Send + Sync + 'static, + F: Fn(&Event, EventSanitizeFields) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, { self.register_event_sanitizer( name, @@ -275,13 +412,14 @@ impl PluginContext { } /// Registers a tool sanitize-request guardrail. - pub fn register_tool_sanitize_request_guardrail( + pub fn register_tool_sanitize_request_guardrail( &mut self, name: &str, priority: i32, callback: F, ) where - F: Fn(&str, Json) -> Json + Send + Sync + 'static, + F: Fn(&str, Json) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, { self.push_registration( name, @@ -289,19 +427,21 @@ impl PluginContext { priority, false, ); - self.handlers - .tool_sanitize_requests - .insert(name.into(), Arc::new(callback)); + self.handlers.tool_sanitize_requests.insert( + name.into(), + Arc::new(move |tool_name, value| Box::pin(callback(tool_name, value))), + ); } /// Registers a tool sanitize-response guardrail. - pub fn register_tool_sanitize_response_guardrail( + pub fn register_tool_sanitize_response_guardrail( &mut self, name: &str, priority: i32, callback: F, ) where - F: Fn(&str, Json) -> Json + Send + Sync + 'static, + F: Fn(&str, Json) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, { self.push_registration( name, @@ -309,9 +449,10 @@ impl PluginContext { priority, false, ); - self.handlers - .tool_sanitize_responses - .insert(name.into(), Arc::new(callback)); + self.handlers.tool_sanitize_responses.insert( + name.into(), + Arc::new(move |tool_name, value| Box::pin(callback(tool_name, value))), + ); } /// Registers a tool conditional-execution guardrail. @@ -382,13 +523,14 @@ impl PluginContext { } /// Registers an LLM sanitize-request guardrail. - pub fn register_llm_sanitize_request_guardrail( + pub fn register_llm_sanitize_request_guardrail( &mut self, name: &str, priority: i32, callback: F, ) where - F: Fn(LlmRequest) -> LlmRequest + Send + Sync + 'static, + F: Fn(LlmRequest, LlmSanitizeRequestContext) -> Fut + Send + Sync + 'static, + Fut: Future>> + Send + 'static, { self.push_registration( name, @@ -396,19 +538,21 @@ impl PluginContext { priority, false, ); - self.handlers - .llm_sanitize_requests - .insert(name.into(), Arc::new(callback)); + self.handlers.llm_sanitize_requests.insert( + name.into(), + Arc::new(move |request, context| Box::pin(callback(request, context))), + ); } /// Registers an LLM sanitize-response guardrail. - pub fn register_llm_sanitize_response_guardrail( + pub fn register_llm_sanitize_response_guardrail( &mut self, name: &str, priority: i32, callback: F, ) where - F: Fn(Json) -> Json + Send + Sync + 'static, + F: Fn(Json, LlmSanitizeResponseContext) -> Fut + Send + Sync + 'static, + Fut: Future>> + Send + 'static, { self.push_registration( name, @@ -416,9 +560,10 @@ impl PluginContext { priority, false, ); - self.handlers - .llm_sanitize_responses - .insert(name.into(), Arc::new(callback)); + self.handlers.llm_sanitize_responses.insert( + name.into(), + Arc::new(move |response, context| Box::pin(callback(response, context))), + ); } /// Registers an LLM conditional-execution guardrail. @@ -541,6 +686,70 @@ pub struct PluginRuntime { } impl PluginRuntime { + async fn decode_llm_codec_request( + &self, + capability_id: &str, + invocation_id: &str, + request: &LlmRequest, + ) -> Result { + let mut client = self.host_client().await?; + let response = client + .decode_llm_codec_request(Request::new(LlmCodecDecodeRequest { + activation_id: self.activation_id.clone(), + auth_token: self.auth_token.clone(), + codec_capability_id: capability_id.into(), + invocation_id: invocation_id.into(), + request: Some(json_envelope(LLM_REQUEST_SCHEMA, request)?), + })) + .await + .map_err(|err| WorkerSdkError::Transport(err.to_string()))? + .into_inner(); + decode_typed_json_result(response, ANNOTATED_LLM_REQUEST_SCHEMA) + } + + async fn encode_llm_codec_request( + &self, + capability_id: &str, + invocation_id: &str, + annotated: &AnnotatedLlmRequest, + original: &LlmRequest, + ) -> Result { + let mut client = self.host_client().await?; + let response = client + .encode_llm_codec_request(Request::new(LlmCodecEncodeRequest { + activation_id: self.activation_id.clone(), + auth_token: self.auth_token.clone(), + codec_capability_id: capability_id.into(), + invocation_id: invocation_id.into(), + annotated_request: Some(json_envelope(ANNOTATED_LLM_REQUEST_SCHEMA, annotated)?), + original_request: Some(json_envelope(LLM_REQUEST_SCHEMA, original)?), + })) + .await + .map_err(|err| WorkerSdkError::Transport(err.to_string()))? + .into_inner(); + decode_typed_json_result(response, LLM_REQUEST_SCHEMA) + } + + async fn decode_llm_codec_response( + &self, + capability_id: &str, + invocation_id: &str, + response: &Json, + ) -> Result { + let mut client = self.host_client().await?; + let response = client + .decode_llm_codec_response(Request::new(LlmCodecDecodeResponse { + activation_id: self.activation_id.clone(), + auth_token: self.auth_token.clone(), + codec_capability_id: capability_id.into(), + invocation_id: invocation_id.into(), + response: Some(json_envelope(JSON_SCHEMA, response)?), + })) + .await + .map_err(|err| WorkerSdkError::Transport(err.to_string()))? + .into_inner(); + decode_typed_json_result(response, JSON_SCHEMA) + } /// Emits a mark event through the host runtime. pub async fn emit_mark( &self, @@ -713,7 +922,7 @@ impl ToolNext { activation_id: self.runtime.activation_id.clone(), auth_token: self.runtime.auth_token.clone(), continuation_id: self.continuation_id.clone(), - value: Some(json_envelope("nemo.relay.Json@1", &value)?), + value: Some(json_envelope(JSON_SCHEMA, &value)?), })) .await .map_err(|err| WorkerSdkError::Transport(err.to_string()))? @@ -738,7 +947,7 @@ impl LlmNext { activation_id: self.runtime.activation_id.clone(), auth_token: self.runtime.auth_token.clone(), continuation_id: self.continuation_id.clone(), - request: Some(json_envelope("nemo.relay.LlmRequest@1", &request)?), + request: Some(json_envelope(LLM_REQUEST_SCHEMA, &request)?), })) .await .map_err(|err| WorkerSdkError::Transport(err.to_string()))? @@ -764,7 +973,7 @@ impl LlmStreamNext { activation_id: self.runtime.activation_id.clone(), auth_token: self.runtime.auth_token.clone(), continuation_id: self.continuation_id.clone(), - request: Some(json_envelope("nemo.relay.LlmRequest@1", &request)?), + request: Some(json_envelope(LLM_REQUEST_SCHEMA, &request)?), })) .await .map_err(|err| WorkerSdkError::Transport(err.to_string()))?; @@ -1222,7 +1431,7 @@ impl PluginWorker for WorkerService { let chunk = match item { Ok(value) => StreamChunk { item: Some(nemo_relay_worker_proto::v1::stream_chunk::Item::Value( - match json_envelope("nemo.relay.Json@1", &value) { + match json_envelope(JSON_SCHEMA, &value) { Ok(value) => value, Err(err) => { let _ = @@ -1431,127 +1640,267 @@ impl WorkerService { let surface = RegistrationSurface::try_from(request.surface) .map_err(|_| WorkerSdkError::InvalidInput("unknown registration surface".into()))?; match surface { - RegistrationSurface::Subscriber => { - let event = event_payload(request.payload)?; - let handler = self.subscriber(&request.registration_name)?; - with_thread_scope(&scope, || handler(&event)); - Ok(empty_response()) - } + RegistrationSurface::Subscriber => self.invoke_subscriber_response(request, &scope), RegistrationSurface::MarkSanitizeGuardrail | RegistrationSurface::ScopeSanitizeStartGuardrail | RegistrationSurface::ScopeSanitizeEndGuardrail => { - let event = event_payload(request.payload)?; - let fields = event.sanitize_fields(); - let handler = self.event_sanitizer(surface, &request.registration_name)?; - let fields = with_thread_scope(&scope, || handler(&event, fields)); - Ok(json_response( - serde_json::to_value(fields) - .expect("event sanitize fields are JSON serializable"), + self.invoke_event_sanitize_response(request, &scope, surface) + .await + } + RegistrationSurface::ToolSanitizeRequestGuardrail + | RegistrationSurface::ToolSanitizeResponseGuardrail + | RegistrationSurface::ToolConditionalExecutionGuardrail + | RegistrationSurface::ToolRequestIntercept + | RegistrationSurface::ToolExecutionIntercept => { + self.invoke_tool_response(request, &scope, surface).await + } + RegistrationSurface::LlmSanitizeRequestGuardrail + | RegistrationSurface::LlmSanitizeResponseGuardrail + | RegistrationSurface::LlmConditionalExecutionGuardrail + | RegistrationSurface::LlmRequestIntercept + | RegistrationSurface::LlmExecutionIntercept => { + self.invoke_llm_response(request, &scope, surface).await + } + RegistrationSurface::LlmStreamExecutionIntercept | RegistrationSurface::Unspecified => { + Err(WorkerSdkError::InvalidInput( + "surface must use InvokeStream or is unspecified".into(), )) } + } + } + + fn invoke_subscriber_response( + &self, + request: InvokeRequest, + scope: &Option, + ) -> Result { + let event = event_payload(request.payload)?; + let handler = self.subscriber(&request.registration_name)?; + with_thread_scope(scope, || handler(&event)); + Ok(empty_response()) + } + + async fn invoke_event_sanitize_response( + &self, + request: InvokeRequest, + scope: &Option, + surface: RegistrationSurface, + ) -> Result { + let event = event_payload(request.payload)?; + let fields = event.sanitize_fields(); + let handler = self.event_sanitizer(surface, &request.registration_name)?; + let fields = with_thread_scope(scope, || handler(&event, fields)).await?; + Ok(json_response( + serde_json::to_value(fields).expect("event sanitize fields are JSON serializable"), + )) + } + + async fn invoke_tool_response( + &self, + request: InvokeRequest, + scope: &Option, + surface: RegistrationSurface, + ) -> Result { + match surface { RegistrationSurface::ToolSanitizeRequestGuardrail => { - let payload = tool_payload(request.payload)?; - let handler = self.tool_sanitize_request(&request.registration_name)?; - Ok(json_response(with_thread_scope(&scope, || { - handler(&payload.tool_name, payload.value) - }))) + self.invoke_tool_sanitize_request_response(request, scope) + .await } RegistrationSurface::ToolSanitizeResponseGuardrail => { - let payload = tool_payload(request.payload)?; - let handler = self.tool_sanitize_response(&request.registration_name)?; - Ok(json_response(with_thread_scope(&scope, || { - handler(&payload.tool_name, payload.value) - }))) + self.invoke_tool_sanitize_response_response(request, scope) + .await } RegistrationSurface::ToolConditionalExecutionGuardrail => { - let payload = tool_payload(request.payload)?; - let handler = self.tool_conditional(&request.registration_name)?; - Ok(guardrail_response(with_thread_scope(&scope, || { - handler(&payload.tool_name, &payload.value) - })?)) + self.invoke_tool_conditional_response(request, scope) } RegistrationSurface::ToolRequestIntercept => { - let payload = tool_payload(request.payload)?; - let handler = self.tool_request(&request.registration_name)?; - Ok(json_response(with_thread_scope(&scope, || { - handler(&payload.tool_name, payload.value) - })?)) + self.invoke_tool_request_response(request, scope) } RegistrationSurface::ToolExecutionIntercept => { - let payload = tool_payload(request.payload)?; - let handler = self.tool_execution(&request.registration_name)?; - let next = ToolNext { - runtime: self.runtime.clone(), - continuation_id: request.continuation_id, - }; - let future = - with_thread_scope(&scope, || handler(&payload.tool_name, payload.value, next)); - Ok(tool_execution_response(future.await?)?) + self.invoke_tool_execution_response(request, scope).await } + _ => unreachable!("tool surface was pre-filtered"), + } + } + + async fn invoke_tool_sanitize_request_response( + &self, + request: InvokeRequest, + scope: &Option, + ) -> Result { + let payload = tool_payload(request.payload)?; + let handler = self.tool_sanitize_request(&request.registration_name)?; + Ok(json_response( + with_thread_scope(scope, || handler(&payload.tool_name, payload.value)).await?, + )) + } + + async fn invoke_tool_sanitize_response_response( + &self, + request: InvokeRequest, + scope: &Option, + ) -> Result { + let payload = tool_payload(request.payload)?; + let handler = self.tool_sanitize_response(&request.registration_name)?; + Ok(json_response( + with_thread_scope(scope, || handler(&payload.tool_name, payload.value)).await?, + )) + } + + fn invoke_tool_conditional_response( + &self, + request: InvokeRequest, + scope: &Option, + ) -> Result { + let payload = tool_payload(request.payload)?; + let handler = self.tool_conditional(&request.registration_name)?; + Ok(guardrail_response(with_thread_scope(scope, || { + handler(&payload.tool_name, &payload.value) + })?)) + } + + fn invoke_tool_request_response( + &self, + request: InvokeRequest, + scope: &Option, + ) -> Result { + let payload = tool_payload(request.payload)?; + let handler = self.tool_request(&request.registration_name)?; + Ok(json_response(with_thread_scope(scope, || { + handler(&payload.tool_name, payload.value) + })?)) + } + + async fn invoke_tool_execution_response( + &self, + request: InvokeRequest, + scope: &Option, + ) -> Result { + let payload = tool_payload(request.payload)?; + let handler = self.tool_execution(&request.registration_name)?; + let next = ToolNext { + runtime: self.runtime.clone(), + continuation_id: request.continuation_id, + }; + let future = with_thread_scope(scope, || handler(&payload.tool_name, payload.value, next)); + tool_execution_response(future.await?) + } + + async fn invoke_llm_response( + &self, + request: InvokeRequest, + scope: &Option, + surface: RegistrationSurface, + ) -> Result { + match surface { RegistrationSurface::LlmSanitizeRequestGuardrail => { - let payload = llm_payload(request.payload)?; - let request_value = required_json::(payload.request, "llm request")?; - let handler = self.llm_sanitize_request(&request.registration_name)?; - let request = with_thread_scope(&scope, || handler(request_value)); - Ok(json_response( - serde_json::to_value(request).expect("LLM request is JSON serializable"), - )) + self.invoke_llm_sanitize_request_response(request, scope) + .await } RegistrationSurface::LlmSanitizeResponseGuardrail => { - let payload = llm_payload(request.payload)?; - let response = required_json::(payload.response, "llm response")?; - let handler = self.llm_sanitize_response(&request.registration_name)?; - Ok(json_response(with_thread_scope(&scope, || { - handler(response) - }))) + self.invoke_llm_sanitize_response_response(request, scope) + .await } RegistrationSurface::LlmConditionalExecutionGuardrail => { - let payload = llm_payload(request.payload)?; - let request_value = required_json::(payload.request, "llm request")?; - let handler = self.llm_conditional(&request.registration_name)?; - Ok(guardrail_response(with_thread_scope(&scope, || { - handler(&request_value) - })?)) + self.invoke_llm_conditional_response(request, scope) } RegistrationSurface::LlmRequestIntercept => { - let payload = llm_payload(request.payload)?; - let request_value = required_json::(payload.request, "llm request")?; - let annotated = payload - .annotated_request - .map(|value| { - decode_expected_json_envelope::( - &value, - "annotated llm request", - ANNOTATED_LLM_REQUEST_SCHEMA, - ) - }) - .transpose()?; - let handler = self.llm_request(&request.registration_name)?; - let outcome = with_thread_scope(&scope, || { - handler(&payload.model_name, request_value, annotated) - })?; - Ok(llm_request_response(outcome)?) + self.invoke_llm_request_response(request, scope) } RegistrationSurface::LlmExecutionIntercept => { - let payload = llm_payload(request.payload)?; - let request_value = required_json::(payload.request, "llm request")?; - let handler = self.llm_execution(&request.registration_name)?; - let next = LlmNext { - runtime: self.runtime.clone(), - continuation_id: request.continuation_id, - }; - let future = - with_thread_scope(&scope, || handler(&payload.model_name, request_value, next)); - Ok(json_response(future.await?)) - } - RegistrationSurface::LlmStreamExecutionIntercept | RegistrationSurface::Unspecified => { - Err(WorkerSdkError::InvalidInput( - "surface must use InvokeStream or is unspecified".into(), - )) + self.invoke_llm_execution_response(request, scope).await } + _ => unreachable!("LLM surface was pre-filtered"), } } + async fn invoke_llm_sanitize_request_response( + &self, + request: InvokeRequest, + scope: &Option, + ) -> Result { + let payload = llm_payload(request.payload)?; + let mut context = payload.sanitize_request_context(&request.invocation_id); + context.runtime = Some(self.runtime.clone()); + let request_value = required_json::(payload.request, "llm request")?; + let handler = self.llm_sanitize_request(&request.registration_name)?; + match with_thread_scope(scope, || handler(request_value, context)).await? { + Some(request) => Ok(json_response( + serde_json::to_value(request).expect("LLM request is JSON serializable"), + )), + None => Ok(empty_response()), + } + } + + async fn invoke_llm_sanitize_response_response( + &self, + request: InvokeRequest, + scope: &Option, + ) -> Result { + let payload = llm_payload(request.payload)?; + let mut context = payload.sanitize_response_context(&request.invocation_id); + context.runtime = Some(self.runtime.clone()); + let response = required_json::(payload.response, "llm response")?; + let handler = self.llm_sanitize_response(&request.registration_name)?; + match with_thread_scope(scope, || handler(response, context)).await? { + Some(response) => Ok(json_response(response)), + None => Ok(empty_response()), + } + } + + fn invoke_llm_conditional_response( + &self, + request: InvokeRequest, + scope: &Option, + ) -> Result { + let payload = llm_payload(request.payload)?; + let request_value = required_json::(payload.request, "llm request")?; + let handler = self.llm_conditional(&request.registration_name)?; + Ok(guardrail_response(with_thread_scope(scope, || { + handler(&request_value) + })?)) + } + + fn invoke_llm_request_response( + &self, + request: InvokeRequest, + scope: &Option, + ) -> Result { + let payload = llm_payload(request.payload)?; + let request_value = required_json::(payload.request, "llm request")?; + let annotated = payload + .annotated_request + .map(|value| { + decode_expected_json_envelope::( + &value, + "annotated llm request", + ANNOTATED_LLM_REQUEST_SCHEMA, + ) + }) + .transpose()?; + let handler = self.llm_request(&request.registration_name)?; + let outcome = with_thread_scope(scope, || { + handler(&payload.model_name, request_value, annotated) + })?; + llm_request_response(outcome) + } + + async fn invoke_llm_execution_response( + &self, + request: InvokeRequest, + scope: &Option, + ) -> Result { + let payload = llm_payload(request.payload)?; + let request_value = required_json::(payload.request, "llm request")?; + let handler = self.llm_execution(&request.registration_name)?; + let next = LlmNext { + runtime: self.runtime.clone(), + continuation_id: request.continuation_id, + }; + let future = with_thread_scope(scope, || handler(&payload.model_name, request_value, next)); + Ok(json_response(future.await?)) + } + fn subscriber(&self, name: &str) -> Result { self.handlers .lock() @@ -1719,6 +2068,65 @@ struct LlmPayload { request: Option, annotated_request: Option, response: Option, + sanitize_context: Option, +} + +impl LlmPayload { + fn sanitize_request_context(&self, invocation_id: &str) -> LlmSanitizeRequestContext { + let codec = match self.sanitize_context.as_ref() { + Some(nemo_relay_worker_proto::v1::llm_invocation::SanitizeContext::RequestSanitizeContext(context)) => context.codec.as_ref(), + _ => None, + }; + LlmSanitizeRequestContext { + codec: codec_identity_from_proto(codec), + runtime: None, + codec_capability_id: match self.sanitize_context.as_ref() { + Some(nemo_relay_worker_proto::v1::llm_invocation::SanitizeContext::RequestSanitizeContext(context)) => context.codec_capability_id.clone(), + _ => None, + }, + invocation_id: Some(invocation_id.to_owned()), + } + } + + fn sanitize_response_context(&self, invocation_id: &str) -> LlmSanitizeResponseContext { + let codec = match self.sanitize_context.as_ref() { + Some(nemo_relay_worker_proto::v1::llm_invocation::SanitizeContext::ResponseSanitizeContext(context)) => context.codec.as_ref(), + _ => None, + }; + LlmSanitizeResponseContext { + codec: codec_identity_from_proto(codec), + runtime: None, + codec_capability_id: match self.sanitize_context.as_ref() { + Some(nemo_relay_worker_proto::v1::llm_invocation::SanitizeContext::ResponseSanitizeContext(context)) => context.codec_capability_id.clone(), + _ => None, + }, + invocation_id: Some(invocation_id.to_owned()), + } + } +} + +fn codec_identity_from_proto( + codec: Option<&nemo_relay_worker_proto::v1::LlmCodecIdentity>, +) -> LlmCodecIdentity { + let codec_kind = codec + .map(|codec| codec.kind) + .unwrap_or(LlmCodecKind::Unspecified as i32); + let codec_id = codec.and_then(|codec| codec.id.clone()); + match LlmCodecKind::try_from(codec_kind).ok() { + Some(LlmCodecKind::Unspecified) => LlmCodecIdentity::None, + Some(LlmCodecKind::Builtin) => match codec_id.as_deref() { + Some("openai_chat") => LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiChat), + Some("openai_responses") => LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiResponses), + Some("anthropic_messages") => { + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::AnthropicMessages) + } + _ => LlmCodecIdentity::Opaque, + }, + Some(LlmCodecKind::Runtime) => codec_id + .filter(|id| !id.is_empty()) + .map_or(LlmCodecIdentity::Opaque, LlmCodecIdentity::Runtime), + Some(LlmCodecKind::Opaque) | None => LlmCodecIdentity::Opaque, + } } fn event_payload( @@ -1758,6 +2166,7 @@ fn llm_payload( request: value.request, annotated_request: value.annotated_request, response: value.response, + sanitize_context: value.sanitize_context, }), _ => Err(WorkerSdkError::InvalidInput("expected llm payload".into())), } @@ -1797,7 +2206,7 @@ fn json_response(value: Json) -> InvokeResponse { InvokeResponse { result: Some(nemo_relay_worker_proto::v1::invoke_response::Result::Json( JsonResult { - value: Some(infallible_json_envelope("nemo.relay.Json@1", &value)), + value: Some(infallible_json_envelope(JSON_SCHEMA, &value)), error: None, }, )), @@ -1863,10 +2272,23 @@ fn json_result_to_sdk(result: JsonResult) -> Result { required_json(result.value, "json result") } +fn decode_typed_json_result( + result: JsonResult, + expected_schema: &str, +) -> Result { + if let Some(error) = result.error { + return Err(worker_error_to_sdk(error)); + } + let value = result + .value + .ok_or_else(|| WorkerSdkError::InvalidInput("json result is missing".into()))?; + decode_expected_json_envelope(&value, "json result", expected_schema) +} + fn optional_json_envelope(value: Option) -> Result> { value .as_ref() - .map(|value| json_envelope("nemo.relay.Json@1", value).map_err(WorkerSdkError::from)) + .map(|value| json_envelope(JSON_SCHEMA, value).map_err(WorkerSdkError::from)) .transpose() } diff --git a/crates/worker/tests/worker_sdk_tests.rs b/crates/worker/tests/worker_sdk_tests.rs index bedf2fb71..80653155f 100644 --- a/crates/worker/tests/worker_sdk_tests.rs +++ b/crates/worker/tests/worker_sdk_tests.rs @@ -200,7 +200,22 @@ async fn worker_service_enforces_auth_and_reports_registrations() { assert_eq!(invalid_register_config.code(), tonic::Code::InvalidArgument); let registrations = register_plugin(&mut client).await; - assert_eq!(registrations.len(), 19); + assert_eq!(registrations.len(), 21); + for local_name in [ + "llm-sanitize-request", + "llm-sanitize-response", + "llm-sanitize-omit-request", + "llm-sanitize-omit-response", + ] { + assert_eq!( + registrations + .iter() + .filter(|registration| registration.local_name == local_name) + .count(), + 1, + "expected exactly one {local_name} registration" + ); + } assert_eq!( registrations .iter() @@ -484,7 +499,20 @@ async fn worker_service_invokes_every_registration_surface() { let plugin = Arc::new(SurfacePlugin::default()); let events = plugin.events.clone(); let (worker_handle, mut client) = spawn_worker(plugin, tcp_endpoint(&host_endpoint)).await; - register_plugin(&mut client).await; + let registrations = register_plugin(&mut client).await; + for name in [ + "llm-sanitize-request", + "llm-sanitize-response", + "llm-sanitize-omit-request", + "llm-sanitize-omit-response", + ] { + assert!( + registrations + .iter() + .any(|registration| registration.local_name == name), + "{name} must be registered" + ); + } let subscriber_response = client .invoke(Request::new(event_invoke("subscriber"))) @@ -632,6 +660,47 @@ async fn worker_service_invokes_every_registration_surface() { "phase", "llm_sanitize_response", ); + let omitted_request = client + .invoke(Request::new(llm_invoke( + "llm-sanitize-omit-request", + RegistrationSurface::LlmSanitizeRequestGuardrail, + llm_request(), + None, + None, + ))) + .await + .expect("request omission invoke") + .into_inner(); + assert_empty_response(omitted_request); + let omitted_response = client + .invoke(Request::new(llm_invoke( + "llm-sanitize-omit-response", + RegistrationSurface::LlmSanitizeResponseGuardrail, + llm_request(), + None, + Some(json!({})), + ))) + .await + .expect("response omission invoke") + .into_inner(); + assert_empty_response(omitted_response); + let codec_request_result = invoke_json(&mut client, llm_request_codec_invoke()).await; + assert_eq!( + codec_request_result, + json!({ + "headers": {}, + "content": {"phase": "llm_sanitize_request"}, + }) + ); + + let codec_response_result = invoke_json(&mut client, llm_response_codec_invoke()).await; + assert_eq!( + codec_response_result, + json!({ + "value": "secret", + "phase": "llm_sanitize_response", + }) + ); assert_eq!( invoke_guardrail( &mut client, @@ -718,6 +787,9 @@ async fn worker_service_invokes_every_registration_surface() { assert!(calls.contains(&"tool_next:next-1".into())); assert!(calls.contains(&"llm_next:next-1".into())); assert!(calls.contains(&"llm_stream_next:next-1".into())); + assert!(calls.contains(&"codec_request_decode:request-capability".into())); + assert!(calls.contains(&"codec_request_encode:request-capability".into())); + assert!(calls.contains(&"codec_response_decode:response-capability".into())); assert!(calls.contains(&"mark:stream-poll:stack-1:parent-1".into())); assert!(calls.contains(&"push:scope-agent:explicit-stack:".into())); assert!(calls.contains(&"push:scope-unknown:explicit-stack:".into())); @@ -1439,6 +1511,43 @@ async fn worker_service_propagates_host_runtime_errors() { "llm next failed", ); + for (failures, request, expected) in [ + ( + MockHostFailures { + codec_request_decode: true, + ..Default::default() + }, + llm_request_codec_invoke(), + "codec request decode failed", + ), + ( + MockHostFailures { + codec_request_encode: true, + ..Default::default() + }, + llm_request_codec_invoke(), + "codec request encode failed", + ), + ( + MockHostFailures { + codec_response_decode: true, + ..Default::default() + }, + llm_response_codec_invoke(), + "codec response decode failed", + ), + ] { + host.set_failures(failures); + assert_worker_error( + client + .invoke(Request::new(request)) + .await + .expect("codec failure returns structured error") + .into_inner(), + expected, + ); + } + for (mode, expected) in [ (MockStreamMode::WorkerError, "stream worker failed"), (MockStreamMode::EmptyChunk, "empty stream chunk"), @@ -1489,8 +1598,8 @@ impl WorkerPlugin for DuplicateEventSanitizerPlugin { } fn register(&self, ctx: &mut PluginContext, _config: &Json) -> Result<()> { - ctx.register_mark_sanitize_guardrail("duplicate", 0, |_, fields| fields); - ctx.register_mark_sanitize_guardrail("duplicate", 1, |_, fields| fields); + ctx.register_mark_sanitize_guardrail("duplicate", 0, |_, fields| async move { Ok(fields) }); + ctx.register_mark_sanitize_guardrail("duplicate", 1, |_, fields| async move { Ok(fields) }); Ok(()) } } @@ -1634,22 +1743,33 @@ impl WorkerPlugin for SurfacePlugin { .push(event.name().into()); }); ctx.register_mark_sanitize_guardrail("event-sanitize", 1, |event, mut fields| { - fields.data = Some(json!({"name": event.name(), "phase": "mark"})); - fields - }); - ctx.register_scope_sanitize_start_guardrail("event-sanitize", 1, |_, mut fields| { - fields.metadata = Some(json!({"phase": "scope_start"})); - fields - }); - ctx.register_scope_sanitize_end_guardrail("event-sanitize", 1, |_, mut fields| { - fields.metadata = Some(json!({"phase": "scope_end"})); - fields + let event_name = event.name().to_owned(); + async move { + fields.data = Some(json!({"name": event_name, "phase": "mark"})); + Ok(fields) + } }); - ctx.register_tool_sanitize_request_guardrail("tool-sanitize", 1, |_, value| { - set_json_field(value, "phase", "tool_sanitize_request") + ctx.register_scope_sanitize_start_guardrail( + "event-sanitize", + 1, + |_, mut fields| async move { + fields.metadata = Some(json!({"phase": "scope_start"})); + Ok(fields) + }, + ); + ctx.register_scope_sanitize_end_guardrail( + "event-sanitize", + 1, + |_, mut fields| async move { + fields.metadata = Some(json!({"phase": "scope_end"})); + Ok(fields) + }, + ); + ctx.register_tool_sanitize_request_guardrail("tool-sanitize", 1, |_, value| async move { + Ok(set_json_field(value, "phase", "tool_sanitize_request")) }); - ctx.register_tool_sanitize_response_guardrail("tool-sanitize", 1, |_, value| { - set_json_field(value, "phase", "tool_sanitize_response") + ctx.register_tool_sanitize_response_guardrail("tool-sanitize", 1, |_, value| async move { + Ok(set_json_field(value, "phase", "tool_sanitize_response")) }); ctx.register_tool_conditional_execution_guardrail("tool-conditional", 1, |_, value| { Ok(value @@ -1730,12 +1850,41 @@ impl WorkerPlugin for SurfacePlugin { } }); - ctx.register_llm_sanitize_request_guardrail("llm-sanitize-request", 1, |request| { - set_llm_phase(request, "llm_sanitize_request") - }); - ctx.register_llm_sanitize_response_guardrail("llm-sanitize-response", 1, |value| { - set_json_field(value, "phase", "llm_sanitize_response") - }); + ctx.register_llm_sanitize_request_guardrail( + "llm-sanitize-request", + 1, + |mut request, context| async move { + if let Some(codec) = context.resolve_codec() { + let annotated = codec.decode(&request).await?; + request = codec.encode(&annotated, &request).await?; + } + Ok(Some(set_llm_phase(request, "llm_sanitize_request"))) + }, + ); + ctx.register_llm_sanitize_response_guardrail( + "llm-sanitize-response", + 1, + |value, context| async move { + if let Some(codec) = context.resolve_codec() { + codec.decode(&value).await?; + } + Ok(Some(set_json_field( + value, + "phase", + "llm_sanitize_response", + ))) + }, + ); + ctx.register_llm_sanitize_request_guardrail( + "llm-sanitize-omit-request", + 1, + |_request, _context| async { Ok(None) }, + ); + ctx.register_llm_sanitize_response_guardrail( + "llm-sanitize-omit-response", + 1, + |_response, _context| async { Ok(None) }, + ); ctx.register_llm_conditional_execution_guardrail("llm-conditional", 1, |request| { Ok(request .content @@ -1854,6 +2003,9 @@ struct MockHostFailures { tool_next: bool, llm_next: bool, llm_stream_mode: MockStreamMode, + codec_request_decode: bool, + codec_request_encode: bool, + codec_response_decode: bool, } #[derive(Clone, Default)] @@ -2062,6 +2214,81 @@ impl RelayHostRuntime for MockHost { let stream = tokio_stream::iter(chunks); Ok(Response::new(Box::pin(stream))) } + + async fn decode_llm_codec_request( + &self, + request: Request, + ) -> std::result::Result, Status> { + let request = request.into_inner(); + authorize_host(&request.activation_id, &request.auth_token)?; + self.record(format!( + "codec_request_decode:{}", + request.codec_capability_id + )); + if self.failures().codec_request_decode { + return Ok(Response::new(JsonResult { + value: None, + error: Some(worker_error("codec request decode failed")), + })); + } + Ok(Response::new(JsonResult { + value: Some( + json_envelope("nemo.relay.AnnotatedLlmRequest@2", &json!({})) + .expect("encode annotated request"), + ), + error: None, + })) + } + + async fn encode_llm_codec_request( + &self, + request: Request, + ) -> std::result::Result, Status> { + let request = request.into_inner(); + authorize_host(&request.activation_id, &request.auth_token)?; + self.record(format!( + "codec_request_encode:{}", + request.codec_capability_id + )); + if self.failures().codec_request_encode { + return Ok(Response::new(JsonResult { + value: None, + error: Some(worker_error("codec request encode failed")), + })); + } + Ok(Response::new(JsonResult { + value: Some( + json_envelope( + "nemo.relay.LlmRequest@1", + &json!({"headers": {}, "content": {}}), + ) + .expect("encode LLM request"), + ), + error: None, + })) + } + + async fn decode_llm_codec_response( + &self, + request: Request, + ) -> std::result::Result, Status> { + let request = request.into_inner(); + authorize_host(&request.activation_id, &request.auth_token)?; + self.record(format!( + "codec_response_decode:{}", + request.codec_capability_id + )); + if self.failures().codec_response_decode { + return Ok(Response::new(JsonResult { + value: None, + error: Some(worker_error("codec response decode failed")), + })); + } + Ok(Response::new(JsonResult { + value: Some(json_env(json!({}))), + error: None, + })) + } } async fn spawn_worker( @@ -2317,6 +2544,7 @@ fn llm_invoke( ), annotated_request: annotated_request.map(json_env), response: response.map(json_env), + sanitize_context: None, }, )), } @@ -2340,6 +2568,7 @@ fn llm_invoke_without_request( request: None, annotated_request: None, response: None, + sanitize_context: None, }, )), } @@ -2404,6 +2633,58 @@ fn llm_request_with_block() -> LlmRequest { } } +fn llm_request_codec_invoke() -> InvokeRequest { + let mut request = llm_invoke( + "llm-sanitize-request", + RegistrationSurface::LlmSanitizeRequestGuardrail, + llm_request(), + None, + None, + ); + if let Some(nemo_relay_worker_proto::v1::invoke_request::Payload::Llm(invocation)) = + request.payload.as_mut() + { + invocation.sanitize_context = Some( + nemo_relay_worker_proto::v1::llm_invocation::SanitizeContext::RequestSanitizeContext( + nemo_relay_worker_proto::v1::LlmSanitizeRequestContext { + codec: Some(nemo_relay_worker_proto::v1::LlmCodecIdentity { + kind: nemo_relay_worker_proto::v1::LlmCodecKind::Builtin as i32, + id: Some("openai_chat".into()), + }), + codec_capability_id: Some("request-capability".into()), + }, + ), + ); + } + request +} + +fn llm_response_codec_invoke() -> InvokeRequest { + let mut request = llm_invoke( + "llm-sanitize-response", + RegistrationSurface::LlmSanitizeResponseGuardrail, + llm_request(), + None, + Some(json!({"value": "secret"})), + ); + if let Some(nemo_relay_worker_proto::v1::invoke_request::Payload::Llm(invocation)) = + request.payload.as_mut() + { + invocation.sanitize_context = Some( + nemo_relay_worker_proto::v1::llm_invocation::SanitizeContext::ResponseSanitizeContext( + nemo_relay_worker_proto::v1::LlmSanitizeResponseContext { + codec: Some(nemo_relay_worker_proto::v1::LlmCodecIdentity { + kind: nemo_relay_worker_proto::v1::LlmCodecKind::Opaque as i32, + id: None, + }), + codec_capability_id: Some("response-capability".into()), + }, + ), + ); + } + request +} + fn set_json_field(mut value: Json, key: &str, field_value: &str) -> Json { value .as_object_mut() diff --git a/docs/about-nemo-relay/concepts/middleware.mdx b/docs/about-nemo-relay/concepts/middleware.mdx index 8b4011b63..8f2b9aca1 100644 --- a/docs/about-nemo-relay/concepts/middleware.mdx +++ b/docs/about-nemo-relay/concepts/middleware.mdx @@ -191,6 +191,124 @@ 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 +## Codec-Aware LLM Sanitizers + +Every LLM sanitize guardrail receives the payload first and a required +directional, per-call context second: + +```text +request: (LlmRequest, LlmSanitizeRequestContext) -> Option +response: (Json, LlmSanitizeResponseContext) -> Option +``` + +`context.codec` is a binding-native codec identity, not a JSON payload. +Python and Node.js expose `kind` and, when applicable, `id` properties. +In-process Rust and the typed native Rust SDK expose enum variants. The raw +native ABI exposes the same information through `codec_kind` and `codec_id`. + +`codec.kind` is `none` for a call with no codec, `builtin` for +Relay's built-in `openai_chat`, `openai_responses`, and `anthropic_messages` +codecs, `runtime` for a named runtime-registered codec, and `opaque` for an +active codec without a registered identity. `codec.id` is present only for +`builtin` and `runtime`. Do not infer a provider from an opaque request shape. + +Return a payload to continue the sanitizer chain. Return `None` (or `null` in +JavaScript) only when observing the payload would be unsafe: Relay omits the +LLM event payload and its annotation, while leaving the client-visible request +and response unchanged. Omission short-circuits later LLM sanitizers. + +All LLM sanitizer callbacks must implement this two-parameter contract. For an +in-process sanitizer, use `resolve_codec()` to access the active codec +implementation. A resolver returns no codec for manual calls without one. +Worker-plugin contexts resolve to an invocation-scoped asynchronous proxy with +the same directional operations: request codecs provide `decode(request)` and +`encode(annotated, original)`, while response codecs provide `decode(response)`. +Python in-process response codecs expose that operation as `decode_response`, +and Node.js exposes it as `decodeResponse`. An active runtime or opaque codec +resolves just like a built-in codec; its identity does not limit the available +operations. Node.js codecs supplied through decode and encode callbacks have +an `opaque` identity but remain resolvable. The worker proxy is valid only +while its sanitizer callback is running; do not retain it after the callback +returns. + + + + +```python +import nemo_relay +from nemo_relay import LLMRequest +from nemo_relay import guardrails + +def redact_request( + request: LLMRequest, + context: nemo_relay.LlmSanitizeRequestContext, +) -> LLMRequest | None: + codec = context.resolve_codec() + if ( + context.codec.kind == "builtin" + and context.codec.id == "openai_chat" + and codec is not None + ): + annotated = codec.decode(request) + # Apply a policy to the normalized request, then preserve the wire shape. + annotated.messages = [] + return codec.encode(annotated, request) + return request + +guardrails.register_llm_sanitize_request("redact-openai-chat", 10, redact_request) +``` + + + + +```js +const relay = require("nemo-relay-node"); + +relay.registerLlmSanitizeRequestGuardrail( + "redact-openai-chat", + 10, + (request, context) => { + const codec = context.resolveCodec(); + if (codec) { + const annotated = codec.decode(request); + annotated.messages = []; + return codec.encode(annotated, request); + } + return request; + }, +); +``` + + + + +```rust +use std::sync::Arc; +use nemo_relay::api::runtime::{BuiltinLlmCodec, LlmCodecIdentity}; +use nemo_relay::api::registry::register_llm_sanitize_request_guardrail; + +register_llm_sanitize_request_guardrail( + "redact-openai-chat", + 10, + Arc::new(|mut request, context| { + if context.codec() == &LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiChat) + && let Some(codec) = context.resolve_codec() + && let Ok(mut annotated) = codec.decode(&request) + { + annotated.messages.clear(); + request = codec.encode(&annotated, &request).ok()?; + } + Some(request) + }), +)?; +``` + + + + +The same registration names support scope-local and plugin-context +registrations. Priority and name tie-break ordering are unchanged. + ## Detailed Execution Flow diff --git a/docs/about-nemo-relay/release-notes/index.mdx b/docs/about-nemo-relay/release-notes/index.mdx index 8d9550fa2..a87eb2176 100644 --- a/docs/about-nemo-relay/release-notes/index.mdx +++ b/docs/about-nemo-relay/release-notes/index.mdx @@ -27,25 +27,50 @@ Relay](/about-nemo-relay/overview). ## Release 0.7 -NVIDIA NeMo Relay 0.7 release notes are in preparation. +NVIDIA NeMo Relay 0.7 release notes are in preparation. The following +compatibility information applies to the current 0.7 prerelease. ### Highlights -- _Highlights will be added for the 0.7 release._ +- LLM observability sanitizers now receive the active request or response codec + for each managed call. Sanitizers can normalize built-in, + runtime-registered, and opaque codec payloads without changing the + client-visible request or response. +- The PII redaction plugin now selects the active codec per call. One component + can safely sanitize mixed OpenAI Chat, OpenAI Responses, and Anthropic + Messages traffic without a fixed provider codec. +- Native and worker plugins can resolve invocation-scoped codec capabilities. + Request sanitizers can decode and encode requests, and response sanitizers + can decode responses. ### Support Matrix and Compatibility Updates + +LLM sanitizer callbacks require migration in 0.7. Callbacks now receive the +payload first and a directional context second, and they return an optional +payload. There is no one-argument compatibility adapter. Rust worker mark, +scope, tool, and LLM sanitizer callbacks are also asynchronous in 0.7. Native +plugins and workers must be rebuilt with the matching 0.7 SDK and protocol +definitions. + + The [Support Matrix](/reference/support-matrix) is the canonical reference for supported platforms and architectures, worker runtimes, coding agents, and integrations. It also records current limitations, including platform-specific worker requirements. -Migration guidance for upgrading from 0.6 to 0.7 will be added to the -[Migration Guides](/reference/migration-guides) before the release. +Before upgrading from 0.6, follow the +[0.7 Migration Guide](/reference/migration-guides#upgrade-to-nemo-relay-07). +For the new callback contract and codec operations, refer to +[Codec-Aware LLM Sanitizers](/about-nemo-relay/concepts/middleware#codec-aware-llm-sanitizers). ### Fixed Known Issues in 0.7 -- _Fixed issues will be added for the 0.7 release._ +- LLM payload redaction now follows the codec active for each call instead of a + codec captured from plugin configuration. Codec-dependent policies omit the + observability payload and annotation when Relay cannot safely normalize the + payload. This change resolves the mixed-provider PII sanitization issue + [#526](https://github.com/NVIDIA/NeMo-Relay/issues/526). ## Known Issues in 0.7 diff --git a/docs/build-plugins/dynamic-plugins/about.mdx b/docs/build-plugins/dynamic-plugins/about.mdx index 68f2331ba..772a0b827 100644 --- a/docs/build-plugins/dynamic-plugins/about.mdx +++ b/docs/build-plugins/dynamic-plugins/about.mdx @@ -13,7 +13,7 @@ two execution lanes: | Lane | Use when | Stable boundary | | --- | --- | --- | -| `rust_dynamic` | Behavior must run in the Relay process. | Native ABI v1 | +| `rust_dynamic` | Behavior must run in the Relay process. | Native ABI v2 | | `worker` | Behavior should run in a separate local process. | `grpc-v1` | The manifest describes compatibility, capabilities, artifact integrity, and the diff --git a/docs/build-plugins/dynamic-plugins/grpc-worker/grpc-worker-protocol.mdx b/docs/build-plugins/dynamic-plugins/grpc-worker/grpc-worker-protocol.mdx index 872a6bde9..b8ca9bfbb 100644 --- a/docs/build-plugins/dynamic-plugins/grpc-worker/grpc-worker-protocol.mdx +++ b/docs/build-plugins/dynamic-plugins/grpc-worker/grpc-worker-protocol.mdx @@ -74,6 +74,61 @@ invocation, or LLM invocation. `InvokeResponse` returns an empty result, JSON result, guardrail result, LLM request-intercept result, tool-execution result, or `WorkerError`. `InvokeStream` emits JSON chunks or `WorkerError` chunks. +Every LLM sanitizer invocation includes a directional context with tagged codec +identity: `none`, `builtin(id)`, `runtime(id)`, or `opaque`. Worker SDKs expose +`context.resolve_codec()` for active codecs. The resulting invocation-scoped +proxy supports request decode/encode or response decode and calls Relay through +the host runtime; the capability identifier is protocol-internal and expires +when the callback completes. `resolve_codec()` returns no proxy only when no +codec is active. Runtime and opaque codecs remain resolvable. + +Request and response sanitizer handlers always receive `(payload, context)`. +Return the sanitized payload to continue the chain, or return no payload to +omit the observability payload and annotation without changing the +client-visible value. Rust worker sanitizer callbacks are async for every +sanitizer surface: mark, scope, tool, and LLM. Python worker sanitizers can +return either an immediate value or an awaitable. + + + + +```rust +ctx.register_llm_sanitize_request_guardrail( + "normalize-request", + 10, + |request, context| async move { + let Some(codec) = context.resolve_codec() else { + return Ok(Some(request)); + }; + let mut annotated = codec.decode(&request).await?; + annotated.messages.clear(); + Ok(Some(codec.encode(&annotated, &request).await?)) + }, +); +``` + + + + +```python +async def normalize_request(request, context): + codec = context.resolve_codec() + if codec is None: + return request + annotated = await codec.decode(request) + annotated["messages"] = [] + return await codec.encode(annotated, request) + +ctx.register_llm_sanitize_request_guardrail( + "normalize-request", + normalize_request, + priority=10, +) +``` + + + + ## RelayHostRuntime RPCs Relay implements these RPCs for worker callbacks: @@ -83,10 +138,15 @@ Relay implements these RPCs for worker callbacks: | Marks and scopes | `EmitMark`, `PushScope`, `PopScope` | | Isolated scope stacks | `CreateScopeStack`, `DropScopeStack` | | Execution continuations | `ToolNext`, `LlmNext`, `LlmStreamNext` | +| Codec capabilities | `DecodeLlmCodecRequest`, `EncodeLlmCodecRequest`, `DecodeLlmCodecResponse` | Every host-runtime request also carries the activation ID and authentication token. Scope operations include a `ScopeContext`; continuation calls include -the continuation ID that Relay supplied for the active intercept. +the continuation ID that Relay supplied for the active intercept. Codec +operations additionally require the unforgeable capability ID supplied for the +current sanitizer invocation. Relay rejects missing, forged, expired, +wrong-direction, or activation-mismatched capabilities. Codec transformation +failures are non-retryable worker errors. ## Authentication and Endpoints diff --git a/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx b/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx index db11f2acc..a4c162368 100644 --- a/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx +++ b/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx @@ -1,6 +1,6 @@ --- title: "Native Dynamic Plugins (Rust)" -description: "Build in-process Rust shared-library plugins against the NeMo Relay Native ABI v1." +description: "Build in-process Rust shared-library plugins against the NeMo Relay Native ABI v2." position: 10 --- {/* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. @@ -116,7 +116,7 @@ path, then replace `` with that library's SHA-256 digest. Use Native Plugin](/build-plugins/dynamic-plugins/native-dynamic/rust-native-plugin-example) for a complete example with validation, middleware, scopes, and configuration schema support. -## Native ABI v1 +## Native ABI v2 The host passes a `NemoRelayNativeHostApiV1` table to the entry symbol. The plugin returns a `NemoRelayNativePluginV1` descriptor: @@ -143,6 +143,35 @@ ABI callbacks can register these runtime surfaces: Relay keeps the library alive while those registrations exist and deregisters them before unloading it. +LLM sanitize callbacks receive their request or response JSON first, followed +by `NemoRelayNativeLlmSanitizeRequestContext` or +`NemoRelayNativeLlmSanitizeResponseContext`. Each context contains structured +codec identity and a borrowed, callback-lifetime codec handle. `codec_kind` is +`None`, `BuiltIn`, `Runtime`, or `Opaque`. `codec_id` is present for `BuiltIn` +(one of `openai_chat`, `openai_responses`, or `anthropic_messages`) and +`Runtime`, and null for `None` and `Opaque`. + +The request handle supports host operations to decode an `LlmRequest` into an +`AnnotatedLlmRequest` and encode normalized changes onto the original request. +The response handle supports decoding response JSON into an +`AnnotatedLlmResponse`. A null handle means no codec is active. Runtime and +opaque codecs still have non-null handles and support the same operations as a +built-in codec. Do not retain a handle or resolved SDK facade after the +sanitizer callback returns. + +The Rust SDK converts the raw structures into +`LlmSanitizeRequestContext` and `LlmSanitizeResponseContext`. Call +`resolve_codec()` to obtain the safe directional facade. Raw ABI consumers use +the `llm_request_codec_decode`, `llm_request_codec_encode`, and +`llm_response_codec_decode` host-table functions. The host owns all returned +JSON string handles, which callers release with the ordinary host string +release operation. + +A successful null sanitizer output omits the LLM observability payload and its +annotation. Returning an error also fails closed rather than exposing the +unsanitized payload. Neither case changes the client-visible request or +response. + Use the `nemo-relay-plugin` crate rather than the host `nemo-relay` runtime crate. Refer to [Build a Rust Native Plugin](/build-plugins/dynamic-plugins/native-dynamic/rust-native-plugin-example) for the SDK-backed example. diff --git a/docs/build-plugins/language-binding/about.mdx b/docs/build-plugins/language-binding/about.mdx index 73359fd44..ec63654b3 100644 --- a/docs/build-plugins/language-binding/about.mdx +++ b/docs/build-plugins/language-binding/about.mdx @@ -73,6 +73,14 @@ A plugin can install one or more of these runtime surfaces: conditional-execution guardrails, request intercepts, execution intercepts, and stream execution intercepts. +LLM sanitize guardrails always receive the request or response first, followed +by a structured context whose `codec` is `none`, `builtin(id)`, `runtime(id)`, +or `opaque`. Every callback declares `(payload, context)` and may omit the LLM +event payload by returning no value. In-process sanitizer contexts can resolve +the active codec for normalized processing. Worker sanitizer contexts resolve +an invocation-scoped asynchronous codec proxy, so workers can perform the same +directional normalization without receiving a host-process object. + Start with one surface. Add a bundle only when one configuration document clearly controls related behavior, such as a subscriber plus the request intercepts needed to add correlation metadata. ## Registration Lifecycle diff --git a/docs/configure-plugins/pii-redaction/about.mdx b/docs/configure-plugins/pii-redaction/about.mdx index 982308d77..a37cd4812 100644 --- a/docs/configure-plugins/pii-redaction/about.mdx +++ b/docs/configure-plugins/pii-redaction/about.mdx @@ -91,9 +91,11 @@ The current backend boundary is intentional: - Managed tool surfaces are sanitized as JSON payloads with exact JSON-pointer targeting. -- Managed LLM surfaces use the selected built-in codec so redaction can target - normalized Relay request and response shapes such as `/messages/0/content` - and `/message`. +- Managed LLM requests use the active resolved codec for each call, including + built-in, runtime, and opaque codecs. Normalized response projection requires + a recognized built-in codec because response codec capabilities are + decode-only. This lets redaction target normalized Relay shapes such as + `/messages/0/content` and `/message`. - `mark` sanitizes `data`, `category_profile`, and `metadata` independently on every mark event. It defaults to `true`; set `mark = false` to opt out. - `input`, `output`, `tool_input`, and `tool_output` sanitize scope metadata on diff --git a/docs/configure-plugins/pii-redaction/configuration.mdx b/docs/configure-plugins/pii-redaction/configuration.mdx index 329f78892..d10f84813 100644 --- a/docs/configure-plugins/pii-redaction/configuration.mdx +++ b/docs/configure-plugins/pii-redaction/configuration.mdx @@ -72,7 +72,7 @@ The top-level PII redaction object contains: | `tool_input` | Enables sanitization of emitted tool-request observability payloads. Defaults to `true`. | | `tool_output` | Enables sanitization of emitted tool-response observability payloads. Defaults to `true`. | | `priority` | Guardrail priority. Lower values run earlier. Defaults to `100`. | -| `codec` | Managed LLM provider codec. Required when `input` or `output` is enabled. | +| `codec` | Optional compatibility fallback for legacy/manual managed LLM calls that do not expose an active codec. | | `profiles` | Ordered redaction profiles. Each profile selects `builtin` or `local_model`; enabled profiles cover every supported sanitization surface. | | `builtin` | Built-in backend settings used when `mode = "builtin"`. | | `local` | Local-backend settings used when `mode = "local_model"`. | @@ -90,6 +90,37 @@ profiles do not require IDs. Disabled profiles are validated but do not register callbacks. The complete array is replaced as one value during config layering. +### Mixed-Provider Gateway Example + +For a gateway that routes requests to more than one supported provider, omit +`codec`. Relay selects the active codec for each managed call, so the same +policy safely applies to OpenAI Chat, OpenAI Responses, and Anthropic Messages +traffic: + +```toml +[[components]] +kind = "pii_redaction" +enabled = true + +[components.config] + +[[components.config.profiles]] +mode = "builtin" +priority = 80 + +[components.config.profiles.builtin] +action = "redact" +detector = "email" +target_paths = ["/messages/0/content", "/message"] +``` + +Use `codec = "openai_chat"` only as a compatibility fallback for a legacy +manual call where Relay cannot identify an active codec. It is ignored when a +managed call supplies any active codec. Request policies use that codec's +resolved decode/encode capability, including runtime and opaque codecs. +Normalized response projection currently requires a recognized built-in codec +because response codec capabilities are decode-only. + ## Backend Support The following table compares the available PII redaction backends: @@ -120,7 +151,7 @@ paths, and supported codecs. To use `mode = "builtin"`: - `builtin` settings are required. -- `codec` is required when `input` or `output` is enabled. +- Set `codec` only when legacy/manual managed LLM calls need a fallback codec. - `builtin.action` must be `remove`, `redact`, `regex_replace`, `hash`, or `mask`. - `builtin.pattern` or `builtin.detector` is required when `builtin.action = "regex_replace"` or `builtin.action = "redact"`. @@ -140,7 +171,6 @@ enabled = true [components.config] version = 1 -codec = "openai_chat" [[components.config.profiles]] mode = "builtin" @@ -220,7 +250,9 @@ The `builtin` section contains: When a target matches: - Object fields are removed. -- Array elements become `null`. +- Array elements become `null` on JSON-native surfaces. If a normalized LLM + array removal cannot be safely re-encoded, Relay omits the event payload and + its annotation instead. - Targeted scalar or root values become `null`. ### `regex_replace` @@ -290,7 +322,10 @@ The plugin uses different payload boundaries for tools and LLMs: - Tools use JSON-native payloads. Paths point into the emitted tool args or tool result shape directly. -- LLMs use the selected built-in codec. Prefer normalized Relay paths such as: +- LLM requests use the active resolved codec for that managed call. LLM + responses use the active built-in codec for normalized projection. Prefer + normalized Relay paths such as: + - `/headers/authorization` for an exact request header field - `/messages/0/content` for request message content - `/message` for the normalized assistant response text - Marks and non-tool, non-LLM scopes sanitize `data`, `category_profile`, and @@ -299,8 +334,14 @@ The plugin uses different payload boundaries for tools and LLMs: `/metadata`. The current implementation also preserves provider-shaped response-path -compatibility for the supported codecs, but normalized LLM paths are the -recommended contract for new configuration. +compatibility for the supported built-in codecs, but normalized LLM paths are +the recommended contract for new configuration. If normalized paths are +configured and Relay cannot resolve a request codec or cannot project a +response through a recognized built-in or compatible fallback codec, the LLM +event payload and its annotation are omitted rather than emitting unsanitized +content. In particular, a manual LLM call with normalized `target_paths`, no +active codec, and no configured `codec` fallback omits both the payload and its +annotation. ## Decision Rule @@ -355,7 +396,7 @@ That means: - The emitted NeMo Relay start, end, or mark event observability fields are sanitized. - `annotated_response` is regenerated from the sanitized end-event payload when - the configured response codec can decode that payload. + the active response codec can decode that payload. - Subscribers and exporters consume only the sanitized canonical event stream. Each exporter can omit or transform fields as part of its normal projection, but raw values removed by the plugin do not bypass the sanitizer. diff --git a/docs/reference/llm-request-intercept-outcomes.mdx b/docs/reference/llm-request-intercept-outcomes.mdx index 5e968a4b8..734a6327d 100644 --- a/docs/reference/llm-request-intercept-outcomes.mdx +++ b/docs/reference/llm-request-intercept-outcomes.mdx @@ -91,12 +91,11 @@ or object shape: - Python callbacks return `LLMRequestInterceptOutcome`. - Rust callbacks return `LlmRequestInterceptOutcome`. -- Go callbacks return `LLMRequestInterceptOutcome`. - Node.js callbacks return `{ request, annotated?, pendingMarks? }`. JavaScript pending-mark DTOs use `categoryProfile`; canonical JSON retains `pending_marks` and `category_profile`. - Public C callbacks return one owned canonical outcome JSON string, and native - ABI v1 callbacks return one host-owned outcome JSON string. + ABI v2 callbacks return one host-owned outcome JSON string. - Rust and Python `grpc-v1` worker SDKs return their canonical outcome in a `JsonEnvelope` with schema `nemo.relay.LlmRequestInterceptOutcome@2`. @@ -121,10 +120,10 @@ codec input, sanitizer input, or start payload. ## Migration -This finalizes unpublished native ABI v1 and `grpc-v1` contracts. Rebuild all +This finalizes unpublished native ABI v2 and `grpc-v1` contracts. Rebuild all development native plugins and workers against the same NeMo Relay release that -hosts them. Replace tuple results, split C/Go -outputs, metadata envelopes, and parallel mark-aware registrations with the +hosts them. Replace tuple results, split outputs, metadata envelopes, and +parallel mark-aware registrations with the canonical outcome and the existing `register_llm_request_intercept` registration name. diff --git a/docs/reference/migration-guides.mdx b/docs/reference/migration-guides.mdx index debc2cda7..edfeada6f 100644 --- a/docs/reference/migration-guides.mdx +++ b/docs/reference/migration-guides.mdx @@ -1,6 +1,6 @@ --- title: "Migration Guides" -description: "Upgrade NeMo Relay integrations, plugins, exporters, and public API consumers." +description: "Upgrade NeMo Relay integrations and migrate LLM sanitizer callbacks, plugins, workers, and PII policy." position: 6 --- {/* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. @@ -13,7 +13,292 @@ intervening release in sequence. ## Upgrade to NeMo Relay 0.7 -Migration guidance will be added before the 0.7 release. +NeMo Relay 0.7 changes the LLM observability sanitizer contract across +in-process bindings, native plugins, raw C FFI consumers, and worker plugins. +Complete the following migrations before you run an existing sanitizer with a +0.7 host. + + +Do not deploy a 0.6 sanitizer plugin or worker against a 0.7 host. The LLM +callback signature, native ABI layout, and worker invocation schema changed. +NeMo Relay does not adapt one-argument LLM sanitizer callbacks. + + +### Update LLM Sanitizer Callbacks + +The registration names remain unchanged for global, plugin-context, and +scope-local sanitizers. Update the callback itself as follows: + +| Surface | 0.6 Contract | 0.7 Contract | +| --- | --- | --- | +| Request | `(LlmRequest) -> LlmRequest` | `(LlmRequest, LlmSanitizeRequestContext) -> Option` | +| Response | `(Json) -> Json` | `(Json, LlmSanitizeResponseContext) -> Option` | + +The payload is always the first argument. The directional context is always the +second argument. Update request sanitizers with the following pattern: + + + + +```rust +use std::sync::Arc; + +use nemo_relay::api::registry::register_llm_sanitize_request_guardrail; + +register_llm_sanitize_request_guardrail( + "redact-request", + 10, + Arc::new(|request, context| { + let _active_codec = context.resolve_codec(); + // Apply policy, using _active_codec when normalized access is required. + Some(request) + }), +)?; +``` + + + + +```python +import nemo_relay +from nemo_relay import LLMRequest +from nemo_relay import guardrails + +def redact_request( + request: LLMRequest, + context: nemo_relay.LlmSanitizeRequestContext, +) -> LLMRequest | None: + _active_codec = context.resolve_codec() + # Apply policy, using _active_codec when normalized access is required. + return request + +guardrails.register_llm_sanitize_request( + "redact-request", + 10, + redact_request, +) +``` + + + + +```js +const relay = require("nemo-relay-node"); + +relay.registerLlmSanitizeRequestGuardrail( + "redact-request", + 10, + (request, context) => { + const _activeCodec = context.resolveCodec(); + // Apply policy, using _activeCodec when normalized access is required. + return request; + }, +); +``` + + + + +Apply the same change to response sanitizers with +`LlmSanitizeResponseContext`. Return the response payload to keep it in the +observability event. + +### Handle Payload Omission + +In 0.7, `None` or `null` has a specific fail-closed meaning for an LLM +sanitizer: Relay omits the observability payload and its annotation. Omission +short-circuits the remaining LLM sanitizer chain but does not change the +request or response returned to the client. + +Check every callback for an implicit empty return. In particular: + +- A Python function that reaches the end without `return` omits the payload. +- A JavaScript callback that returns `null` or `undefined` omits the payload. +- A Rust callback must return `Some(payload)` to retain the payload. +- A sanitizer error reported through a plugin or binding boundary also omits + the payload and annotation. + +Use omission only when recording the payload would be unsafe. + +### Resolve the Active Codec + +The request and response contexts expose `codec`, which has one of the +following identities: + +| Identity | Meaning | `resolve_codec()` | +| --- | --- | --- | +| `none` | No codec is active for this payload direction. | Returns no codec. | +| `builtin(id)` | A built-in codec is active. | Returns the active codec. | +| `runtime(id)` | A named runtime-registered codec is active. | Returns the active codec. | +| `opaque` | A codec without a public ID is active. | Returns the active codec. | + +Treat identity as descriptive and codec resolution as authoritative. Do not +infer a provider from the request shape. A request codec supports +`decode(request)` and `encode(annotated, original)`. A response codec supports +`decode_response(response)` in in-process Rust and Python, +`decode(response).await` in Rust workers, `await decode(response)` in Python +workers, and `decodeResponse(response)` in Node.js; response encoding is not +available. + +Worker codec proxies are scoped to the sanitizer callback and must not be +stored for later work. Worker codec operations are asynchronous because they +call the host over `grpc-v1`. In-process resolvers return the active codec +facade directly. + +For complete in-process examples, refer to +[Codec-Aware LLM Sanitizers](/about-nemo-relay/concepts/middleware#codec-aware-llm-sanitizers). + +### Migrate Worker Sanitizers + +All Rust worker sanitizer registrations now require callbacks that return +futures. This change applies to mark, scope-start, scope-end, tool-request, +tool-response, LLM-request, and LLM-response sanitizers. Conditional guardrails +and request intercepts keep their existing synchronous contracts. + +Update Rust worker callbacks to use `async move` and return `Result` from the +future. Python worker sanitizers can return either an immediate value or an +awaitable, but Python LLM sanitizers must still accept both the payload and +directional context. + + + + +```rust +ctx.register_llm_sanitize_request_guardrail( + "redact-request", + 10, + |request, context| async move { + if let Some(codec) = context.resolve_codec() { + let annotated = codec.decode(&request).await?; + let request = codec.encode(&annotated, &request).await?; + return Ok(Some(request)); + } + Ok(Some(request)) + }, +); +``` + + + + +```python +from nemo_relay_plugin import LlmSanitizeRequestContext + +async def redact_request(request, context: LlmSanitizeRequestContext): + codec = context.resolve_codec() + if codec is None: + return request + annotated = await codec.decode(request) + return await codec.encode(annotated, request) + +ctx.register_llm_sanitize_request_guardrail( + "redact-request", + redact_request, + priority=10, +) +``` + + + + +Regenerate worker bindings from the 0.7 +`nemo.relay.worker.v1` definition and deploy the 0.7 host and worker SDK +together. The protocol identifier remains `grpc-v1`, but the wire contract now +includes: + +- `LlmCodecIdentity` +- `LlmSanitizeRequestContext` +- `LlmSanitizeResponseContext` +- Host RPCs for request decode, request encode, and response decode + +The previous codec fields on `LlmInvocation` are reserved. Read codec identity +and capability data from the directional sanitizer context instead. + +The codec capability ID in the protocol is SDK-internal. Do not expose or +persist it in plugin code. The host rejects forged, expired, unauthorized, and +wrong-direction capability IDs. + +### Rebuild Native and Raw FFI Plugins + +NeMo Relay 0.7 uses native ABI v2. Recompile native plugins against the 0.7 +`nemo-relay-plugin` crate and rebuild raw FFI consumers against the generated +0.7 header. + +If you already built a plugin against an earlier 0.7 ABI v2 prerelease, rebuild +it again. The ABI version remains 2, but the prerelease LLM sanitizer callback +slots and context layouts changed before release. + +The plugin manifest value remains `compat.native_api = "1"`. This manifest +contract version is separate from the host ABI version; do not change it to +`"2"`. + +Request and response callbacks now receive distinct context structures. Each +structure contains structured codec identity and a borrowed directional codec +handle. The request handle supports decode and encode host operations. The +response handle supports decode. A successful null sanitizer output omits the +observability payload and annotation. + +Do not retain the codec handle, SDK facade, input pointers, or borrowed codec ID +after the callback returns. Release host-owned output strings with the standard +host string release operation. + +For the complete ABI contract, refer to +[Native ABI v2](/build-plugins/dynamic-plugins/native-dynamic/about#native-abi-v2). + +### Update PII Redaction Configuration + +The PII `codec` field is now an optional compatibility fallback. Relay uses it +only for a legacy or manual managed LLM call with no active codec. Any active +codec takes precedence, including runtime and opaque codecs. + +For a mixed-provider gateway, remove the fixed `codec` field: + + + + +```toml +[[components]] +kind = "pii_redaction" +enabled = true + +[components.config] + +[[components.config.profiles]] +mode = "builtin" +priority = 80 + +[components.config.profiles.builtin] +action = "redact" +detector = "email" +target_paths = ["/messages/0/content", "/message"] +``` + + + + +Provider-agnostic all-leaf and `trajectory_context` policies do not require +`codec`. A normalized `target_paths` policy fails closed when Relay has neither +a recognized active codec nor a compatible fallback: Relay omits the LLM +payload and annotation instead of recording unsanitized data. Normalized +`target_paths` response policies for runtime and opaque codecs also omit the +payload because response codec capabilities are decode-only. + +For the complete policy behavior, refer to +[PII Redaction Configuration](/configure-plugins/pii-redaction/configuration). + +### Verify the Upgrade + +Before deployment: + +1. Search for every LLM sanitize-request and sanitize-response registration. +2. Add the required directional context argument and optional payload result. +3. Check that no callback implicitly returns `None`, `null`, or `undefined`. +4. Convert every Rust worker sanitizer callback to an asynchronous callback. +5. Rebuild native and raw FFI plugins against 0.7. +6. Regenerate and redeploy workers with the matching 0.7 protocol and SDK. +7. Remove fixed PII `codec` values from mixed-provider configurations. +8. Test buffered and streaming calls for every provider and custom codec that + your deployment uses. ## Related Release Information diff --git a/examples/rust-native-plugin/src/lib.rs b/examples/rust-native-plugin/src/lib.rs index af19ee8f7..84f51410c 100644 --- a/examples/rust-native-plugin/src/lib.rs +++ b/examples/rust-native-plugin/src/lib.rs @@ -234,11 +234,19 @@ impl NativePlugin for ExampleNativePlugin { ctx.register_llm_sanitize_request_guardrail("example_llm_sanitize_request", 10, { let tag = config.tag.clone(); - move |request| tag_llm_request(request, "native_llm_sanitize_request", &tag) + move |request, _context| { + Some(tag_llm_request( + request, + "native_llm_sanitize_request", + &tag, + )) + } })?; ctx.register_llm_sanitize_response_guardrail("example_llm_sanitize_response", 10, { let tag = config.tag.clone(); - move |response| tag_json(response, "native_llm_sanitize_response", &tag) + move |response, _context| { + Some(tag_json(response, "native_llm_sanitize_response", &tag)) + } })?; ctx.register_llm_conditional_execution_guardrail("example_llm_conditional", 10, { let block_llms = config.block_llms; diff --git a/go/nemo_relay/adaptive_plugin_test.go b/go/nemo_relay/adaptive_plugin_test.go index 58da74050..9931d0187 100644 --- a/go/nemo_relay/adaptive_plugin_test.go +++ b/go/nemo_relay/adaptive_plugin_test.go @@ -90,8 +90,8 @@ func registerLifecycleGuardrails(ctx *PluginContext) error { if err := ctx.RegisterLlmSanitizeRequestGuardrail( "llm_sanitize_request", 7, - func(headers, content json.RawMessage) (json.RawMessage, json.RawMessage) { - return headers, content + func(request LLMRequestDTO, _ LLMSanitizeRequestContext) (LLMRequestDTO, bool) { + return request, false }, ); err != nil { return err @@ -99,7 +99,9 @@ func registerLifecycleGuardrails(ctx *PluginContext) error { if err := ctx.RegisterLlmSanitizeResponseGuardrail( "llm_sanitize_response", 7, - func(responseJSON json.RawMessage) json.RawMessage { return responseJSON }, + func(responseJSON json.RawMessage, _ LLMSanitizeResponseContext) (json.RawMessage, bool) { + return responseJSON, false + }, ); err != nil { return err } @@ -532,10 +534,12 @@ func TestPluginFuncsAndClosedContextBranches(t *testing.T) { return closed.RegisterToolConditionalExecutionGuardrail("tool_conditional", 1, func(name string, args json.RawMessage) *string { return nil }) }}, {"llm sanitize request", func() error { - return closed.RegisterLlmSanitizeRequestGuardrail("llm_sanitize_request", 1, func(headers, content json.RawMessage) (json.RawMessage, json.RawMessage) { return headers, content }) + return closed.RegisterLlmSanitizeRequestGuardrail("llm_sanitize_request", 1, func(request LLMRequestDTO, _ LLMSanitizeRequestContext) (LLMRequestDTO, bool) { return request, false }) }}, {"llm sanitize response", func() error { - return closed.RegisterLlmSanitizeResponseGuardrail("llm_sanitize_response", 1, func(response json.RawMessage) json.RawMessage { return response }) + return closed.RegisterLlmSanitizeResponseGuardrail("llm_sanitize_response", 1, func(response json.RawMessage, _ LLMSanitizeResponseContext) (json.RawMessage, bool) { + return response, false + }) }}, {"llm conditional", func() error { return closed.RegisterLlmConditionalExecutionGuardrail("llm_conditional", 1, func(headers, content json.RawMessage) *string { return nil }) diff --git a/go/nemo_relay/callbacks.go b/go/nemo_relay/callbacks.go index f8069eda0..8096c7447 100644 --- a/go/nemo_relay/callbacks.go +++ b/go/nemo_relay/callbacks.go @@ -26,15 +26,27 @@ typedef struct FfiToolHandle FfiToolHandle; typedef struct FfiLLMHandle FfiLLMHandle; typedef struct FfiLLMRequest FfiLLMRequest; typedef struct FfiEvent FfiEvent; +typedef struct FfiLlmSanitizeRequestCodec FfiLlmSanitizeRequestCodec; +typedef struct FfiLlmSanitizeResponseCodec FfiLlmSanitizeResponseCodec; +typedef struct NemoRelayLlmSanitizeRequestContext { + uint32_t codec_kind; + const char* codec_id; + const FfiLlmSanitizeRequestCodec* codec; +} NemoRelayLlmSanitizeRequestContext; +typedef struct NemoRelayLlmSanitizeResponseContext { + uint32_t codec_kind; + const char* codec_id; + const FfiLlmSanitizeResponseCodec* codec; +} NemoRelayLlmSanitizeResponseContext; typedef void (*NemoRelayFreeFn)(void* user_data); typedef char* (*NemoRelayToolSanitizeFn)(void* user_data, const char* name, const char* args_json); typedef char* (*NemoRelayToolConditionalFn)(void* user_data, const char* name, const char* args_json); typedef char* (*NemoRelayToolExecFn)(void* user_data, const char* args_json); -typedef FfiLLMRequest* (*NemoRelayLlmRequestCb)(void* user_data, const FfiLLMRequest* request); +typedef FfiLLMRequest* (*NemoRelayLlmSanitizeRequestCb)(void* user_data, const FfiLLMRequest* request, NemoRelayLlmSanitizeRequestContext context); typedef char* (*NemoRelayLlmConditionalCb)(void* user_data, const FfiLLMRequest* request); typedef char* (*NemoRelayLlmExecFn)(void* user_data, const char* native_json); -typedef char* (*NemoRelayLlmResponseFn)(void* user_data, const char* response_json); +typedef char* (*NemoRelayLlmSanitizeResponseCb)(void* user_data, const char* response_json, NemoRelayLlmSanitizeResponseContext context); typedef void (*NemoRelayEventSubscriberFn)(void* user_data, const FfiEvent* event); typedef char* (*NemoRelayEventSanitizeFn)(void* user_data, const FfiEvent* event, const char* fields_json); typedef struct FfiPluginContext FfiPluginContext; @@ -57,10 +69,14 @@ static inline char* callLlmExecNext(NemoRelayLlmExecNextFn next_fn, const char* // LLMRequest accessors (also declared in types.go, needed here for trampolines) extern FfiLLMRequest* nemo_relay_llm_request_new(const char* headers_json, const char* content_json); +extern void nemo_relay_llm_request_free(FfiLLMRequest* request); extern char* nemo_relay_llm_request_headers(const FfiLLMRequest* ptr); extern char* nemo_relay_llm_request_content(const FfiLLMRequest* ptr); extern void nemo_relay_string_free(char* ptr); extern void nemo_relay_set_last_error_message(const char* msg); +extern char* nemo_relay_llm_sanitize_request_codec_decode(const FfiLlmSanitizeRequestCodec*, const FfiLLMRequest*); +extern FfiLLMRequest* nemo_relay_llm_sanitize_request_codec_encode(const FfiLlmSanitizeRequestCodec*, const char*, const FfiLLMRequest*); +extern char* nemo_relay_llm_sanitize_response_codec_decode(const FfiLlmSanitizeResponseCodec*, const char*); // Codec callback typedefs (kept for trampoline use at execute time) typedef char* (*NemoRelayCodecDecodeCb)(void* user_data, const FfiLLMRequest* request); @@ -72,6 +88,7 @@ import "C" import ( "encoding/json" + "errors" "sync" "sync/atomic" "unsafe" @@ -162,16 +179,110 @@ type ToolExecutionFunc func(args json.RawMessage) (json.RawMessage, error) // the canonical outcome containing the tool result and any pending marks. type ToolExecutionInterceptFunc func(args json.RawMessage, next func(json.RawMessage) (json.RawMessage, error)) (ToolExecutionInterceptOutcome, error) -// LLMResponseFunc is a callback that transforms an LLM response. It receives -// the response as plain JSON and must return the (possibly modified) response -// JSON. -type LLMResponseFunc func(responseJSON json.RawMessage) json.RawMessage +// LLMCodecKind identifies the active codec state supplied to a sanitizer. +type LLMCodecKind string + +const ( + // LLMCodecNone means no codec was active. + LLMCodecNone LLMCodecKind = "none" + // LLMCodecBuiltin means a Relay built-in codec was active. + LLMCodecBuiltin LLMCodecKind = "builtin" + // LLMCodecRuntime means a runtime-registered codec was active. + LLMCodecRuntime LLMCodecKind = "runtime" + // LLMCodecOpaque means an active codec has no registered identity. + LLMCodecOpaque LLMCodecKind = "opaque" +) + +// LLMCodec identifies the codec active for one managed LLM sanitizer callback. +// ID is present for built-in and runtime codec identities. +type LLMCodec struct { + CodecKind LLMCodecKind + CodecID *string +} + +// LLMSanitizeRequestContext provides request codec context for one sanitizer call. +type LLMSanitizeRequestContext struct { + Codec LLMCodec + resolved *LLMRequestSanitizeCodec +} + +// ResolveCodec returns the active callback-scoped request codec, if any. +func (context LLMSanitizeRequestContext) ResolveCodec() *LLMRequestSanitizeCodec { + return context.resolved +} + +// LLMSanitizeResponseContext provides response codec context for one sanitizer call. +type LLMSanitizeResponseContext struct { + Codec LLMCodec + resolved *LLMResponseSanitizeCodec +} + +// ResolveCodec returns the active callback-scoped response codec, if any. +func (context LLMSanitizeResponseContext) ResolveCodec() *LLMResponseSanitizeCodec { + return context.resolved +} + +// ErrLLMSanitizeCodecExpired is returned when a callback-scoped codec +// capability is used after its sanitizer callback has returned. +var ErrLLMSanitizeCodecExpired = errors.New("LLM sanitizer codec capability is no longer active") + +type llmSanitizeCodecInvocation struct { + mu sync.RWMutex + active bool +} + +func newLLMSanitizeCodecInvocation() *llmSanitizeCodecInvocation { + return &llmSanitizeCodecInvocation{active: true} +} + +func (invocation *llmSanitizeCodecInvocation) acquire() (func(), error) { + invocation.mu.RLock() + if !invocation.active { + invocation.mu.RUnlock() + return nil, ErrLLMSanitizeCodecExpired + } + return invocation.mu.RUnlock, nil +} + +func (invocation *llmSanitizeCodecInvocation) invalidate() { + invocation.mu.Lock() + invocation.active = false + invocation.mu.Unlock() +} + +// LLMRequestSanitizeCodec is a callback-scoped request codec capability. +type LLMRequestSanitizeCodec struct { + ptr unsafe.Pointer + invocation *llmSanitizeCodecInvocation +} + +// LLMResponseSanitizeCodec is a callback-scoped response codec capability. +type LLMResponseSanitizeCodec struct { + ptr unsafe.Pointer + invocation *llmSanitizeCodecInvocation +} + +func acquireLLMSanitizeCodec( + ptr unsafe.Pointer, + invocation *llmSanitizeCodecInvocation, +) (unsafe.Pointer, func(), error) { + if ptr == nil || invocation == nil { + return nil, nil, ErrLLMSanitizeCodecExpired + } + release, err := invocation.acquire() + if err != nil { + return nil, nil, err + } + return ptr, release, nil +} -// LLMRequestFunc is a callback that transforms an LLM request. It receives -// the headers JSON and content JSON from the FfiLLMRequest, and returns the -// (possibly modified) versions of each. The Go binding uses JSON -// serialization rather than opaque C pointers for ergonomics. -type LLMRequestFunc func(headers, content json.RawMessage) (headers2, content2 json.RawMessage) +// LLMRequestFunc sanitizes an emitted LLM request. It receives the request +// first and codec context second; returning omit true removes observability. +type LLMRequestFunc func(request LLMRequestDTO, context LLMSanitizeRequestContext) (sanitized LLMRequestDTO, omit bool) + +// LLMResponseFunc sanitizes an emitted LLM response. It receives the response +// first and codec context second; returning omit true removes observability. +type LLMResponseFunc func(response json.RawMessage, context LLMSanitizeResponseContext) (sanitized json.RawMessage, omit bool) // LLMConditionalFunc is a callback that decides whether an LLM call should // proceed. It returns nil to allow execution, or a non-nil pointer to an error @@ -236,6 +347,86 @@ type LLMRequestDTO struct { Content json.RawMessage `json:"content"` } +// Decode normalizes an opaque request through the active codec. +func (codec *LLMRequestSanitizeCodec) Decode(request LLMRequestDTO) (json.RawMessage, error) { + if codec == nil { + return nil, ErrLLMSanitizeCodecExpired + } + ptr, release, err := acquireLLMSanitizeCodec(codec.ptr, codec.invocation) + if err != nil { + return nil, err + } + defer release() + cHeaders := C.CString(string(request.Headers)) + defer C.free(unsafe.Pointer(cHeaders)) + cContent := C.CString(string(request.Content)) + defer C.free(unsafe.Pointer(cContent)) + cRequest := C.nemo_relay_llm_request_new(cHeaders, cContent) + if cRequest == nil { + return nil, lastError() + } + defer C.nemo_relay_llm_request_free(cRequest) + out := C.nemo_relay_llm_sanitize_request_codec_decode((*C.FfiLlmSanitizeRequestCodec)(ptr), cRequest) + if out == nil { + return nil, lastError() + } + defer C.nemo_relay_string_free(out) + return json.RawMessage(C.GoString(out)), nil +} + +// Encode merges normalized changes onto the original opaque request. +func (codec *LLMRequestSanitizeCodec) Encode(annotated json.RawMessage, original LLMRequestDTO) (LLMRequestDTO, error) { + if codec == nil { + return LLMRequestDTO{}, ErrLLMSanitizeCodecExpired + } + ptr, release, err := acquireLLMSanitizeCodec(codec.ptr, codec.invocation) + if err != nil { + return LLMRequestDTO{}, err + } + defer release() + cHeaders := C.CString(string(original.Headers)) + defer C.free(unsafe.Pointer(cHeaders)) + cContent := C.CString(string(original.Content)) + defer C.free(unsafe.Pointer(cContent)) + cOriginal := C.nemo_relay_llm_request_new(cHeaders, cContent) + if cOriginal == nil { + return LLMRequestDTO{}, lastError() + } + defer C.nemo_relay_llm_request_free(cOriginal) + cAnnotated := C.CString(string(annotated)) + defer C.free(unsafe.Pointer(cAnnotated)) + out := C.nemo_relay_llm_sanitize_request_codec_encode((*C.FfiLlmSanitizeRequestCodec)(ptr), cAnnotated, cOriginal) + if out == nil { + return LLMRequestDTO{}, lastError() + } + defer C.nemo_relay_llm_request_free(out) + headers := C.nemo_relay_llm_request_headers(out) + defer C.nemo_relay_string_free(headers) + content := C.nemo_relay_llm_request_content(out) + defer C.nemo_relay_string_free(content) + return LLMRequestDTO{Headers: json.RawMessage(C.GoString(headers)), Content: json.RawMessage(C.GoString(content))}, nil +} + +// Decode normalizes an opaque response through the active codec. +func (codec *LLMResponseSanitizeCodec) Decode(response json.RawMessage) (json.RawMessage, error) { + if codec == nil { + return nil, ErrLLMSanitizeCodecExpired + } + ptr, release, err := acquireLLMSanitizeCodec(codec.ptr, codec.invocation) + if err != nil { + return nil, err + } + defer release() + cResponse := C.CString(string(response)) + defer C.free(unsafe.Pointer(cResponse)) + out := C.nemo_relay_llm_sanitize_response_codec_decode((*C.FfiLlmSanitizeResponseCodec)(ptr), cResponse) + if out == nil { + return nil, lastError() + } + defer C.nemo_relay_string_free(out) + return json.RawMessage(C.GoString(out)), nil +} + // PendingMarkSpec describes a mark Relay materializes under a managed lifecycle. type PendingMarkSpec struct { Name string `json:"name"` @@ -452,34 +643,97 @@ func goFreeTrampoline(userData unsafe.Pointer) { } //export goLlmRequestTrampoline -func goLlmRequestTrampoline(userData unsafe.Pointer, request *C.FfiLLMRequest) *C.FfiLLMRequest { +func goLlmRequestTrampoline(userData unsafe.Pointer, request *C.FfiLLMRequest, context C.NemoRelayLlmSanitizeRequestContext) *C.FfiLLMRequest { fn := lookupClosure(userData).(LLMRequestFunc) - - // Extract headers and content from the incoming FfiLLMRequest cHeaders := C.nemo_relay_llm_request_headers(request) cContent := C.nemo_relay_llm_request_content(request) - goHeaders := json.RawMessage(C.GoString(cHeaders)) - goContent := json.RawMessage(C.GoString(cContent)) - C.nemo_relay_string_free(cHeaders) - C.nemo_relay_string_free(cContent) - - // Call the Go callback - newHeaders, newContent := fn(goHeaders, goContent) + defer C.nemo_relay_string_free(cHeaders) + defer C.nemo_relay_string_free(cContent) + invocation := newLLMSanitizeCodecInvocation() + defer invocation.invalidate() + sanitized, omit := fn(LLMRequestDTO{ + Headers: json.RawMessage(C.GoString(cHeaders)), + Content: json.RawMessage(C.GoString(cContent)), + }, llmSanitizeRequestContextFromC(context, invocation)) + if omit { + return nil + } + cHeaders = C.CString(string(sanitized.Headers)) + cContent = C.CString(string(sanitized.Content)) + defer C.free(unsafe.Pointer(cHeaders)) + defer C.free(unsafe.Pointer(cContent)) + return C.nemo_relay_llm_request_new(cHeaders, cContent) +} + +func llmSanitizeRequestContextFromC( + context C.NemoRelayLlmSanitizeRequestContext, + invocation *llmSanitizeCodecInvocation, +) LLMSanitizeRequestContext { + var codecID *string + if context.codec_id != nil { + id := C.GoString(context.codec_id) + codecID = &id + } + result := LLMSanitizeRequestContext{Codec: llmCodecIdentity(uint32(context.codec_kind), codecID)} + if context.codec != nil { + result.resolved = &LLMRequestSanitizeCodec{ + ptr: unsafe.Pointer(context.codec), + invocation: invocation, + } + } + return result +} - // Create a new FfiLLMRequest from the result - cNewHeaders := C.CString(string(newHeaders)) - cNewContent := C.CString(string(newContent)) - defer C.free(unsafe.Pointer(cNewHeaders)) - defer C.free(unsafe.Pointer(cNewContent)) - return C.nemo_relay_llm_request_new(cNewHeaders, cNewContent) +func llmCodecIdentity(codecKind uint32, codecID *string) LLMCodec { + kind := LLMCodecOpaque + switch codecKind { + case 0: + kind = LLMCodecNone + case 1: + kind = LLMCodecBuiltin + case 2: + kind = LLMCodecRuntime + case 3: + kind = LLMCodecOpaque + } + return LLMCodec{CodecKind: kind, CodecID: codecID} } //export goLlmResponseTrampoline -func goLlmResponseTrampoline(userData unsafe.Pointer, responseJSON *C.char) *C.char { +func goLlmResponseTrampoline(userData unsafe.Pointer, responseJSON *C.char, context C.NemoRelayLlmSanitizeResponseContext) *C.char { fn := lookupClosure(userData).(LLMResponseFunc) - goJSON := json.RawMessage(C.GoString(responseJSON)) - result := fn(goJSON) - return C.CString(string(result)) + invocation := newLLMSanitizeCodecInvocation() + defer invocation.invalidate() + sanitized, omit := fn( + json.RawMessage(C.GoString(responseJSON)), + llmSanitizeResponseContextFromC(context, invocation), + ) + if omit { + return nil + } + return C.CString(string(sanitized)) +} + +func llmSanitizeResponseContextFromC( + context C.NemoRelayLlmSanitizeResponseContext, + invocation *llmSanitizeCodecInvocation, +) LLMSanitizeResponseContext { + var codecID *string + if context.codec_id != nil { + id := C.GoString(context.codec_id) + codecID = &id + } + resolved := (*LLMResponseSanitizeCodec)(nil) + if context.codec != nil { + resolved = &LLMResponseSanitizeCodec{ + ptr: unsafe.Pointer(context.codec), + invocation: invocation, + } + } + return LLMSanitizeResponseContext{ + Codec: llmCodecIdentity(uint32(context.codec_kind), codecID), + resolved: resolved, + } } //export goLlmConditionalTrampoline diff --git a/go/nemo_relay/callbacks_test.go b/go/nemo_relay/callbacks_test.go index 2655608d0..4be9b5637 100644 --- a/go/nemo_relay/callbacks_test.go +++ b/go/nemo_relay/callbacks_test.go @@ -5,9 +5,40 @@ package nemo_relay import ( "encoding/json" + "errors" "testing" + "time" ) +func TestLLMSanitizeCodecInvocationInvalidationWaitsForInflightCall(t *testing.T) { + invocation := newLLMSanitizeCodecInvocation() + release, err := invocation.acquire() + if err != nil { + t.Fatalf("acquire active invocation: %v", err) + } + + invalidated := make(chan struct{}) + go func() { + invocation.invalidate() + close(invalidated) + }() + + select { + case <-invalidated: + t.Fatal("invalidation returned while a codec operation was in flight") + case <-time.After(20 * time.Millisecond): + } + release() + select { + case <-invalidated: + case <-time.After(time.Second): + t.Fatal("invalidation did not finish after the codec operation completed") + } + if _, err := invocation.acquire(); !errors.Is(err, ErrLLMSanitizeCodecExpired) { + t.Fatalf("expired invocation acquire returned %v", err) + } +} + func toolExecutionOutcome(result json.RawMessage, err error) (ToolExecutionInterceptOutcome, error) { return ToolExecutionInterceptOutcome{Result: result}, err } @@ -36,3 +67,43 @@ func TestRegisterAndUnregisterClosure(t *testing.T) { t.Fatal("closure registry still contains callback after unregister") } } + +func TestLlmSanitizeDirectionalContextsPreserveEveryCodecIdentity(t *testing.T) { + openAIChat := "openai_chat" + openAIResponses := "openai_responses" + anthropicMessages := "anthropic_messages" + runtimeCodec := "com.example.chat.v1" + + cases := []struct { + name string + kind uint32 + id *string + want LLMCodecKind + }{ + {"none", 0, nil, LLMCodecNone}, + {"openai chat", 1, &openAIChat, LLMCodecBuiltin}, + {"openai responses", 1, &openAIResponses, LLMCodecBuiltin}, + {"anthropic messages", 1, &anthropicMessages, LLMCodecBuiltin}, + {"runtime", 2, &runtimeCodec, LLMCodecRuntime}, + {"opaque", 3, nil, LLMCodecOpaque}, + {"unknown", 99, nil, LLMCodecOpaque}, + } + + for _, test := range cases { + t.Run(test.name, func(t *testing.T) { + codec := llmCodecIdentity(test.kind, test.id) + if codec.CodecKind != test.want { + t.Fatalf("codec kind = %q, want %q", codec.CodecKind, test.want) + } + if codec.CodecID == nil && test.id != nil { + t.Fatal("codec ID was lost") + } + if codec.CodecID != nil && test.id == nil { + t.Fatalf("unexpected codec ID %q", *codec.CodecID) + } + if codec.CodecID != nil && test.id != nil && *codec.CodecID != *test.id { + t.Fatalf("codec ID = %q, want %q", *codec.CodecID, *test.id) + } + }) + } +} diff --git a/go/nemo_relay/deregister_test.go b/go/nemo_relay/deregister_test.go index ab6446e2e..e243bed6b 100644 --- a/go/nemo_relay/deregister_test.go +++ b/go/nemo_relay/deregister_test.go @@ -159,8 +159,8 @@ func TestRegisterDeregisterReregisterToolExecutionIntercept(t *testing.T) { func TestRegisterDeregisterReregisterLlmSanitizeRequestGuardrail(t *testing.T) { name := "go_reregister_llm_san_req" - fn := func(headers, content json.RawMessage) (json.RawMessage, json.RawMessage) { - return headers, content + fn := func(request LLMRequestDTO, _ LLMSanitizeRequestContext) (LLMRequestDTO, bool) { + return request, false } err := RegisterLlmSanitizeRequestGuardrail(name, 1, fn) diff --git a/go/nemo_relay/error_test.go b/go/nemo_relay/error_test.go index 1d32b1087..197edef89 100644 --- a/go/nemo_relay/error_test.go +++ b/go/nemo_relay/error_test.go @@ -126,7 +126,7 @@ func TestAlreadyExistsErrorOnDuplicateSubscriber(t *testing.T) { func TestAlreadyExistsErrorOnDuplicateLlmGuardrails(t *testing.T) { t.Run("LlmSanitizeRequest", func(t *testing.T) { name := "go_err_dup_llm_san_req" - fn := func(h, c json.RawMessage) (json.RawMessage, json.RawMessage) { return h, c } + fn := func(request LLMRequestDTO, _ LLMSanitizeRequestContext) (LLMRequestDTO, bool) { return request, false } err := RegisterLlmSanitizeRequestGuardrail(name, 1, fn) if err != nil { @@ -142,7 +142,7 @@ func TestAlreadyExistsErrorOnDuplicateLlmGuardrails(t *testing.T) { t.Run("LlmSanitizeResponse", func(t *testing.T) { name := "go_err_dup_llm_san_resp" - fn := func(r json.RawMessage) json.RawMessage { return r } + fn := func(r json.RawMessage, _ LLMSanitizeResponseContext) (json.RawMessage, bool) { return r, false } err := RegisterLlmSanitizeResponseGuardrail(name, 1, fn) if err != nil { diff --git a/go/nemo_relay/guardrails/guardrails_test.go b/go/nemo_relay/guardrails/guardrails_test.go index 002906950..c7bc5470d 100644 --- a/go/nemo_relay/guardrails/guardrails_test.go +++ b/go/nemo_relay/guardrails/guardrails_test.go @@ -120,12 +120,13 @@ func runGlobalLLMGuardrailShorthandChecks(t *testing.T, output func() json.RawMe t.Helper() if err := guardrails.RegisterLlmSanitizeRequest("guardrails_llm_req", 1, - func(headers, content json.RawMessage) (json.RawMessage, json.RawMessage) { + func(request nemo_relay.LLMRequestDTO, _ nemo_relay.LLMSanitizeRequestContext) (nemo_relay.LLMRequestDTO, bool) { var payload map[string]interface{} - _ = json.Unmarshal(content, &payload) + _ = json.Unmarshal(request.Content, &payload) payload["request_sanitized"] = true out, _ := json.Marshal(payload) - return headers, out + request.Content = out + return request, false }, ); err != nil { t.Fatalf("RegisterLlmSanitizeRequest failed: %v", err) @@ -135,12 +136,12 @@ func runGlobalLLMGuardrailShorthandChecks(t *testing.T, output func() json.RawMe }) if err := guardrails.RegisterLlmSanitizeResponse("guardrails_llm_resp", 1, - func(response json.RawMessage) json.RawMessage { + func(response json.RawMessage, _ nemo_relay.LLMSanitizeResponseContext) (json.RawMessage, bool) { var payload map[string]interface{} _ = json.Unmarshal(response, &payload) payload["guarded"] = true out, _ := json.Marshal(payload) - return out + return out, false }, ); err != nil { t.Fatalf("RegisterLlmSanitizeResponse failed: %v", err) @@ -212,14 +213,16 @@ func runScopeLocalLLMGuardrailShorthandChecks(t *testing.T, scopeUUID string) { t.Helper() if err := guardrails.ScopeRegisterLlmSanitizeRequest(scopeUUID, "guardrails_scope_llm_req", 1, - func(headers, content json.RawMessage) (json.RawMessage, json.RawMessage) { - return headers, content + func(request nemo_relay.LLMRequestDTO, _ nemo_relay.LLMSanitizeRequestContext) (nemo_relay.LLMRequestDTO, bool) { + return request, false }, ); err != nil { t.Fatalf("ScopeRegisterLlmSanitizeRequest failed: %v", err) } if err := guardrails.ScopeRegisterLlmSanitizeResponse(scopeUUID, "guardrails_scope_llm_resp", 1, - func(response json.RawMessage) json.RawMessage { return response }, + func(response json.RawMessage, _ nemo_relay.LLMSanitizeResponseContext) (json.RawMessage, bool) { + return response, false + }, ); err != nil { t.Fatalf("ScopeRegisterLlmSanitizeResponse failed: %v", err) } diff --git a/go/nemo_relay/llm_test.go b/go/nemo_relay/llm_test.go index 9bb402b33..bf2a60fe4 100644 --- a/go/nemo_relay/llm_test.go +++ b/go/nemo_relay/llm_test.go @@ -282,6 +282,141 @@ func TestLlmCallExecuteWithRequestAndResponseCodecs(t *testing.T) { } } +func TestLlmSanitizersResolveDirectionalCodecs(t *testing.T) { + const requestGuard = "go_llm_resolved_request_codec" + const responseGuard = "go_llm_resolved_response_codec" + _ = DeregisterLlmSanitizeRequestGuardrail(requestGuard) + _ = DeregisterLlmSanitizeResponseGuardrail(responseGuard) + defer DeregisterLlmSanitizeRequestGuardrail(requestGuard) + defer DeregisterLlmSanitizeResponseGuardrail(responseGuard) + + var callbackState struct { + sync.Mutex + requestResolved bool + responseResolved bool + retainedRequestCodec *LLMRequestSanitizeCodec + retainedResponseCodec *LLMResponseSanitizeCodec + retainedRequest LLMRequestDTO + errors []string + } + recordCallbackError := func(format string, args ...any) { + callbackState.Lock() + defer callbackState.Unlock() + callbackState.errors = append(callbackState.errors, fmt.Sprintf(format, args...)) + } + if err := RegisterLlmSanitizeRequestGuardrail( + requestGuard, + 0, + func(request LLMRequestDTO, context LLMSanitizeRequestContext) (LLMRequestDTO, bool) { + if context.Codec.CodecKind != LLMCodecOpaque || + context.Codec.CodecID != nil { + recordCallbackError("unexpected request codec identity: %#v", context.Codec) + } + codec := context.ResolveCodec() + if codec == nil { + recordCallbackError("active request codec did not resolve") + return request, false + } + callbackState.Lock() + callbackState.retainedRequestCodec = codec + callbackState.retainedRequest = request + callbackState.Unlock() + annotated, err := codec.Decode(request) + if err != nil { + recordCallbackError("request codec decode failed: %v", err) + return request, false + } + encoded, err := codec.Encode(annotated, request) + if err != nil { + recordCallbackError("request codec encode failed: %v", err) + return request, false + } + callbackState.Lock() + callbackState.requestResolved = true + callbackState.Unlock() + return encoded, false + }, + ); err != nil { + t.Fatalf("request sanitizer registration failed: %v", err) + } + if err := RegisterLlmSanitizeResponseGuardrail( + responseGuard, + 0, + func(response json.RawMessage, context LLMSanitizeResponseContext) (json.RawMessage, bool) { + if context.Codec.CodecKind != LLMCodecBuiltin || + context.Codec.CodecID == nil || + *context.Codec.CodecID != "openai_chat" { + recordCallbackError("unexpected response codec identity: %#v", context.Codec) + } + codec := context.ResolveCodec() + if codec == nil { + recordCallbackError("active response codec did not resolve") + return response, false + } + callbackState.Lock() + callbackState.retainedResponseCodec = codec + callbackState.Unlock() + if _, err := codec.Decode(response); err != nil { + recordCallbackError("response codec decode failed: %v", err) + return response, false + } + callbackState.Lock() + callbackState.responseResolved = true + callbackState.Unlock() + return response, false + }, + ); err != nil { + t.Fatalf("response sanitizer registration failed: %v", err) + } + + response := json.RawMessage(`{ + "id":"chatcmpl-test", + "model":"test-model", + "choices":[{ + "index":0, + "message":{"role":"assistant","content":"ok"}, + "finish_reason":"stop" + }] + }`) + _, err := LlmCallExecute( + "resolved_codec_llm", + makeRequest(), + func(json.RawMessage) (json.RawMessage, error) { return response, nil }, + WithLLMCodec(llmRequestResponseCodec()), + WithLLMResponseCodec(NewOpenAIChatCodec()), + ) + if err != nil { + t.Fatalf(llmCallExecuteFailed, err) + } + callbackState.Lock() + sanitizerErrors := append([]string(nil), callbackState.errors...) + requestResolved := callbackState.requestResolved + responseResolved := callbackState.responseResolved + retainedRequestCodec := callbackState.retainedRequestCodec + retainedResponseCodec := callbackState.retainedResponseCodec + retainedRequest := callbackState.retainedRequest + callbackState.Unlock() + if len(sanitizerErrors) != 0 { + t.Fatalf("sanitizer callbacks failed: %v", sanitizerErrors) + } + if !requestResolved || !responseResolved { + t.Fatalf( + "expected both codec capabilities to resolve, request=%t response=%t", + requestResolved, + responseResolved, + ) + } + if _, err := retainedRequestCodec.Decode(retainedRequest); !errors.Is(err, ErrLLMSanitizeCodecExpired) { + t.Fatalf("retained request codec must expire after callback, got %v", err) + } + if _, err := retainedRequestCodec.Encode(json.RawMessage(`{}`), retainedRequest); !errors.Is(err, ErrLLMSanitizeCodecExpired) { + t.Fatalf("retained request codec encode must expire after callback, got %v", err) + } + if _, err := retainedResponseCodec.Decode(response); !errors.Is(err, ErrLLMSanitizeCodecExpired) { + t.Fatalf("retained response codec must expire after callback, got %v", err) + } +} + func llmRequestResponseCodec() CodecFunc { return CodecFunc{ Decode: func(headersJSON, contentJSON json.RawMessage) (json.RawMessage, error) { @@ -365,8 +500,8 @@ func requireLlmScopeEvents(t *testing.T, events []Event) (*ScopeEvent, *ScopeEve func TestLlmSanitizeRequestGuardrail(t *testing.T) { err := RegisterLlmSanitizeRequestGuardrail("go_llm_san_req", 1, - func(headers, content json.RawMessage) (json.RawMessage, json.RawMessage) { - return headers, content + func(request LLMRequestDTO, _ LLMSanitizeRequestContext) (LLMRequestDTO, bool) { + return request, false }, ) if err != nil { @@ -377,7 +512,9 @@ func TestLlmSanitizeRequestGuardrail(t *testing.T) { func TestLlmSanitizeResponseGuardrail(t *testing.T) { err := RegisterLlmSanitizeResponseGuardrail("go_llm_san_resp", 1, - func(responseJSON json.RawMessage) json.RawMessage { return responseJSON }, + func(responseJSON json.RawMessage, _ LLMSanitizeResponseContext) (json.RawMessage, bool) { + return responseJSON, false + }, ) if err != nil { t.Fatalf(llmRegisterFailed, err) @@ -385,6 +522,81 @@ func TestLlmSanitizeResponseGuardrail(t *testing.T) { DeregisterLlmSanitizeResponseGuardrail("go_llm_san_resp") } +func TestLlmSanitizeGuardrailsReceiveContext(t *testing.T) { + var capturedInput, capturedOutput json.RawMessage + var mu sync.Mutex + var callbackErrorsMu sync.Mutex + var callbackErrors []string + recordCallbackError := func(message string) { + callbackErrorsMu.Lock() + defer callbackErrorsMu.Unlock() + callbackErrors = append(callbackErrors, message) + } + if err := RegisterSubscriber("go_contextual_llm_sanitize_events", func(event Event) { + mu.Lock() + defer mu.Unlock() + if event.Kind() == "scope" && event.Category() == "llm" && event.ScopeCategory() == "start" { + capturedInput = append(json.RawMessage(nil), event.Input()...) + } + if event.Kind() == "scope" && event.Category() == "llm" && event.ScopeCategory() == "end" { + capturedOutput = append(json.RawMessage(nil), event.Output()...) + } + }); err != nil { + t.Fatalf("RegisterSubscriber failed: %v", err) + } + defer DeregisterSubscriber("go_contextual_llm_sanitize_events") + + if err := RegisterLlmSanitizeRequestGuardrail("go_contextual_llm_request", 1, + func(request LLMRequestDTO, context LLMSanitizeRequestContext) (LLMRequestDTO, bool) { + if context.Codec.CodecKind != LLMCodecNone { + recordCallbackError("manual registration received an active codec identity") + } + return request, true + }, + ); err != nil { + t.Fatalf(llmRegisterFailed, err) + } + defer DeregisterLlmSanitizeRequestGuardrail("go_contextual_llm_request") + + if err := RegisterLlmSanitizeResponseGuardrail("go_contextual_llm_response", 1, + func(response json.RawMessage, context LLMSanitizeResponseContext) (json.RawMessage, bool) { + if context.Codec.CodecID != nil { + recordCallbackError("manual registration received a codec ID") + } + return response, true + }, + ); err != nil { + t.Fatalf(llmRegisterFailed, err) + } + defer DeregisterLlmSanitizeResponseGuardrail("go_contextual_llm_response") + + result, err := LlmCallExecute("go_contextual_llm_sanitize", makeRequest(), + func(nativeJSON json.RawMessage) (json.RawMessage, error) { + return json.RawMessage(`{"response":"client-visible"}`), nil + }, + ) + if err != nil { + t.Fatalf(llmCallExecuteFailed, err) + } + callbackErrorsMu.Lock() + sanitizerErrors := append([]string(nil), callbackErrors...) + callbackErrorsMu.Unlock() + if len(sanitizerErrors) != 0 { + t.Fatalf("sanitizer callbacks failed: %v", sanitizerErrors) + } + if string(result) != `{"response":"client-visible"}` { + t.Fatalf("contextual sanitizers must not change the client result: %s", result) + } + if err := FlushSubscribers(); err != nil { + t.Fatalf(llmFlushSubscribersFailed, err) + } + mu.Lock() + defer mu.Unlock() + if capturedInput != nil || capturedOutput != nil { + t.Fatalf("contextual omission must remove observability payloads, got input=%s output=%s", capturedInput, capturedOutput) + } +} + func TestLlmConditionalExecutionGuardrail(t *testing.T) { err := RegisterLlmConditionalExecutionGuardrail("go_llm_cond", 1, func(headers, content json.RawMessage) *string { @@ -399,13 +611,13 @@ func TestLlmConditionalExecutionGuardrail(t *testing.T) { func TestLlmDuplicateGuardrailFails(t *testing.T) { RegisterLlmSanitizeRequestGuardrail("go_llm_dup", 1, - func(headers, content json.RawMessage) (json.RawMessage, json.RawMessage) { - return headers, content + func(request LLMRequestDTO, _ LLMSanitizeRequestContext) (LLMRequestDTO, bool) { + return request, false }, ) err := RegisterLlmSanitizeRequestGuardrail("go_llm_dup", 1, - func(headers, content json.RawMessage) (json.RawMessage, json.RawMessage) { - return headers, content + func(request LLMRequestDTO, _ LLMSanitizeRequestContext) (LLMRequestDTO, bool) { + return request, false }, ) if err == nil { @@ -660,12 +872,13 @@ func TestLlmSanitizeRequestGuardrailModifiesEventInput(t *testing.T) { defer DeregisterSubscriber("go_llm_san_evt_sub") RegisterLlmSanitizeRequestGuardrail("go_llm_content_mod", 1, - func(headers, content json.RawMessage) (json.RawMessage, json.RawMessage) { + func(request LLMRequestDTO, _ LLMSanitizeRequestContext) (LLMRequestDTO, bool) { var m map[string]interface{} - json.Unmarshal(content, &m) + json.Unmarshal(request.Content, &m) m["system_prompt_injected"] = true out, _ := json.Marshal(m) - return headers, out + request.Content = out + return request, false }, ) defer DeregisterLlmSanitizeRequestGuardrail("go_llm_content_mod") diff --git a/go/nemo_relay/nemo_relay.go b/go/nemo_relay/nemo_relay.go index 727f8380c..937427e5d 100644 --- a/go/nemo_relay/nemo_relay.go +++ b/go/nemo_relay/nemo_relay.go @@ -38,6 +38,10 @@ typedef struct FfiLLMRequest FfiLLMRequest; typedef struct FfiEvent FfiEvent; typedef struct FfiStream FfiStream; typedef struct FfiCodecHandle FfiCodecHandle; +typedef struct FfiLlmSanitizeRequestCodec FfiLlmSanitizeRequestCodec; +typedef struct FfiLlmSanitizeResponseCodec FfiLlmSanitizeResponseCodec; +typedef struct NemoRelayLlmSanitizeRequestContext { uint32_t codec_kind; const char* codec_id; const FfiLlmSanitizeRequestCodec* codec; } NemoRelayLlmSanitizeRequestContext; +typedef struct NemoRelayLlmSanitizeResponseContext { uint32_t codec_kind; const char* codec_id; const FfiLlmSanitizeResponseCodec* codec; } NemoRelayLlmSanitizeResponseContext; typedef void (*NemoRelayFreeFn)(void* user_data); @@ -135,12 +139,12 @@ extern int32_t nemo_relay_register_tool_execution_intercept(const char* name, in extern int32_t nemo_relay_deregister_tool_execution_intercept(const char* name); // LLM guardrails -typedef FfiLLMRequest* (*NemoRelayLlmRequestCb)(void* user_data, const FfiLLMRequest* request); -extern int32_t nemo_relay_register_llm_sanitize_request_guardrail(const char* name, int32_t priority, NemoRelayLlmRequestCb cb, void* user_data, NemoRelayFreeFn free_fn); +typedef FfiLLMRequest* (*NemoRelayLlmSanitizeRequestCb)(void* user_data, const FfiLLMRequest* request, NemoRelayLlmSanitizeRequestContext context); +extern int32_t nemo_relay_register_llm_sanitize_request_guardrail(const char* name, int32_t priority, NemoRelayLlmSanitizeRequestCb cb, void* user_data, NemoRelayFreeFn free_fn); extern int32_t nemo_relay_deregister_llm_sanitize_request_guardrail(const char* name); -typedef char* (*NemoRelayLlmResponseFn)(void* user_data, const char* response_json); -extern int32_t nemo_relay_register_llm_sanitize_response_guardrail(const char* name, int32_t priority, NemoRelayLlmResponseFn cb, void* user_data, NemoRelayFreeFn free_fn); +typedef char* (*NemoRelayLlmSanitizeResponseCb)(void* user_data, const char* response_json, NemoRelayLlmSanitizeResponseContext context); +extern int32_t nemo_relay_register_llm_sanitize_response_guardrail(const char* name, int32_t priority, NemoRelayLlmSanitizeResponseCb cb, void* user_data, NemoRelayFreeFn free_fn); extern int32_t nemo_relay_deregister_llm_sanitize_response_guardrail(const char* name); typedef char* (*NemoRelayLlmConditionalCb)(void* user_data, const FfiLLMRequest* request); @@ -193,9 +197,9 @@ extern int32_t nemo_relay_scope_register_tool_execution_intercept(const char* sc extern int32_t nemo_relay_scope_deregister_tool_execution_intercept(const char* scope_uuid, const char* name); // Scope-local LLM guardrails -extern int32_t nemo_relay_scope_register_llm_sanitize_request_guardrail(const char* scope_uuid, const char* name, int32_t priority, NemoRelayLlmRequestCb cb, void* user_data, NemoRelayFreeFn free_fn); +extern int32_t nemo_relay_scope_register_llm_sanitize_request_guardrail(const char* scope_uuid, const char* name, int32_t priority, NemoRelayLlmSanitizeRequestCb cb, void* user_data, NemoRelayFreeFn free_fn); extern int32_t nemo_relay_scope_deregister_llm_sanitize_request_guardrail(const char* scope_uuid, const char* name); -extern int32_t nemo_relay_scope_register_llm_sanitize_response_guardrail(const char* scope_uuid, const char* name, int32_t priority, NemoRelayLlmResponseFn cb, void* user_data, NemoRelayFreeFn free_fn); +extern int32_t nemo_relay_scope_register_llm_sanitize_response_guardrail(const char* scope_uuid, const char* name, int32_t priority, NemoRelayLlmSanitizeResponseCb cb, void* user_data, NemoRelayFreeFn free_fn); extern int32_t nemo_relay_scope_deregister_llm_sanitize_response_guardrail(const char* scope_uuid, const char* name); extern int32_t nemo_relay_scope_register_llm_conditional_execution_guardrail(const char* scope_uuid, const char* name, int32_t priority, NemoRelayLlmConditionalCb cb, void* user_data, NemoRelayFreeFn free_fn); extern int32_t nemo_relay_scope_deregister_llm_conditional_execution_guardrail(const char* scope_uuid, const char* name); @@ -274,8 +278,8 @@ extern char* goToolConditionalTrampoline(void*, const char*, const char*); extern char* goToolExecTrampoline(void*, const char*); extern void goEventSubscriberTrampoline(void*, const FfiEvent*); extern void goFreeTrampoline(void*); -extern FfiLLMRequest* goLlmRequestTrampoline(void*, const FfiLLMRequest*); -extern char* goLlmResponseTrampoline(void*, const char*); +extern FfiLLMRequest* goLlmRequestTrampoline(void*, const FfiLLMRequest*, NemoRelayLlmSanitizeRequestContext); +extern char* goLlmResponseTrampoline(void*, const char*, NemoRelayLlmSanitizeResponseContext); extern char* goLlmConditionalTrampoline(void*, const FfiLLMRequest*); extern char* goLlmExecTrampoline(void*, const char*); extern char* goToolExecInterceptTrampoline(void*, const char*, NemoRelayToolExecNextFn, void*); @@ -1367,17 +1371,15 @@ func DeregisterToolExecutionIntercept(name string) error { // Guardrail/Intercept registration (LLM) // --------------------------------------------------------------------------- -// RegisterLlmSanitizeRequestGuardrail registers a guardrail that sanitizes LLM -// request data before the call is made. The callback receives the request -// headers and content JSON and must return the (possibly modified) versions. -// Guardrails are invoked in priority order (lower values run first). +// RegisterLlmSanitizeRequestGuardrail registers a codec-aware LLM request +// sanitizer. Returning omit true removes only the emitted payload. func RegisterLlmSanitizeRequestGuardrail(name string, priority int32, fn LLMRequestFunc) error { id := registerClosure(fn) cName := C.CString(name) defer C.free(unsafe.Pointer(cName)) return checkStatus(C.nemo_relay_register_llm_sanitize_request_guardrail( cName, C.int32_t(priority), - C.NemoRelayLlmRequestCb(C.goLlmRequestTrampoline), + C.NemoRelayLlmSanitizeRequestCb(C.goLlmRequestTrampoline), id, C.NemoRelayFreeFn(C.goFreeTrampoline), )) @@ -1391,17 +1393,15 @@ func DeregisterLlmSanitizeRequestGuardrail(name string) error { return checkStatus(C.nemo_relay_deregister_llm_sanitize_request_guardrail(cName)) } -// RegisterLlmSanitizeResponseGuardrail registers a guardrail that sanitizes -// LLM response data before it is returned to the caller. The callback receives -// the response as plain JSON and must return the (possibly modified) response -// JSON. Guardrails are invoked in priority order (lower values run first). +// RegisterLlmSanitizeResponseGuardrail registers a codec-aware LLM response +// sanitizer. Returning omit true removes only the emitted payload. func RegisterLlmSanitizeResponseGuardrail(name string, priority int32, fn LLMResponseFunc) error { id := registerClosure(fn) cName := C.CString(name) defer C.free(unsafe.Pointer(cName)) return checkStatus(C.nemo_relay_register_llm_sanitize_response_guardrail( cName, C.int32_t(priority), - C.NemoRelayLlmResponseFn(C.goLlmResponseTrampoline), + C.NemoRelayLlmSanitizeResponseCb(C.goLlmResponseTrampoline), id, C.NemoRelayFreeFn(C.goFreeTrampoline), )) @@ -2426,7 +2426,7 @@ func ScopeRegisterLlmSanitizeRequestGuardrail(scopeUUID, name string, priority i defer C.free(unsafe.Pointer(cName)) return checkStatus(C.nemo_relay_scope_register_llm_sanitize_request_guardrail( cScopeUUID, cName, C.int32_t(priority), - C.NemoRelayLlmRequestCb(C.goLlmRequestTrampoline), + C.NemoRelayLlmSanitizeRequestCb(C.goLlmRequestTrampoline), id, C.NemoRelayFreeFn(C.goFreeTrampoline), )) @@ -2452,7 +2452,7 @@ func ScopeRegisterLlmSanitizeResponseGuardrail(scopeUUID, name string, priority defer C.free(unsafe.Pointer(cName)) return checkStatus(C.nemo_relay_scope_register_llm_sanitize_response_guardrail( cScopeUUID, cName, C.int32_t(priority), - C.NemoRelayLlmResponseFn(C.goLlmResponseTrampoline), + C.NemoRelayLlmSanitizeResponseCb(C.goLlmResponseTrampoline), id, C.NemoRelayFreeFn(C.goFreeTrampoline), )) diff --git a/go/nemo_relay/plugin.go b/go/nemo_relay/plugin.go index 7da42b475..8a7f1942e 100644 --- a/go/nemo_relay/plugin.go +++ b/go/nemo_relay/plugin.go @@ -5,10 +5,15 @@ package nemo_relay /* #include +#include #include typedef struct FfiPluginContext FfiPluginContext; typedef struct FfiPluginActivation FfiPluginActivation; +typedef struct FfiLlmSanitizeRequestCodec FfiLlmSanitizeRequestCodec; +typedef struct FfiLlmSanitizeResponseCodec FfiLlmSanitizeResponseCodec; +typedef struct NemoRelayLlmSanitizeRequestContext { uint32_t codec_kind; const char* codec_id; const FfiLlmSanitizeRequestCodec* codec; } NemoRelayLlmSanitizeRequestContext; +typedef struct NemoRelayLlmSanitizeResponseContext { uint32_t codec_kind; const char* codec_id; const FfiLlmSanitizeResponseCodec* codec; } NemoRelayLlmSanitizeResponseContext; typedef void (*NemoRelayFreeFn)(void* user_data); typedef char* (*NemoRelayPluginValidateCb)(void* user_data, const char* plugin_config_json); @@ -17,8 +22,8 @@ typedef void (*NemoRelayEventSubscriberFn)(void* user_data, const void* event); typedef char* (*NemoRelayEventSanitizeFn)(void* user_data, const void* event, const char* fields_json); typedef char* (*NemoRelayToolSanitizeFn)(void* user_data, const char* name, const char* args_json); typedef char* (*NemoRelayToolConditionalFn)(void* user_data, const char* name, const char* args_json); -typedef void* (*NemoRelayLlmRequestCb)(void* user_data, const void* request); -typedef char* (*NemoRelayLlmResponseFn)(void* user_data, const char* response_json); +typedef void* (*NemoRelayLlmSanitizeRequestCb)(void* user_data, const void* request, NemoRelayLlmSanitizeRequestContext context); +typedef char* (*NemoRelayLlmSanitizeResponseCb)(void* user_data, const char* response_json, NemoRelayLlmSanitizeResponseContext context); typedef char* (*NemoRelayLlmConditionalCb)(void* user_data, const void* request); typedef int32_t (*NemoRelayLlmRequestInterceptCb)(void* user_data, const char* name, const void* request, const char* annotated_json, char** out_outcome_json); typedef char* (*NemoRelayLlmExecNextFn)(const char* native_json, void* next_ctx); @@ -45,8 +50,8 @@ extern int32_t nemo_relay_plugin_context_register_scope_sanitize_end_guardrail(F extern int32_t nemo_relay_plugin_context_register_tool_sanitize_request_guardrail(FfiPluginContext* ctx, const char* name, int32_t priority, NemoRelayToolSanitizeFn cb, void* user_data, NemoRelayFreeFn free_fn); extern int32_t nemo_relay_plugin_context_register_tool_sanitize_response_guardrail(FfiPluginContext* ctx, const char* name, int32_t priority, NemoRelayToolSanitizeFn cb, void* user_data, NemoRelayFreeFn free_fn); extern int32_t nemo_relay_plugin_context_register_tool_conditional_execution_guardrail(FfiPluginContext* ctx, const char* name, int32_t priority, NemoRelayToolConditionalFn cb, void* user_data, NemoRelayFreeFn free_fn); -extern int32_t nemo_relay_plugin_context_register_llm_sanitize_request_guardrail(FfiPluginContext* ctx, const char* name, int32_t priority, NemoRelayLlmRequestCb cb, void* user_data, NemoRelayFreeFn free_fn); -extern int32_t nemo_relay_plugin_context_register_llm_sanitize_response_guardrail(FfiPluginContext* ctx, const char* name, int32_t priority, NemoRelayLlmResponseFn cb, void* user_data, NemoRelayFreeFn free_fn); +extern int32_t nemo_relay_plugin_context_register_llm_sanitize_request_guardrail(FfiPluginContext* ctx, const char* name, int32_t priority, NemoRelayLlmSanitizeRequestCb cb, void* user_data, NemoRelayFreeFn free_fn); +extern int32_t nemo_relay_plugin_context_register_llm_sanitize_response_guardrail(FfiPluginContext* ctx, const char* name, int32_t priority, NemoRelayLlmSanitizeResponseCb cb, void* user_data, NemoRelayFreeFn free_fn); extern int32_t nemo_relay_plugin_context_register_llm_conditional_execution_guardrail(FfiPluginContext* ctx, const char* name, int32_t priority, NemoRelayLlmConditionalCb cb, void* user_data, NemoRelayFreeFn free_fn); extern int32_t nemo_relay_plugin_context_register_llm_request_intercept(FfiPluginContext* ctx, const char* name, int32_t priority, _Bool break_chain, NemoRelayLlmRequestInterceptCb cb, void* user_data, NemoRelayFreeFn free_fn); extern int32_t nemo_relay_plugin_context_register_tool_request_intercept(FfiPluginContext* ctx, const char* name, int32_t priority, _Bool break_chain, NemoRelayToolSanitizeFn cb, void* user_data, NemoRelayFreeFn free_fn); @@ -61,8 +66,8 @@ extern char* goEventSanitizeTrampoline(void*, const void*, const char*); extern void goFreeTrampoline(void*); extern char* goToolSanitizeTrampoline(void*, const char*, const char*); extern char* goToolConditionalTrampoline(void*, const char*, const char*); -extern void* goLlmRequestTrampoline(void*, const void*); -extern char* goLlmResponseTrampoline(void*, const char*); +extern void* goLlmRequestTrampoline(void*, const void*, NemoRelayLlmSanitizeRequestContext); +extern char* goLlmResponseTrampoline(void*, const char*, NemoRelayLlmSanitizeResponseContext); extern char* goLlmConditionalTrampoline(void*, const void*); extern char* goLlmExecInterceptTrampoline(void*, const char*, NemoRelayLlmExecNextFn, void*); extern int32_t goLlmRequestInterceptTrampoline(void*, const char*, const void*, const char*, char**); @@ -668,7 +673,7 @@ func (ctx *PluginContext) RegisterLlmSanitizeRequestGuardrail(name string, prior ctx.ptr, cName, C.int32_t(priority), - (C.NemoRelayLlmRequestCb)(C.goLlmRequestTrampoline), + (C.NemoRelayLlmSanitizeRequestCb)(C.goLlmRequestTrampoline), userData, (C.NemoRelayFreeFn)(C.goFreeTrampoline), )) @@ -686,7 +691,7 @@ func (ctx *PluginContext) RegisterLlmSanitizeResponseGuardrail(name string, prio ctx.ptr, cName, C.int32_t(priority), - (C.NemoRelayLlmResponseFn)(C.goLlmResponseTrampoline), + (C.NemoRelayLlmSanitizeResponseCb)(C.goLlmResponseTrampoline), userData, (C.NemoRelayFreeFn)(C.goFreeTrampoline), )) diff --git a/go/nemo_relay/scope_local_test.go b/go/nemo_relay/scope_local_test.go index 3e553d458..565d0552e 100644 --- a/go/nemo_relay/scope_local_test.go +++ b/go/nemo_relay/scope_local_test.go @@ -923,12 +923,13 @@ func TestScopeLocalLlmSanitizeRequestGuardrailAffectsEvent(t *testing.T) { handle, _ := PushScope("llm_scope_guard", ScopeTypeAgent) defer PopScope(handle) err := ScopeRegisterLlmSanitizeRequestGuardrail(handle.UUID(), "scope_llm_san_req", 1, - func(headers, content json.RawMessage) (json.RawMessage, json.RawMessage) { + func(request LLMRequestDTO, _ LLMSanitizeRequestContext) (LLMRequestDTO, bool) { var m map[string]interface{} - json.Unmarshal(content, &m) + json.Unmarshal(request.Content, &m) m["scope_llm_sanitized"] = true out, _ := json.Marshal(m) - return headers, out + request.Content = out + return request, false }) if err != nil { t.Fatalf("ScopeRegisterLlmSanitizeRequestGuardrail failed: %v", err) @@ -1080,9 +1081,9 @@ func assertScopeLocalLLMWrappersDeregister(t *testing.T, scopeUUID string, reque &sanitizeRequestCalls, func() error { return ScopeRegisterLlmSanitizeRequestGuardrail(scopeUUID, "llm_scope_san_req", 1, - func(headers, content json.RawMessage) (json.RawMessage, json.RawMessage) { + func(request LLMRequestDTO, _ LLMSanitizeRequestContext) (LLMRequestDTO, bool) { sanitizeRequestCalls++ - return headers, content + return request, false }, ) }, @@ -1098,9 +1099,9 @@ func assertScopeLocalLLMWrappersDeregister(t *testing.T, scopeUUID string, reque &sanitizeResponseCalls, func() error { return ScopeRegisterLlmSanitizeResponseGuardrail(scopeUUID, "llm_scope_san_resp", 1, - func(responseJSON json.RawMessage) json.RawMessage { + func(responseJSON json.RawMessage, _ LLMSanitizeResponseContext) (json.RawMessage, bool) { sanitizeResponseCalls++ - return responseJSON + return responseJSON, false }, ) }, diff --git a/python/nemo_relay/__init__.py b/python/nemo_relay/__init__.py index b54404246..fb8c89152 100644 --- a/python/nemo_relay/__init__.py +++ b/python/nemo_relay/__init__.py @@ -96,9 +96,14 @@ async def main(): AtofExporterMode, AtofStreamSinkConfig, LLMAttributes, + LlmCodecIdentity, LLMHandle, LLMRequest, LLMRequestInterceptOutcome, + LlmSanitizeRequestCodec, + LlmSanitizeRequestContext, + LlmSanitizeResponseCodec, + LlmSanitizeResponseContext, MarkEvent, OpenInferenceConfig, OpenInferenceSubscriber, @@ -157,12 +162,13 @@ class EventSanitizeFields(TypedDict): #: message. Returning ``None`` allows execution to continue. ToolConditionalExecutionGuardrail: TypeAlias = Callable[[str, Json], Optional[str]] #: Guardrail callback that sanitizes an ``LLMRequest`` used for emitted events. -#: The returned request is recorded for observability and does not replace the -#: caller-visible request value unless the managed LLM API documents otherwise. -LlmSanitizeRequestGuardrail: TypeAlias = Callable[[LLMRequest], LLMRequest] -#: Guardrail callback that sanitizes an emitted JSON LLM response payload. The -#: returned object is recorded on the event; callback exceptions propagate. -LlmSanitizeResponseGuardrail: TypeAlias = Callable[[JsonObject], JsonObject] +#: Callbacks receive ``(request, context)``. Returning ``None`` omits the LLM observability +#: payload and annotation without changing the caller-visible request. +LlmSanitizeRequestGuardrail: TypeAlias = Callable[[LLMRequest, "LlmSanitizeRequestContext"], Optional[LLMRequest]] +#: Guardrail callback that sanitizes an emitted JSON LLM response payload. +#: Callbacks receive ``(response, context)`` and can return ``None`` to omit +#: observability payload and annotation without changing the caller response. +LlmSanitizeResponseGuardrail: TypeAlias = Callable[[Json, "LlmSanitizeResponseContext"], Optional[Json]] #: Guardrail callback that can block an LLM call by returning a rejection #: message. Returning ``None`` allows execution to continue. LlmConditionalExecutionGuardrail: TypeAlias = Callable[[LLMRequest], Optional[str]] @@ -497,6 +503,11 @@ def worker() -> None: "ToolConditionalExecutionGuardrail", "LlmSanitizeRequestGuardrail", "LlmSanitizeResponseGuardrail", + "LlmCodecIdentity", + "LlmSanitizeRequestContext", + "LlmSanitizeResponseContext", + "LlmSanitizeRequestCodec", + "LlmSanitizeResponseCodec", "LlmConditionalExecutionGuardrail", "ToolRequestIntercept", "ToolExecutionIntercept", diff --git a/python/nemo_relay/__init__.pyi b/python/nemo_relay/__init__.pyi index 04cb3429d..7c2372a57 100644 --- a/python/nemo_relay/__init__.pyi +++ b/python/nemo_relay/__init__.pyi @@ -63,6 +63,9 @@ from nemo_relay._native import ( from nemo_relay._native import ( LLMAttributes as LLMAttributes, ) +from nemo_relay._native import ( + LlmCodecIdentity as LlmCodecIdentity, +) from nemo_relay._native import ( LLMHandle as LLMHandle, ) @@ -72,6 +75,18 @@ from nemo_relay._native import ( from nemo_relay._native import ( LLMRequestInterceptOutcome as LLMRequestInterceptOutcome, ) +from nemo_relay._native import ( + LlmSanitizeRequestCodec as LlmSanitizeRequestCodec, +) +from nemo_relay._native import ( + LlmSanitizeRequestContext as LlmSanitizeRequestContext, +) +from nemo_relay._native import ( + LlmSanitizeResponseCodec as LlmSanitizeResponseCodec, +) +from nemo_relay._native import ( + LlmSanitizeResponseContext as LlmSanitizeResponseContext, +) from nemo_relay._native import ( MarkEvent as MarkEvent, ) @@ -172,23 +187,31 @@ Arguments: Return: ``None`` to allow execution, or a rejection message to block it. """ -LlmSanitizeRequestGuardrail: TypeAlias = Callable[[LLMRequest], LLMRequest] +LlmSanitizeRequestGuardrail: TypeAlias = Callable[[LLMRequest, "LlmSanitizeRequestContext"], Optional[LLMRequest]] """Guardrail callback that sanitizes an ``LLMRequest`` used for emitted events. Arguments: - The current LLM request. + The current LLM request and a context object containing a tagged ``codec`` + identity. Its ``kind`` is ``none``, ``builtin``, ``runtime``, or ``opaque``; + ``builtin`` and ``runtime`` identities include ``id``. Use + ``context.resolve_codec()`` to access the active in-process codec. Return: - Request object recorded on the emitted lifecycle event. + Request object recorded on the emitted lifecycle event, or ``None`` to omit + the LLM observability payload and annotation. """ -LlmSanitizeResponseGuardrail: TypeAlias = Callable[[JsonObject], JsonObject] +LlmSanitizeResponseGuardrail: TypeAlias = Callable[[Json, "LlmSanitizeResponseContext"], Optional[Json]] """Guardrail callback that sanitizes an emitted JSON LLM response payload. Arguments: - The response object to sanitize for observability. + The response object and a context object containing a tagged ``codec`` + identity. Its ``kind`` is ``none``, ``builtin``, ``runtime``, or ``opaque``; + ``builtin`` and ``runtime`` identities include ``id``. Use + ``context.resolve_codec()`` to access the active in-process codec. Return: - Response object recorded on the emitted lifecycle event. + Response object recorded on the emitted lifecycle event, or ``None`` to + omit the LLM observability payload and annotation. """ LlmConditionalExecutionGuardrail: TypeAlias = Callable[[LLMRequest], Optional[str]] """Guardrail callback that can block an LLM call. diff --git a/python/nemo_relay/_native.pyi b/python/nemo_relay/_native.pyi index aaec3027a..a76ee5240 100644 --- a/python/nemo_relay/_native.pyi +++ b/python/nemo_relay/_native.pyi @@ -39,8 +39,8 @@ class _EventSanitizeFields(TypedDict): _ToolSanitizeGuardrail: TypeAlias = Callable[[str, _Json], _Json] _ToolConditionalExecutionGuardrail: TypeAlias = Callable[[str, _Json], Optional[str]] -_LlmSanitizeRequestGuardrail: TypeAlias = Callable[["LLMRequest"], "LLMRequest"] -_LlmSanitizeResponseGuardrail: TypeAlias = Callable[[_JsonObject], _JsonObject] +_LlmSanitizeRequestGuardrail: TypeAlias = Callable[["LLMRequest", "LlmSanitizeRequestContext"], Optional["LLMRequest"]] +_LlmSanitizeResponseGuardrail: TypeAlias = Callable[[_Json, "LlmSanitizeResponseContext"], Optional[_Json]] _EventSanitizeGuardrail: TypeAlias = Callable[[ScopeEvent | MarkEvent, _EventSanitizeFields], _EventSanitizeFields] _LlmConditionalExecutionGuardrail: TypeAlias = Callable[["LLMRequest"], Optional[str]] _ToolRequestIntercept: TypeAlias = Callable[[str, _Json], _Json] @@ -61,6 +61,35 @@ _LlmStreamExecutionIntercept: TypeAlias = Callable[ AsyncIterator[_Json] | Awaitable[AsyncIterator[_Json]], ] +class LlmCodecIdentity: + """Structured identity of the active managed LLM codec.""" + + @property + def kind(self) -> Literal["none", "builtin", "runtime", "opaque"]: ... + @property + def id(self) -> str | None: ... + +class LlmSanitizeRequestContext: + """Per-call context passed to an LLM request sanitizer callback.""" + + @property + def codec(self) -> LlmCodecIdentity: ... + def resolve_codec(self) -> LlmSanitizeRequestCodec | None: ... + +class LlmSanitizeResponseContext: + """Per-call context passed to an LLM response sanitizer callback.""" + + @property + def codec(self) -> LlmCodecIdentity: ... + def resolve_codec(self) -> LlmSanitizeResponseCodec | None: ... + +class LlmSanitizeRequestCodec: + def decode(self, request: LLMRequest) -> AnnotatedLLMRequest: ... + def encode(self, annotated: AnnotatedLLMRequest, original: LLMRequest) -> LLMRequest: ... + +class LlmSanitizeResponseCodec: + def decode_response(self, response: _Json) -> AnnotatedLLMResponse: ... + class ScopeAttributes: """Bitflags describing scope properties. diff --git a/python/nemo_relay/guardrails.py b/python/nemo_relay/guardrails.py index 92bc8d9b5..94cb2fd63 100644 --- a/python/nemo_relay/guardrails.py +++ b/python/nemo_relay/guardrails.py @@ -248,8 +248,14 @@ def register_llm_sanitize_request(name: str, priority: int, guardrail: LlmSaniti Args: name: Unique guardrail name used for later replacement or removal. priority: Execution order for the guardrail. Lower values run first. - guardrail: Callable invoked as ``guardrail(request)`` that must return - the sanitized request recorded on the emitted start event. + guardrail: Callable invoked as ``guardrail(request, context)`` that + returns the sanitized request, or ``None`` to omit the LLM + observability payload and its annotation. ``context`` contains + a ``LlmSanitizeRequestContext`` whose ``codec`` is a structured identity + with ``kind`` of ``none``, ``builtin``, ``runtime``, or ``opaque``. + ``builtin`` and ``runtime`` identities include ``id``. Use + ``context.resolve_codec()`` to access the active codec for normalized + processing. Returns: None: This function returns after the guardrail is registered. @@ -263,7 +269,7 @@ def register_llm_sanitize_request(name: str, priority: int, guardrail: LlmSaniti import nemo_relay - def strip_auth(request): + def strip_auth(request, context): headers = {k: v for k, v in request.headers.items() if k.lower() != "authorization"} return nemo_relay.LLMRequest(headers, request.content) @@ -295,8 +301,14 @@ def register_llm_sanitize_response(name: str, priority: int, guardrail: LlmSanit Args: name: Unique guardrail name used for later replacement or removal. priority: Execution order for the guardrail. Lower values run first. - guardrail: Callable invoked as ``guardrail(response)`` that must return - the sanitized payload recorded on the emitted end event. + guardrail: Callable invoked as ``guardrail(response, context)`` that + returns the sanitized payload, or ``None`` to omit the LLM + observability payload and its annotation. ``context`` contains + a ``LlmSanitizeResponseContext`` whose ``codec`` is a structured identity + with ``kind`` of ``none``, ``builtin``, ``runtime``, or ``opaque``. + ``builtin`` and ``runtime`` identities include ``id``. Use + ``context.resolve_codec()`` to access the active codec for normalized + processing. Returns: None: This function returns after the guardrail is registered. diff --git a/python/plugin/src/nemo_relay_plugin/__init__.py b/python/plugin/src/nemo_relay_plugin/__init__.py index f3052e42a..54466a5b5 100644 --- a/python/plugin/src/nemo_relay_plugin/__init__.py +++ b/python/plugin/src/nemo_relay_plugin/__init__.py @@ -20,6 +20,11 @@ Event: A Relay event represented as a JSON object. EventSanitizeFields: Mutable event observability fields. LlmRequest: A Relay LLM request represented as a JSON object. + LlmCodecIdentity: Typed discriminator for the active LLM codec. + LlmSanitizeRequestContext: Per-call context supplied to an LLM request sanitizer. + LlmSanitizeResponseContext: Per-call context supplied to an LLM response sanitizer. + WorkerRequestCodec: Invocation-scoped async proxy for an active request codec. + WorkerResponseCodec: Invocation-scoped async proxy for an active response codec. AnnotatedLlmRequest: An annotated Relay LLM request represented as a JSON object. PendingMarkSpec: A mark Relay emits under its managed lifecycle scope. @@ -71,6 +76,7 @@ EventSanitizeCallback, EventSanitizeFields, Json, + LlmCodecIdentity, LlmConditionalCallback, LlmExecutionCallback, LlmNext, @@ -85,7 +91,9 @@ LlmRequestCallback, LlmRequestInterceptOutcome, LlmSanitizeRequestCallback, + LlmSanitizeRequestContext, LlmSanitizeResponseCallback, + LlmSanitizeResponseContext, LlmStreamExecutionCallback, LlmStreamNext, PendingMarkSpec, @@ -100,6 +108,8 @@ ToolRequestCallback, ToolSanitizeCallback, WorkerPlugin, + WorkerRequestCodec, + WorkerResponseCodec, WorkerSdkError, serve_plugin, ) @@ -113,6 +123,7 @@ "EventSanitizeFields", "Json", "LlmConditionalCallback", + "LlmCodecIdentity", "LlmExecutionCallback", "LlmOptimizationContribution", "LlmOptimizationDataSchema", @@ -123,6 +134,8 @@ "LlmOptimizationTokens", "LlmNext", "LlmRequest", + "LlmSanitizeRequestContext", + "LlmSanitizeResponseContext", "LlmRequestCallback", "LlmRequestInterceptOutcome", "LlmSanitizeRequestCallback", @@ -140,6 +153,8 @@ "ToolNext", "ToolRequestCallback", "ToolSanitizeCallback", + "WorkerRequestCodec", + "WorkerResponseCodec", "WorkerPlugin", "WorkerSdkError", "serve_plugin", diff --git a/python/plugin/src/nemo_relay_plugin/_api.py b/python/plugin/src/nemo_relay_plugin/_api.py index b8facbd76..7f8b6ccb8 100644 --- a/python/plugin/src/nemo_relay_plugin/_api.py +++ b/python/plugin/src/nemo_relay_plugin/_api.py @@ -102,6 +102,105 @@ class EventSanitizeFields(TypedDict): #: An annotated Relay LLM request represented as a JSON object. AnnotatedLlmRequest: TypeAlias = dict[str, Any] + +@dataclass(frozen=True) +class LlmCodecIdentity: + """Structured identity of the codec active for one worker invocation.""" + + kind: str + id: str | None = None + + +@dataclass(frozen=True) +class LlmSanitizeRequestContext: + """Structured per-call context provided to an LLM request sanitizer.""" + + codec: LlmCodecIdentity + _runtime: "PluginRuntime | None" = field(default=None, repr=False, compare=False) + _capability_id: str | None = field(default=None, repr=False, compare=False) + _invocation_id: str | None = field(default=None, repr=False, compare=False) + + def resolve_codec(self) -> "WorkerRequestCodec | None": + """Return the active request-codec proxy for this invocation.""" + if self._runtime is None or self._capability_id is None: + return None + if self._invocation_id is None: + return None + return WorkerRequestCodec(self._runtime, self._capability_id, self._invocation_id) + + +@dataclass(frozen=True) +class LlmSanitizeResponseContext: + """Structured per-call context provided to an LLM response sanitizer.""" + + codec: LlmCodecIdentity + _runtime: "PluginRuntime | None" = field(default=None, repr=False, compare=False) + _capability_id: str | None = field(default=None, repr=False, compare=False) + _invocation_id: str | None = field(default=None, repr=False, compare=False) + + def resolve_codec(self) -> "WorkerResponseCodec | None": + """Return the active response-codec proxy for this invocation.""" + if self._runtime is None or self._capability_id is None: + return None + if self._invocation_id is None: + return None + return WorkerResponseCodec(self._runtime, self._capability_id, self._invocation_id) + + +@dataclass(frozen=True) +class WorkerRequestCodec: + """Invocation-scoped async proxy for an active request codec.""" + + _runtime: "PluginRuntime" + _capability_id: str + _invocation_id: str + + async def decode(self, request: LlmRequest) -> AnnotatedLlmRequest: + """Decode a wire request with the active host codec.""" + return await self._runtime._decode_llm_codec_request(self._capability_id, self._invocation_id, request) + + async def encode(self, annotated: AnnotatedLlmRequest, original: LlmRequest) -> LlmRequest: + """Encode an annotated request with the active host codec.""" + return await self._runtime._encode_llm_codec_request( + self._capability_id, self._invocation_id, annotated, original + ) + + +@dataclass(frozen=True) +class WorkerResponseCodec: + """Invocation-scoped async proxy for an active response codec.""" + + _runtime: "PluginRuntime" + _capability_id: str + _invocation_id: str + + async def decode(self, response: Json) -> Json: + """Decode a wire response with the active host codec.""" + return await self._runtime._decode_llm_codec_response(self._capability_id, self._invocation_id, response) + + +def _llm_codec_identity(invocation: pb.LlmInvocation) -> LlmCodecIdentity: + """Return the codec identity from a worker invocation.""" + context = getattr(invocation, invocation.WhichOneof("sanitize_context") or "", None) + proto_codec = context.codec if context is not None and context.HasField("codec") else None + codec_id = proto_codec.id if proto_codec is not None and proto_codec.HasField("id") else None + codec_kind = proto_codec.kind if proto_codec is not None else pb.LLM_CODEC_KIND_UNSPECIFIED + if codec_kind == pb.LLM_CODEC_KIND_UNSPECIFIED: + identity = LlmCodecIdentity("none") + elif codec_kind == pb.LLM_CODEC_KIND_BUILTIN and codec_id: + identity = LlmCodecIdentity("builtin", codec_id) + elif codec_kind == pb.LLM_CODEC_KIND_RUNTIME and codec_id: + identity = LlmCodecIdentity("runtime", codec_id) + else: + identity = LlmCodecIdentity("opaque") + return identity + + +def _llm_codec_capability(invocation: pb.LlmInvocation) -> str | None: + context = getattr(invocation, invocation.WhichOneof("sanitize_context") or "", None) + return context.codec_capability_id if context is not None and context.HasField("codec_capability_id") else None + + WORKER_PROTOCOL = "grpc-v1" JSON_SCHEMA = "nemo.relay.Json@1" EVENT_SCHEMA = "nemo.relay.Event@1" @@ -730,8 +829,12 @@ def register(self, ctx: PluginContext, config: Json) -> None | Awaitable[None]: [str, Json, "ToolNext"], ToolExecutionInterceptOutcome | Awaitable[ToolExecutionInterceptOutcome], ] -LlmSanitizeRequestCallback: TypeAlias = Callable[[LlmRequest], LlmRequest | Awaitable[LlmRequest]] -LlmSanitizeResponseCallback: TypeAlias = Callable[[Json], Json | Awaitable[Json]] +LlmSanitizeRequestCallback: TypeAlias = Callable[ + [LlmRequest, LlmSanitizeRequestContext], LlmRequest | None | Awaitable[LlmRequest | None] +] +LlmSanitizeResponseCallback: TypeAlias = Callable[ + [Json, LlmSanitizeResponseContext], Json | None | Awaitable[Json | None] +] LlmConditionalCallback: TypeAlias = Callable[[LlmRequest], str | None | Awaitable[str | None]] LlmRequestCallback: TypeAlias = Callable[ [str, LlmRequest, AnnotatedLlmRequest | None], @@ -808,6 +911,15 @@ class PluginContext: pb.SCOPE_SANITIZE_START_GUARDRAIL: "scope_start_sanitizers", pb.SCOPE_SANITIZE_END_GUARDRAIL: "scope_end_sanitizers", } + _LLM_HANDLER_ATTRIBUTES: ClassVar[frozenset[int]] = frozenset( + { + pb.LLM_SANITIZE_REQUEST_GUARDRAIL, + pb.LLM_SANITIZE_RESPONSE_GUARDRAIL, + pb.LLM_CONDITIONAL_EXECUTION_GUARDRAIL, + pb.LLM_REQUEST_INTERCEPT, + pb.LLM_EXECUTION_INTERCEPT, + } + ) def __init__(self, runtime: PluginRuntime | None = None) -> None: self._runtime = runtime @@ -1030,10 +1142,9 @@ def register_llm_sanitize_request_guardrail( Args: name: Component-local registration name. - callback: Function receiving an :data:`LlmRequest` and returning + callback: Function receiving ``(request, context)`` and returning the request recorded on the LLM start event, directly or - through an awaitable. It does not change the request sent to - the model. + through an awaitable. Return ``None`` to omit observability. priority: Execution order. Lower values run first. """ self._push_registration(name, pb.LLM_SANITIZE_REQUEST_GUARDRAIL, priority, False) @@ -1050,9 +1161,9 @@ def register_llm_sanitize_response_guardrail( Args: name: Component-local registration name. - callback: Function receiving response JSON and returning the value - recorded on the LLM end event, directly or through an - awaitable. It does not change the real model response. + callback: Function receiving ``(response, context)`` and returning + the value recorded on the LLM end event, directly or through an + awaitable. Return ``None`` to omit observability. priority: Execution order. Lower values run first. """ self._push_registration(name, pb.LLM_SANITIZE_RESPONSE_GUARDRAIL, priority, False) @@ -1190,6 +1301,47 @@ def __init__(self, *, activation_id: str, auth_token: str, host_stub: Any) -> No self._auth_token = auth_token self._host_stub = host_stub + async def _decode_llm_codec_request( + self, capability_id: str, invocation_id: str, request: LlmRequest + ) -> AnnotatedLlmRequest: + result = await self._host_stub.DecodeLlmCodecRequest( + pb.LlmCodecDecodeRequest( + activation_id=self._activation_id, + auth_token=self._auth_token, + codec_capability_id=capability_id, + invocation_id=invocation_id, + request=_json_envelope(LLM_REQUEST_SCHEMA, request), + ) + ) + return _json_result_to_value(result, ANNOTATED_LLM_REQUEST_SCHEMA) + + async def _encode_llm_codec_request( + self, capability_id: str, invocation_id: str, annotated: AnnotatedLlmRequest, original: LlmRequest + ) -> LlmRequest: + result = await self._host_stub.EncodeLlmCodecRequest( + pb.LlmCodecEncodeRequest( + activation_id=self._activation_id, + auth_token=self._auth_token, + codec_capability_id=capability_id, + invocation_id=invocation_id, + annotated_request=_json_envelope(ANNOTATED_LLM_REQUEST_SCHEMA, annotated), + original_request=_json_envelope(LLM_REQUEST_SCHEMA, original), + ) + ) + return _json_result_to_value(result, LLM_REQUEST_SCHEMA) + + async def _decode_llm_codec_response(self, capability_id: str, invocation_id: str, response: Json) -> Json: + result = await self._host_stub.DecodeLlmCodecResponse( + pb.LlmCodecDecodeResponse( + activation_id=self._activation_id, + auth_token=self._auth_token, + codec_capability_id=capability_id, + invocation_id=invocation_id, + response=_json_envelope(JSON_SCHEMA, response), + ) + ) + return _json_result_to_value(result) + async def emit_mark( self, name: str, @@ -1837,6 +1989,8 @@ async def Shutdown(self, request: Any, context: Any) -> Any: async def _invoke_result(self, request: Any) -> Any: with _bind_invocation_scope(request): + if request.surface in PluginContext._LLM_HANDLER_ATTRIBUTES: + return await self._invoke_llm_result(request) if request.surface == pb.SUBSCRIBER: event = _decode_required_envelope(request.event, "event", EVENT_SCHEMA) await _maybe_await(self._handler(self._handlers.subscribers, request.registration_name)(event)) @@ -1906,69 +2060,80 @@ async def _invoke_result(self, request: Any) -> Any: ), ) ) - if request.surface == pb.LLM_SANITIZE_REQUEST_GUARDRAIL: - return _json_response( - await _maybe_await( - self._handler(self._handlers.llm_sanitize_requests, request.registration_name)( - _decode_required_envelope(request.llm.request, "llm request", LLM_REQUEST_SCHEMA) - ) - ) - ) - if request.surface == pb.LLM_SANITIZE_RESPONSE_GUARDRAIL: - return _json_response( - await _maybe_await( - self._handler(self._handlers.llm_sanitize_responses, request.registration_name)( - _decode_required_envelope(request.llm.response, "llm response") - ) - ) - ) - if request.surface == pb.LLM_CONDITIONAL_EXECUTION_GUARDRAIL: - result = await _maybe_await( - self._handler(self._handlers.llm_conditionals, request.registration_name)( - _decode_required_envelope(request.llm.request, "llm request", LLM_REQUEST_SCHEMA) - ) + raise WorkerSdkError(f"unsupported registration surface {request.surface}") + + async def _invoke_llm_result(self, request: Any) -> Any: + if request.surface == pb.LLM_SANITIZE_REQUEST_GUARDRAIL: + callback = self._handler(self._handlers.llm_sanitize_requests, request.registration_name) + payload = _decode_required_envelope(request.llm.request, "llm request", LLM_REQUEST_SCHEMA) + return await self._invoke_llm_sanitizer(callback, payload, request, LlmSanitizeRequestContext) + if request.surface == pb.LLM_SANITIZE_RESPONSE_GUARDRAIL: + callback = self._handler(self._handlers.llm_sanitize_responses, request.registration_name) + payload = _decode_required_envelope(request.llm.response, "llm response") + return await self._invoke_llm_sanitizer(callback, payload, request, LlmSanitizeResponseContext) + if request.surface == pb.LLM_CONDITIONAL_EXECUTION_GUARDRAIL: + result = await _maybe_await( + self._handler(self._handlers.llm_conditionals, request.registration_name)( + _decode_required_envelope(request.llm.request, "llm request", LLM_REQUEST_SCHEMA) ) - return pb.InvokeResponse(guardrail=pb.GuardrailResult(block_reason=result or "")) - if request.surface == pb.LLM_REQUEST_INTERCEPT: - payload = request.llm - llm_request = _decode_required_envelope(payload.request, "llm request", LLM_REQUEST_SCHEMA) - annotated = ( - _decode_required_envelope( - payload.annotated_request, - "annotated llm request", - ANNOTATED_LLM_REQUEST_SCHEMA, - ) - if payload.HasField("annotated_request") - else None + ) + return pb.InvokeResponse(guardrail=pb.GuardrailResult(block_reason=result or "")) + if request.surface == pb.LLM_REQUEST_INTERCEPT: + payload = request.llm + llm_request = _decode_required_envelope(payload.request, "llm request", LLM_REQUEST_SCHEMA) + annotated = ( + _decode_required_envelope( + payload.annotated_request, + "annotated llm request", + ANNOTATED_LLM_REQUEST_SCHEMA, ) - result = await _maybe_await( - self._handler(self._handlers.llm_requests, request.registration_name)( - payload.model_name, - llm_request, - annotated, - ) + if payload.HasField("annotated_request") + else None + ) + result = await _maybe_await( + self._handler(self._handlers.llm_requests, request.registration_name)( + payload.model_name, + llm_request, + annotated, ) - if not isinstance(result, LlmRequestInterceptOutcome): - raise WorkerSdkError("LLM request intercept must return LlmRequestInterceptOutcome") - return pb.InvokeResponse( - llm_request=pb.LlmRequestInterceptResult( - outcome=_json_envelope( - LLM_REQUEST_INTERCEPT_OUTCOME_SCHEMA, - result.to_json(), - ), - ) + ) + if not isinstance(result, LlmRequestInterceptOutcome): + raise WorkerSdkError("LLM request intercept must return LlmRequestInterceptOutcome") + return pb.InvokeResponse( + llm_request=pb.LlmRequestInterceptResult( + outcome=_json_envelope( + LLM_REQUEST_INTERCEPT_OUTCOME_SCHEMA, + result.to_json(), + ), ) - if request.surface == pb.LLM_EXECUTION_INTERCEPT: - payload = request.llm - result = await _maybe_await( - self._handler(self._handlers.llm_executions, request.registration_name)( - payload.model_name, - _decode_required_envelope(payload.request, "llm request", LLM_REQUEST_SCHEMA), - LlmNext(self._runtime, request.continuation_id), - ) + ) + if request.surface == pb.LLM_EXECUTION_INTERCEPT: + payload = request.llm + result = await _maybe_await( + self._handler(self._handlers.llm_executions, request.registration_name)( + payload.model_name, + _decode_required_envelope(payload.request, "llm request", LLM_REQUEST_SCHEMA), + LlmNext(self._runtime, request.continuation_id), ) - return _json_response(result) - raise WorkerSdkError(f"unsupported registration surface {request.surface}") + ) + return _json_response(result) + raise WorkerSdkError(f"unsupported LLM registration surface {request.surface}") + + async def _invoke_llm_sanitizer( + self, + callback: Any, + payload: Any, + request: Any, + context_type: type[LlmSanitizeRequestContext] | type[LlmSanitizeResponseContext], + ) -> Any: + context = context_type( + _llm_codec_identity(request.llm), + self._runtime, + _llm_codec_capability(request.llm), + request.invocation_id, + ) + result = await _maybe_await(callback(payload, context)) + return pb.InvokeResponse(empty=pb.EmptyResult()) if result is None else _json_response(result) def _start_invocation( self, @@ -2116,10 +2281,10 @@ def _json_response(value: Json) -> Any: return pb.InvokeResponse(json=pb.JsonResult(value=_json_envelope(JSON_SCHEMA, value))) -def _json_result_to_value(result: Any) -> Json: +def _json_result_to_value(result: Any, expected_schema: str = JSON_SCHEMA) -> Json: if result.HasField("error"): raise _worker_error_to_sdk(result.error) - return _decode_required_envelope(result.value, "json result") + return _decode_required_envelope(result.value, "json result", expected_schema) def _stream_chunk_to_value(chunk: Any) -> Json: diff --git a/python/tests/plugin/test_worker_sdk.py b/python/tests/plugin/test_worker_sdk.py index 5a521fb1e..7ab918d20 100644 --- a/python/tests/plugin/test_worker_sdk.py +++ b/python/tests/plugin/test_worker_sdk.py @@ -29,6 +29,8 @@ Json, LlmOptimizationContribution, LlmRequestInterceptOutcome, + LlmSanitizeRequestContext, + LlmSanitizeResponseContext, PendingMarkSpec, PluginContext, PluginRuntime, @@ -246,6 +248,37 @@ async def stream() -> AsyncIterator[Any]: return stream() + async def DecodeLlmCodecRequest(self, request: Any) -> Any: + self.requests.append(request) + if self.failures.get("DecodeLlmCodecRequest") == "error": + return pb.JsonResult(error=_worker_error("DecodeLlmCodecRequest failed")) + value = _envelope_value(request.request) + model = value.get("content", {}).get("model") + return pb.JsonResult( + value=_json_envelope( + ANNOTATED_LLM_REQUEST_SCHEMA, + {"messages": [], "model": model}, + ) + ) + + async def EncodeLlmCodecRequest(self, request: Any) -> Any: + self.requests.append(request) + if self.failures.get("EncodeLlmCodecRequest") == "error": + return pb.JsonResult(error=_worker_error("EncodeLlmCodecRequest failed")) + return pb.JsonResult(value=request.original_request) + + async def DecodeLlmCodecResponse(self, request: Any) -> Any: + self.requests.append(request) + if self.failures.get("DecodeLlmCodecResponse") == "error": + return pb.JsonResult(error=_worker_error("DecodeLlmCodecResponse failed")) + value = _envelope_value(request.response) + return pb.JsonResult( + value=_json_envelope( + JSON_SCHEMA, + {"model": value.get("model"), "message": value.get("message")}, + ) + ) + def _host_ack(self, method: str) -> Any: failure = self.failures.get(method) if failure == "empty": @@ -301,10 +334,12 @@ async def tool_execution(name: str, value: Json, next_call: ToolNext) -> ToolExe pending_marks=[PendingMarkSpec("worker.tool.execution")], ) - def llm_sanitize_request(request: Json) -> Json: + def llm_sanitize_request(request: Json, context: LlmSanitizeRequestContext) -> Json: + del context return _tag_llm_request(request, "llm_sanitize_request") - async def llm_sanitize_response(response: Json) -> Json: + async def llm_sanitize_response(response: Json, context: LlmSanitizeResponseContext) -> Json: + del context return _tag(response, "llm_sanitize_response") def llm_block(request: Json) -> str | None: @@ -699,6 +734,266 @@ def callback(tool_name: str, value: Json) -> Json: assert ("shared", pb.TOOL_SANITIZE_RESPONSE_GUARDRAIL) in registrations +def test_plugin_context_registers_llm_sanitizers_under_standard_names(): + context = PluginContext() + + context.register_llm_sanitize_request_guardrail( + "request", + lambda request, codec_context: request if codec_context.codec.kind != "none" else None, + ) + context.register_llm_sanitize_response_guardrail( + "response", + lambda response, codec_context: response if codec_context.codec.kind == "builtin" else None, + ) + + assert [(registration.local_name, registration.surface) for registration in context._handlers.registrations] == [ + ("request", pb.LLM_SANITIZE_REQUEST_GUARDRAIL), + ("response", pb.LLM_SANITIZE_RESPONSE_GUARDRAIL), + ] + + +async def test_llm_sanitizers_receive_codec_context_and_can_omit_payloads(): + seen: list[tuple[str, LlmSanitizeRequestContext | LlmSanitizeResponseContext]] = [] + + class ContextualSanitizerPlugin(WorkerPlugin): + plugin_id = "tests.llm_sanitizer" + + def register(self, ctx: PluginContext, config: Json) -> None: + del config + + def request_sanitizer(request: Json, codec_context: LlmSanitizeRequestContext) -> None: + del request + seen.append(("request", codec_context)) + + def response_sanitizer(response: Json, codec_context: LlmSanitizeResponseContext) -> None: + del response + seen.append(("response", codec_context)) + + ctx.register_llm_sanitize_request_guardrail("request", request_sanitizer) + ctx.register_llm_sanitize_response_guardrail("response", response_sanitizer) + + service = _service(ContextualSanitizerPlugin(), RecordingHostStub()) + await _register(service) + for codec_kind, codec_id in [ + (pb.LLM_CODEC_KIND_UNSPECIFIED, None), + (pb.LLM_CODEC_KIND_BUILTIN, "openai_chat"), + (pb.LLM_CODEC_KIND_BUILTIN, "openai_responses"), + (pb.LLM_CODEC_KIND_BUILTIN, "anthropic_messages"), + (pb.LLM_CODEC_KIND_RUNTIME, "com.example.chat.v1"), + (pb.LLM_CODEC_KIND_OPAQUE, None), + ]: + codec = pb.LlmCodecIdentity(kind=codec_kind) + if codec_id is not None: + codec.id = codec_id + request_payload = _llm_payload( + request={"content": {"prompt": "secret"}}, + response={"secret": "value"}, + ) + request_payload.request_sanitize_context.CopyFrom(pb.LlmSanitizeRequestContext(codec=codec)) + response_payload = _llm_payload( + request={"content": {"prompt": "secret"}}, + response={"secret": "value"}, + ) + response_payload.response_sanitize_context.CopyFrom(pb.LlmSanitizeResponseContext(codec=codec)) + + request_response = await service.Invoke( + _invoke_request( + "request", + pb.LLM_SANITIZE_REQUEST_GUARDRAIL, + llm=request_payload, + ), + AbortContext(), + ) + response_response = await service.Invoke( + _invoke_request( + "response", + pb.LLM_SANITIZE_RESPONSE_GUARDRAIL, + llm=response_payload, + ), + AbortContext(), + ) + assert request_response.WhichOneof("result") == "empty" + assert response_response.WhichOneof("result") == "empty" + + assert seen == [ + ("request", LlmSanitizeRequestContext(plugin_api.LlmCodecIdentity("none"))), + ("response", LlmSanitizeResponseContext(plugin_api.LlmCodecIdentity("none"))), + ("request", LlmSanitizeRequestContext(plugin_api.LlmCodecIdentity("builtin", "openai_chat"))), + ("response", LlmSanitizeResponseContext(plugin_api.LlmCodecIdentity("builtin", "openai_chat"))), + ("request", LlmSanitizeRequestContext(plugin_api.LlmCodecIdentity("builtin", "openai_responses"))), + ("response", LlmSanitizeResponseContext(plugin_api.LlmCodecIdentity("builtin", "openai_responses"))), + ( + "request", + LlmSanitizeRequestContext(plugin_api.LlmCodecIdentity("builtin", "anthropic_messages")), + ), + ( + "response", + LlmSanitizeResponseContext(plugin_api.LlmCodecIdentity("builtin", "anthropic_messages")), + ), + ("request", LlmSanitizeRequestContext(plugin_api.LlmCodecIdentity("runtime", "com.example.chat.v1"))), + ("response", LlmSanitizeResponseContext(plugin_api.LlmCodecIdentity("runtime", "com.example.chat.v1"))), + ("request", LlmSanitizeRequestContext(plugin_api.LlmCodecIdentity("opaque"))), + ("response", LlmSanitizeResponseContext(plugin_api.LlmCodecIdentity("opaque"))), + ] + + +async def test_llm_sanitizers_resolve_directional_codec_proxies(): + calls: list[tuple[str, Json]] = [] + + class CodecSanitizerPlugin(WorkerPlugin): + plugin_id = "tests.llm_codec_proxy" + + def register(self, ctx: PluginContext, config: Json) -> None: + del config + + async def request_sanitizer( + request: Json, + context: LlmSanitizeRequestContext, + ) -> Json: + codec = context.resolve_codec() + assert codec is not None + annotated = await codec.decode(request) + calls.append(("request_decode", annotated)) + return await codec.encode(annotated, request) + + async def response_sanitizer( + response: Json, + context: LlmSanitizeResponseContext, + ) -> Json: + codec = context.resolve_codec() + assert codec is not None + calls.append(("response_decode", await codec.decode(response))) + return response + + ctx.register_llm_sanitize_request_guardrail("request", request_sanitizer) + ctx.register_llm_sanitize_response_guardrail("response", response_sanitizer) + + host = RecordingHostStub() + service = _service(CodecSanitizerPlugin(), host) + await _register(service) + + request_payload = _llm_payload( + request={"headers": {}, "content": {"model": "runtime-model"}}, + ) + request_payload.request_sanitize_context.CopyFrom( + pb.LlmSanitizeRequestContext( + codec=pb.LlmCodecIdentity( + kind=pb.LLM_CODEC_KIND_RUNTIME, + id="com.example.runtime", + ), + codec_capability_id="request-capability", + ) + ) + request_result = await service.Invoke( + _invoke_request( + "request", + pb.LLM_SANITIZE_REQUEST_GUARDRAIL, + llm=request_payload, + ), + AbortContext(), + ) + assert request_result.WhichOneof("result") == "json", request_result + + response_payload = _llm_payload(response={"model": "opaque-model", "message": "secret"}) + response_payload.response_sanitize_context.CopyFrom( + pb.LlmSanitizeResponseContext( + codec=pb.LlmCodecIdentity(kind=pb.LLM_CODEC_KIND_OPAQUE), + codec_capability_id="response-capability", + ) + ) + response_result = await service.Invoke( + _invoke_request( + "response", + pb.LLM_SANITIZE_RESPONSE_GUARDRAIL, + llm=response_payload, + ), + AbortContext(), + ) + assert response_result.WhichOneof("result") == "json", response_result + assert calls == [ + ("request_decode", {"messages": [], "model": "runtime-model"}), + ( + "response_decode", + {"model": "opaque-model", "message": "secret"}, + ), + ] + assert [request.codec_capability_id for request in host.requests if hasattr(request, "codec_capability_id")] == [ + "request-capability", + "request-capability", + "response-capability", + ] + + +@pytest.mark.parametrize( + ("failure", "surface", "registration_name"), + [ + ("DecodeLlmCodecRequest", pb.LLM_SANITIZE_REQUEST_GUARDRAIL, "request"), + ("EncodeLlmCodecRequest", pb.LLM_SANITIZE_REQUEST_GUARDRAIL, "request"), + ("DecodeLlmCodecResponse", pb.LLM_SANITIZE_RESPONSE_GUARDRAIL, "response"), + ], +) +async def test_llm_sanitizer_codec_rpc_failures_return_invocation_errors( + failure: str, + surface: int, + registration_name: str, +): + class CodecFailurePlugin(WorkerPlugin): + plugin_id = "tests.llm_codec_failure" + + def register(self, ctx: PluginContext, config: Json) -> None: + del config + + async def request_sanitizer( + request: Json, + context: LlmSanitizeRequestContext, + ) -> Json: + codec = context.resolve_codec() + assert codec is not None + annotated = await codec.decode(request) + return await codec.encode(annotated, request) + + async def response_sanitizer( + response: Json, + context: LlmSanitizeResponseContext, + ) -> Json: + codec = context.resolve_codec() + assert codec is not None + await codec.decode(response) + return response + + ctx.register_llm_sanitize_request_guardrail("request", request_sanitizer) + ctx.register_llm_sanitize_response_guardrail("response", response_sanitizer) + + host = RecordingHostStub() + host.failures[failure] = "error" + service = _service(CodecFailurePlugin(), host) + await _register(service) + + if surface == pb.LLM_SANITIZE_REQUEST_GUARDRAIL: + payload = _llm_payload(request={"headers": {}, "content": {"model": "test"}}) + payload.request_sanitize_context.CopyFrom( + pb.LlmSanitizeRequestContext( + codec=pb.LlmCodecIdentity(kind=pb.LLM_CODEC_KIND_OPAQUE), + codec_capability_id="request-capability", + ) + ) + else: + payload = _llm_payload(response={"model": "test", "message": "secret"}) + payload.response_sanitize_context.CopyFrom( + pb.LlmSanitizeResponseContext( + codec=pb.LlmCodecIdentity(kind=pb.LLM_CODEC_KIND_OPAQUE), + codec_capability_id="response-capability", + ) + ) + + result = await service.Invoke( + _invoke_request(registration_name, surface, llm=payload), + AbortContext(), + ) + assert result.WhichOneof("result") == "error" + assert f"{failure} failed" in result.error.message + + @pytest.mark.parametrize( "surface", [ diff --git a/python/tests/test_builtin_codecs.py b/python/tests/test_builtin_codecs.py index 99beb9054..156122583 100644 --- a/python/tests/test_builtin_codecs.py +++ b/python/tests/test_builtin_codecs.py @@ -394,7 +394,8 @@ def test_manual_call_end_response_codec_uses_sanitized_payload(self): def capture(event): captured_events.append(event) - def sanitize_response(response): + def sanitize_response(response, context): + del context return { "id": "chatcmpl-sanitized", "model": "gpt-4", diff --git a/python/tests/test_llm.py b/python/tests/test_llm.py index f2575c9aa..9a3662be4 100644 --- a/python/tests/test_llm.py +++ b/python/tests/test_llm.py @@ -22,6 +22,7 @@ scope, subscribers, ) +from nemo_relay.codecs import OpenAIChatCodec def make_request(): @@ -115,8 +116,22 @@ async def func(request): class TestLLMGuardrails: + @pytest.mark.parametrize( + ("register", "callback"), + [ + (guardrails.register_llm_sanitize_request, lambda request: request), + (guardrails.register_llm_sanitize_response, lambda response: response), + (guardrails.register_llm_sanitize_request, object()), + (guardrails.register_llm_sanitize_response, object()), + ], + ) + def test_sanitizer_registration_rejects_legacy_or_uninspectable_callbacks(self, register, callback): + with pytest.raises(TypeError, match="payload, context"): + register("py_llm_invalid_signature", 1, callback) + def test_sanitize_request_guardrail(self): - def sanitizer(request): + def sanitizer(request, context): + del context # request is an LLMRequest object; must return a new LLMRequest headers = request.headers headers["X-Sanitized"] = "true" @@ -126,7 +141,8 @@ def sanitizer(request): guardrails.deregister_llm_sanitize_request("py_llm_san_req") def test_sanitize_response_guardrail(self): - def sanitizer(response): + def sanitizer(response, context): + del context # response is a plain dict response["cleaned"] = True return response @@ -134,6 +150,112 @@ def sanitizer(response): guardrails.register_llm_sanitize_response("py_llm_san_resp", 1, sanitizer) guardrails.deregister_llm_sanitize_response("py_llm_san_resp") + def test_sanitizers_receive_a_structured_codec_context(self): + request_contexts = [] + response_contexts = [] + + def sanitize_request(request, context): + request_contexts.append(context) + return request + + def sanitize_response(response, context): + response_contexts.append(context) + return response + + guardrails.register_llm_sanitize_request("py_llm_structured_context_request", 1, sanitize_request) + guardrails.register_llm_sanitize_response("py_llm_structured_context_response", 1, sanitize_response) + try: + handle = llm.call("py_llm_structured_context", make_request()) + llm.call_end(handle, {"response": "ok"}) + finally: + guardrails.deregister_llm_sanitize_request("py_llm_structured_context_request") + guardrails.deregister_llm_sanitize_response("py_llm_structured_context_response") + + assert len(request_contexts) == 1 + assert len(response_contexts) == 1 + for context in [*request_contexts, *response_contexts]: + assert context.codec.kind == "none" + assert context.codec.id is None + + async def test_sanitizers_resolve_active_builtin_codecs(self): + request_codec_used = False + response_codec_used = False + + def sanitize_request(request, context): + nonlocal request_codec_used + assert context.codec.kind == "builtin" + assert context.codec.id == "openai_chat" + codec = context.resolve_codec() + assert codec is not None + annotated = codec.decode(request) + request_codec_used = True + return codec.encode(annotated, request) + + def sanitize_response(response, context): + nonlocal response_codec_used + assert context.codec.kind == "builtin" + assert context.codec.id == "openai_chat" + codec = context.resolve_codec() + assert codec is not None + assert codec.decode_response(response).model == "test-model" + response_codec_used = True + return response + + codec = OpenAIChatCodec() + response = { + "id": "chatcmpl-python", + "model": "test-model", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}}], + } + guardrails.register_llm_sanitize_request("py_llm_builtin_context_request", 1, sanitize_request) + guardrails.register_llm_sanitize_response("py_llm_builtin_context_response", 1, sanitize_response) + try: + result = await llm.execute( + "py_llm_builtin_context", + make_request(), + lambda request: response, + codec=codec, + response_codec=codec, + ) + finally: + guardrails.deregister_llm_sanitize_request("py_llm_builtin_context_request") + guardrails.deregister_llm_sanitize_response("py_llm_builtin_context_response") + + assert result == response + assert request_codec_used is True + assert response_codec_used is True + + def test_none_omits_payload_and_short_circuits_later_sanitizers(self): + events = [] + later_called = False + + def omit(request, context): + return None + + def later(request, context): + nonlocal later_called + later_called = True + return request + + subscribers.register("py_llm_none_sanitize_sub", events.append) + guardrails.register_llm_sanitize_request("py_llm_none_sanitize_first", 1, omit) + guardrails.register_llm_sanitize_request("py_llm_none_sanitize_later", 2, later) + try: + handle = llm.call("py_llm_none_sanitize", make_request()) + llm.call_end(handle, {"ok": True}) + finally: + guardrails.deregister_llm_sanitize_request("py_llm_none_sanitize_first") + guardrails.deregister_llm_sanitize_request("py_llm_none_sanitize_later") + try: + subscribers.flush() + finally: + subscribers.deregister("py_llm_none_sanitize_sub") + + start = _llm_event(events, "py_llm_none_sanitize", "start") + assert start.data is None + assert start.annotated_request is None + assert later_called is False + def test_conditional_execution_guardrail(self): def checker(request): return None @@ -148,18 +270,18 @@ def test_conditional_execution_direct(self): guardrails.deregister_llm_conditional_execution("py_llm_cond_direct") def test_duplicate_raises(self): - guardrails.register_llm_sanitize_request("py_llm_dup", 1, lambda r: r) + guardrails.register_llm_sanitize_request("py_llm_dup", 1, lambda r, context: r) with pytest.raises(RuntimeError): - guardrails.register_llm_sanitize_request("py_llm_dup", 1, lambda r: r) + guardrails.register_llm_sanitize_request("py_llm_dup", 1, lambda r, context: r) guardrails.deregister_llm_sanitize_request("py_llm_dup") - def test_sanitize_request_callable_error_falls_back_to_original_input(self): + def test_sanitize_request_callable_error_omits_observability_input(self): events = [] subscribers.register("py_llm_sanitize_req_sub", lambda event: events.append(event)) guardrails.register_llm_sanitize_request( "py_llm_sanitize_req_fail", 1, - lambda request: raise_runtime_error("boom"), + lambda request, context: raise_runtime_error("boom"), ) try: request = make_request() @@ -173,16 +295,16 @@ def test_sanitize_request_callable_error_falls_back_to_original_input(self): subscribers.deregister("py_llm_sanitize_req_sub") start = _llm_event(events, "llm_sanitize_req_fail", "start") - request = make_request() - assert start.data == {"headers": request.headers, "content": request.content} + assert start.data is None + assert start.annotated_request is None - def test_sanitize_request_invalid_return_falls_back_to_original_input(self): + def test_sanitize_request_invalid_return_omits_observability_input(self): events = [] subscribers.register("py_llm_sanitize_req_bad_sub", lambda event: events.append(event)) guardrails.register_llm_sanitize_request( "py_llm_sanitize_req_bad", 1, - cast(guardrails.LlmSanitizeRequestGuardrail, lambda request: object()), + cast(guardrails.LlmSanitizeRequestGuardrail, lambda request, context: object()), ) try: request = make_request() @@ -196,16 +318,16 @@ def test_sanitize_request_invalid_return_falls_back_to_original_input(self): subscribers.deregister("py_llm_sanitize_req_bad_sub") start = _llm_event(events, "llm_sanitize_req_bad", "start") - request = make_request() - assert start.data == {"headers": request.headers, "content": request.content} + assert start.data is None + assert start.annotated_request is None - def test_sanitize_response_callable_error_falls_back_to_original_output(self): + def test_sanitize_response_callable_error_omits_observability_output(self): events = [] subscribers.register("py_llm_sanitize_resp_sub", lambda event: events.append(event)) guardrails.register_llm_sanitize_response( "py_llm_sanitize_resp_fail", 1, - lambda response: raise_runtime_error("boom"), + lambda response, context: raise_runtime_error("boom"), ) try: handle = llm.call("llm_sanitize_resp_fail", make_request()) @@ -218,15 +340,16 @@ def test_sanitize_response_callable_error_falls_back_to_original_output(self): subscribers.deregister("py_llm_sanitize_resp_sub") end = _llm_event(events, "llm_sanitize_resp_fail", "end") - assert end.data == {"ok": True} + assert end.data is None + assert end.annotated_response is None - def test_sanitize_response_invalid_return_falls_back_to_original_output(self): + def test_sanitize_response_invalid_return_omits_observability_output(self): events = [] subscribers.register("py_llm_sanitize_resp_bad_sub", lambda event: events.append(event)) guardrails.register_llm_sanitize_response( "py_llm_sanitize_resp_bad", 1, - cast(guardrails.LlmSanitizeResponseGuardrail, lambda response: object()), + cast(guardrails.LlmSanitizeResponseGuardrail, lambda response, context: object()), ) try: handle = llm.call("llm_sanitize_resp_bad", make_request()) @@ -239,7 +362,29 @@ def test_sanitize_response_invalid_return_falls_back_to_original_output(self): subscribers.deregister("py_llm_sanitize_resp_bad_sub") end = _llm_event(events, "llm_sanitize_resp_bad", "end") - assert end.data == {"ok": True} + assert end.data is None + assert end.annotated_response is None + + def test_sanitize_response_guardrail_accepts_scalar_json_payloads(self): + events = [] + subscribers.register("py_llm_sanitize_scalar_sub", lambda event: events.append(event)) + guardrails.register_llm_sanitize_response( + "py_llm_sanitize_scalar", + 1, + lambda response, context: f"sanitized:{response}", + ) + try: + handle = llm.call("llm_sanitize_scalar", make_request()) + llm.call_end(handle, "raw-response") + finally: + guardrails.deregister_llm_sanitize_response("py_llm_sanitize_scalar") + try: + subscribers.flush() + finally: + subscribers.deregister("py_llm_sanitize_scalar_sub") + + end = _llm_event(events, "llm_sanitize_scalar", "end") + assert end.data == "sanitized:raw-response" def test_deregister_nonexistent(self): assert not guardrails.deregister_llm_sanitize_request("nope") diff --git a/python/tests/test_scope_local.py b/python/tests/test_scope_local.py index 654400805..d20c7382d 100644 --- a/python/tests/test_scope_local.py +++ b/python/tests/test_scope_local.py @@ -576,6 +576,20 @@ def test_deregister_nonexistent_returns_false(self): class TestScopeLocalLlmWrappers: + @pytest.mark.parametrize( + ("register", "callback"), + [ + (scope_local.register_llm_sanitize_request, lambda request: request), + (scope_local.register_llm_sanitize_response, lambda response: response), + (scope_local.register_llm_sanitize_request, object()), + (scope_local.register_llm_sanitize_response, object()), + ], + ) + def test_sanitizer_registration_rejects_legacy_or_uninspectable_callbacks(self, register, callback): + with scope.scope("invalid_scope_local_sanitizer", ScopeType.Agent) as handle: + with pytest.raises(TypeError, match="payload, context"): + register(handle, "invalid_scope_local_sanitizer", 1, callback) + def test_register_and_deregister_scope_local_wrappers(self): """Scope-local wrapper functions round-trip through the native API for both tool and LLM middleware.""" request = LLMRequest({}, {"messages": [], "model": "scope-local"}) @@ -602,10 +616,20 @@ async def stream_intercept(request_inner, next_fn): ) assert scope_local.deregister_tool_execution(handle, "sl_tool_exec_cov") is True - scope_local.register_llm_sanitize_request(handle, "sl_llm_req_cov", 1, lambda req: req) + scope_local.register_llm_sanitize_request( + handle, + "sl_llm_req_cov", + 1, + lambda req, _context: req, + ) assert scope_local.deregister_llm_sanitize_request(handle, "sl_llm_req_cov") is True - scope_local.register_llm_sanitize_response(handle, "sl_llm_resp_cov", 1, lambda response: response) + scope_local.register_llm_sanitize_response( + handle, + "sl_llm_resp_cov", + 1, + lambda response, _context: response, + ) assert scope_local.deregister_llm_sanitize_response(handle, "sl_llm_resp_cov") is True scope_local.register_llm_conditional_execution(handle, "sl_llm_cond_cov", 1, lambda req: None) @@ -635,7 +659,8 @@ async def test_scope_local_llm_sanitize_request_rewrites_event_input(self): events = [] request = LLMRequest({}, {"messages": [], "model": "scope-local"}) - def sanitize_request(req): + def sanitize_request(req, context): + del context return LLMRequest({"X-Scope-Local": "yes"}, req.content) with scope.scope("sl_llm_sanitize_scope", ScopeType.Agent) as handle: From e5db478faca03bc773fde2eee678f7ca6b0d9f7e Mon Sep 17 00:00:00 2001 From: Will Killian Date: Mon, 27 Jul 2026 10:21:27 -0400 Subject: [PATCH 9/9] feat: add cross-boundary Relay propagation Signed-off-by: Will Killian --- crates/core/src/api/llm.rs | 19 ++- crates/core/src/api/runtime.rs | 10 +- crates/core/src/api/runtime/scope_stack.rs | 136 ++++++++++++++++++ crates/core/src/api/tool.rs | 4 +- .../core/src/observability/openinference.rs | 31 +++- crates/core/src/observability/otel.rs | 31 +++- .../integration/context_isolation_tests.rs | 38 ++++- crates/ffi/nemo_relay.h | 33 +++++ crates/ffi/src/api/scope_stack.rs | 123 +++++++++++++++- crates/node/src/api/mod.rs | 92 +++++++++++- crates/node/tests/context_tests.mjs | 23 +++ crates/python/src/py_api/mod.rs | 44 +++++- crates/python/src/py_types/core.rs | 47 +++++- crates/python/src/py_types/mod.rs | 1 + docs/about-nemo-relay/concepts/scopes.mdx | 18 +++ .../observability/opentelemetry.mdx | 6 + go/nemo_relay/context_test.go | 24 ++++ go/nemo_relay/nemo_relay.go | 59 ++++++++ python/nemo_relay/__init__.py | 45 ++++++ python/nemo_relay/__init__.pyi | 7 + python/nemo_relay/_native.pyi | 13 ++ python/tests/test_context_isolation.py | 15 ++ 22 files changed, 786 insertions(+), 33 deletions(-) diff --git a/crates/core/src/api/llm.rs b/crates/core/src/api/llm.rs index 29e4655a9..ea0b49af0 100644 --- a/crates/core/src/api/llm.rs +++ b/crates/core/src/api/llm.rs @@ -22,6 +22,7 @@ use crate::api::runtime::global_context; use crate::api::runtime::{ EventSubscriberFn, LlmCollectorFn, LlmExecutionNextFn, LlmFinalizerFn, LlmJsonStream, LlmSanitizeRequestContext, LlmSanitizeResponseContext, LlmStreamExecutionNextFn, + with_active_event_uuid, }; use crate::api::runtime::{ScopeStackHandle, current_scope_stack}; use crate::api::scope::event; @@ -1025,7 +1026,9 @@ pub async fn llm_call_execute(params: LlmCallExecuteParams) -> Result { emit_optimization_marks(&handle, &lifecycle_subscribers); let execution_name = name.clone(); - let execution = + let event_uuid = handle.uuid; + let execution = with_active_event_uuid( + event_uuid, scope_llm_optimization_recorder(handle.optimization_recorder.clone(), async move { let execution = { let scope_stack = current_scope_stack(); @@ -1040,8 +1043,9 @@ pub async fn llm_call_execute(params: LlmCallExecuteParams) -> Result { state.llm_build_execution_chain(&execution_name, func, &scope_locals) }; execution(intercepted_request).await - }) - .await; + }), + ) + .await; match execution { Ok(response) => { @@ -1215,7 +1219,9 @@ pub async fn llm_stream_call_execute(params: LlmStreamCallExecuteParams) -> Resu emit_optimization_marks(&handle, &lifecycle_subscribers); let execution_name = name.clone(); - let execution = + let event_uuid = handle.uuid; + let execution = with_active_event_uuid( + event_uuid, scope_llm_optimization_recorder(handle.optimization_recorder.clone(), async move { let execution = { let scope_stack = current_scope_stack(); @@ -1230,8 +1236,9 @@ pub async fn llm_stream_call_execute(params: LlmStreamCallExecuteParams) -> Resu state.llm_stream_build_execution_chain(&execution_name, func, &scope_locals) }; execution(intercepted_request).await - }) - .await; + }), + ) + .await; match execution { Ok(raw_stream) => { diff --git a/crates/core/src/api/runtime.rs b/crates/core/src/api/runtime.rs index 0261657ba..77670c804 100644 --- a/crates/core/src/api/runtime.rs +++ b/crates/core/src/api/runtime.rs @@ -19,10 +19,12 @@ pub use callbacks::{ }; pub use global::global_context; pub use scope_stack::{ - ScopeStack, ScopeStackHandle, TASK_SCOPE_STACK, ThreadScopeStackBinding, - capture_thread_scope_stack, create_scope_stack, current_scope_stack, propagate_scope_to_thread, - restore_thread_scope_stack, scope_stack_active, set_thread_scope_stack, - sync_thread_scope_stack, task_scope_push, task_scope_remove, task_scope_top, with_scope_stack, + PropagationContext, ScopeStack, ScopeStackHandle, TASK_SCOPE_STACK, ThreadScopeStackBinding, + capture_propagation_context, capture_propagation_context_with_root, capture_thread_scope_stack, + create_scope_stack, create_scope_stack_from_propagation, current_scope_stack, + propagate_scope_to_thread, restore_thread_scope_stack, scope_stack_active, + set_thread_scope_stack, sync_thread_scope_stack, task_scope_push, task_scope_remove, + task_scope_top, with_active_event_uuid, with_scope_stack, }; pub use state::NemoRelayContextState; pub use subscriber_dispatcher::flush_subscribers; diff --git a/crates/core/src/api/runtime/scope_stack.rs b/crates/core/src/api/runtime/scope_stack.rs index f62f7191b..ac2a70899 100644 --- a/crates/core/src/api/runtime/scope_stack.rs +++ b/crates/core/src/api/runtime/scope_stack.rs @@ -10,8 +10,10 @@ use std::cell::RefCell; use std::collections::{HashMap, HashSet}; +use std::future::Future; use std::sync::{Arc, RwLock}; +use serde::{Deserialize, Serialize}; use uuid::Uuid; use crate::api::runtime::callbacks::EventSubscriberFn; @@ -31,6 +33,51 @@ pub struct ScopeStack { stack: Vec, scope_registries: HashMap, fresh_agents: HashSet, + propagated_parent_uuid: Option, +} + +/// Versioned, transport-neutral causal context for crossing a Relay boundary. +/// +/// Applications are responsible for serializing, transporting, authenticating, +/// and trusting this value. It intentionally contains only Relay identifiers; +/// OpenTelemetry `traceparent` and `tracestate` remain transport sidecars. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct PropagationContext { + /// Wire-format version. Version 1 is the only currently supported value. + pub version: u16, + /// Stable session root when the sending application knows one. + #[serde(skip_serializing_if = "Option::is_none")] + pub root_uuid: Option, + /// Immediate Relay event or scope that caused the boundary crossing. + pub parent_uuid: Uuid, +} + +impl PropagationContext { + /// The current wire-format version. + pub const VERSION: u16 = 1; + + /// Validate a context received from an untrusted transport. + pub fn validate(&self) -> Result<()> { + if self.version != Self::VERSION { + return Err(FlowError::InvalidArgument(format!( + "unsupported propagation context version {}; expected {}", + self.version, + Self::VERSION + ))); + } + for (name, uuid) in [("parent_uuid", self.parent_uuid)] + .into_iter() + .chain(self.root_uuid.map(|uuid| ("root_uuid", uuid))) + { + let bytes = uuid.as_bytes(); + if bytes.iter().all(|byte| *byte == 0) || bytes[8..].iter().all(|byte| *byte == 0) { + return Err(FlowError::InvalidArgument(format!( + "propagation context {name} is not a usable Relay identifier" + ))); + } + } + Ok(()) + } } impl ScopeStack { @@ -49,7 +96,52 @@ impl ScopeStack { stack: vec![root], scope_registries: HashMap::new(), fresh_agents: HashSet::from([root_uuid]), + propagated_parent_uuid: None, + } + } + + fn from_propagation(context: &PropagationContext) -> Result { + context.validate()?; + let (root, parent) = match context.root_uuid { + Some(root_uuid) => { + let root = ScopeHandle::builder() + .uuid(root_uuid) + .name("propagated-root") + .scope_type(ScopeType::Agent) + .build(); + let parent = (root_uuid != context.parent_uuid).then(|| { + ScopeHandle::builder() + .uuid(context.parent_uuid) + .parent_uuid(root_uuid) + .name("propagated-parent") + .scope_type(ScopeType::Unknown) + .build() + }); + (root, parent) + } + None => ( + ScopeHandle::builder() + .uuid(context.parent_uuid) + .name("propagated-root") + .scope_type(ScopeType::Agent) + .build(), + None, + ), + }; + let root_uuid = root.uuid; + let mut stack = vec![root]; + if let Some(parent) = parent { + stack.push(parent); } + Ok(Self { + stack, + scope_registries: HashMap::new(), + fresh_agents: HashSet::from([root_uuid]), + propagated_parent_uuid: context + .root_uuid + .filter(|root_uuid| *root_uuid != context.parent_uuid) + .map(|_| context.parent_uuid), + }) } /// Push a scope handle onto the top of the stack. @@ -98,6 +190,11 @@ impl ScopeStack { .uuid } + /// Whether `uuid` is the synthetic parent imported from propagation. + pub fn is_propagated_parent(&self, uuid: Uuid) -> bool { + self.propagated_parent_uuid == Some(uuid) + } + /// Return the full ordered stack of scope handles. /// /// # Returns @@ -290,9 +387,48 @@ pub fn create_scope_stack() -> ScopeStackHandle { Arc::new(RwLock::new(ScopeStack::new())) } +/// Create an isolated scope stack rooted below a supplied propagation context. +/// +/// The imported handles are synthetic bookkeeping only; Relay never emits their +/// lifecycle events or transfers scope-local registrations across the boundary. +pub fn create_scope_stack_from_propagation( + context: &PropagationContext, +) -> Result { + Ok(Arc::new(RwLock::new(ScopeStack::from_propagation( + context, + )?))) +} + +/// Capture the current causal parent without asserting a session root. +pub fn capture_propagation_context() -> Result { + capture_propagation_context_with_root(None) +} + +/// Capture the current causal parent and an application-supplied session root. +pub fn capture_propagation_context_with_root( + root_uuid: Option, +) -> Result { + let context = PropagationContext { + version: PropagationContext::VERSION, + root_uuid, + parent_uuid: ACTIVE_EVENT_UUID + .try_with(|uuid| *uuid) + .unwrap_or_else(|_| task_scope_top().uuid), + }; + context.validate()?; + Ok(context) +} + tokio::task_local! { /// Task-local scope stack handle used by async execution contexts. pub static TASK_SCOPE_STACK: ScopeStackHandle; + /// Managed tool or LLM event currently executing in this task. + static ACTIVE_EVENT_UUID: Uuid; +} + +/// Run a future with `uuid` as the causally active managed event. +pub async fn with_active_event_uuid(uuid: Uuid, future: impl Future) -> T { + ACTIVE_EVENT_UUID.scope(uuid, future).await } thread_local! { diff --git a/crates/core/src/api/tool.rs b/crates/core/src/api/tool.rs index 2655fc857..6fbe6cd70 100644 --- a/crates/core/src/api/tool.rs +++ b/crates/core/src/api/tool.rs @@ -7,7 +7,7 @@ use crate::api::event::{BaseEvent, Event, MarkEvent, PendingMarkSpec}; use crate::api::runtime::NemoRelayContextState; use crate::api::runtime::current_scope_stack; use crate::api::runtime::global_context; -use crate::api::runtime::{EventSubscriberFn, ToolExecutionNextFn}; +use crate::api::runtime::{EventSubscriberFn, ToolExecutionNextFn, with_active_event_uuid}; use crate::api::scope::event; use crate::api::scope::{EmitMarkEventParams, ScopeHandle}; use crate::api::shared::{ @@ -551,7 +551,7 @@ pub async fn tool_call_execute(params: ToolCallExecuteParams) -> Result { state.tool_build_execution_chain(&name, func, &scope_locals) }; - match execution(intercepted_args).await { + match with_active_event_uuid(handle.uuid, execution(intercepted_args)).await { Ok(outcome) => { let ToolExecutionInterceptOutcome { result, diff --git a/crates/core/src/observability/openinference.rs b/crates/core/src/observability/openinference.rs index 362015853..05d980869 100644 --- a/crates/core/src/observability/openinference.rs +++ b/crates/core/src/observability/openinference.rs @@ -29,7 +29,7 @@ use super::{ relay_trace_id, validate_attribute_mappings, }; use crate::api::event::{Event, EventNormalizationExt, ScopeCategory}; -use crate::api::runtime::EventSubscriberFn; +use crate::api::runtime::{EventSubscriberFn, current_scope_stack}; use crate::api::scope::ScopeType; use crate::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; use crate::codec::request::{ @@ -42,7 +42,8 @@ use chrono::{DateTime, Utc}; use openinference_semantic_conventions::SpanKind as OpenInferenceSpanKind; use openinference_semantic_conventions::attributes as oi; use opentelemetry::trace::{ - Span as _, SpanContext, SpanKind, TraceContextExt, Tracer, TracerProvider as _, + Span as _, SpanContext, SpanKind, TraceContextExt, TraceFlags, TraceState, Tracer, + TracerProvider as _, }; use opentelemetry::{Context, KeyValue}; use opentelemetry_otlp::{Protocol, SpanExporter, WithExportConfig, WithHttpConfig}; @@ -788,11 +789,31 @@ impl OpenInferenceEventProcessor { if let Some(active_span) = self.find_parent_span(event) { return Context::new().with_remote_span_context(active_span.span_context.clone()); } - event + if let Some(span_context) = event .parent_uuid() .and_then(|uuid| self.completed_span_contexts.get(&uuid)) - .map(|span_context| Context::new().with_remote_span_context(span_context.clone())) - .unwrap_or_default() + { + return Context::new().with_remote_span_context(span_context.clone()); + } + let Some(parent_uuid) = event.parent_uuid() else { + return Context::new(); + }; + let stack = current_scope_stack(); + let stack = stack.read().expect("scope stack lock poisoned"); + if !stack.is_propagated_parent(parent_uuid) { + return Context::new(); + } + let root_uuid = stack.root_uuid(); + if parent_uuid == root_uuid { + return Context::new(); + } + Context::new().with_remote_span_context(SpanContext::new( + relay_trace_id(root_uuid), + relay_span_id(parent_uuid), + TraceFlags::SAMPLED, + true, + TraceState::default(), + )) } fn parent_span_uuid(&self, event: &Event) -> Option { diff --git a/crates/core/src/observability/otel.rs b/crates/core/src/observability/otel.rs index d9b19a1d7..841e58f03 100644 --- a/crates/core/src/observability/otel.rs +++ b/crates/core/src/observability/otel.rs @@ -29,14 +29,15 @@ use super::{ relay_trace_id, validate_attribute_mappings, }; use crate::api::event::{Event, EventNormalizationExt, ScopeCategory}; -use crate::api::runtime::EventSubscriberFn; +use crate::api::runtime::{EventSubscriberFn, current_scope_stack}; use crate::api::scope::ScopeType; use crate::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; use crate::codec::response::CostEstimate; use crate::error::FlowError; use chrono::{DateTime, Utc}; use opentelemetry::trace::{ - Span as _, SpanContext, SpanKind, TraceContextExt, Tracer, TracerProvider as _, + Span as _, SpanContext, SpanKind, TraceContextExt, TraceFlags, TraceState, Tracer, + TracerProvider as _, }; use opentelemetry::{Context, KeyValue}; use opentelemetry_otlp::{Protocol, SpanExporter, WithExportConfig, WithHttpConfig}; @@ -774,11 +775,31 @@ impl OtelEventProcessor { if let Some(active_span) = self.find_parent_span(event) { return Context::new().with_remote_span_context(active_span.span_context.clone()); } - event + if let Some(span_context) = event .parent_uuid() .and_then(|uuid| self.completed_span_contexts.get(&uuid)) - .map(|span_context| Context::new().with_remote_span_context(span_context.clone())) - .unwrap_or_default() + { + return Context::new().with_remote_span_context(span_context.clone()); + } + let Some(parent_uuid) = event.parent_uuid() else { + return Context::new(); + }; + let stack = current_scope_stack(); + let stack = stack.read().expect("scope stack lock poisoned"); + if !stack.is_propagated_parent(parent_uuid) { + return Context::new(); + } + let root_uuid = stack.root_uuid(); + if parent_uuid == root_uuid { + return Context::new(); + } + Context::new().with_remote_span_context(SpanContext::new( + relay_trace_id(root_uuid), + relay_span_id(parent_uuid), + TraceFlags::SAMPLED, + true, + TraceState::default(), + )) } fn parent_span_uuid(&self, event: &Event) -> Option { diff --git a/crates/core/tests/integration/context_isolation_tests.rs b/crates/core/tests/integration/context_isolation_tests.rs index 22cf33ba2..33a28e530 100644 --- a/crates/core/tests/integration/context_isolation_tests.rs +++ b/crates/core/tests/integration/context_isolation_tests.rs @@ -6,9 +6,10 @@ use std::sync::Arc; use nemo_relay::api::runtime::{ - ScopeStack, TASK_SCOPE_STACK, create_scope_stack, current_scope_stack, - propagate_scope_to_thread, scope_stack_active, set_thread_scope_stack, sync_thread_scope_stack, - task_scope_push, task_scope_remove, task_scope_top, + PropagationContext, ScopeStack, TASK_SCOPE_STACK, create_scope_stack, + create_scope_stack_from_propagation, current_scope_stack, propagate_scope_to_thread, + scope_stack_active, set_thread_scope_stack, sync_thread_scope_stack, task_scope_push, + task_scope_remove, task_scope_top, }; use nemo_relay::api::scope::{ PopScopeParams, PushScopeParams, ScopeHandle, ScopeType, pop_scope, push_scope, @@ -55,6 +56,37 @@ fn test_two_scope_stacks_are_independent() { assert_ne!(root_a_uuid, root_b_uuid); // scope_a != scope_b } +#[test] +fn test_propagation_context_seeds_a_synthetic_root_and_parent() { + let root_uuid = Uuid::now_v7(); + let parent_uuid = Uuid::now_v7(); + let stack = create_scope_stack_from_propagation(&PropagationContext { + version: PropagationContext::VERSION, + root_uuid: Some(root_uuid), + parent_uuid, + }) + .unwrap(); + let stack = stack.read().unwrap(); + assert_eq!(stack.root_uuid(), root_uuid); + assert_eq!(stack.top().uuid, parent_uuid); + assert_eq!(stack.scopes().len(), 2); +} + +#[test] +fn test_rootless_propagation_context_uses_the_parent_as_root() { + let parent_uuid = Uuid::now_v7(); + let stack = create_scope_stack_from_propagation(&PropagationContext { + version: PropagationContext::VERSION, + root_uuid: None, + parent_uuid, + }) + .unwrap(); + let stack = stack.read().unwrap(); + assert_eq!(stack.root_uuid(), parent_uuid); + assert_eq!(stack.top().uuid, parent_uuid); + assert_eq!(stack.scopes().len(), 1); +} + #[test] fn test_pop_scope_rejects_non_top_and_unknown_handles() { set_thread_scope_stack(create_scope_stack()); diff --git a/crates/ffi/nemo_relay.h b/crates/ffi/nemo_relay.h index 3c381015a..be8822bc4 100644 --- a/crates/ffi/nemo_relay.h +++ b/crates/ffi/nemo_relay.h @@ -2341,6 +2341,39 @@ NemoRelayStatus nemo_relay_scope_deregister_subscriber(const char *scope_uuid, c */ NemoRelayStatus nemo_relay_scope_stack_create(struct FfiScopeStack **out); +/** + * Serialize the current causal parent as a versioned propagation context. + * + * The returned JSON must be freed with `nemo_relay_string_free`. + * + * # Safety + * `out` must be a valid, writable pointer to a C-string output slot. + */ +NemoRelayStatus nemo_relay_capture_propagation_context_json(char **out); + +/** + * Serialize the current causal parent with an application-supplied root UUID. + * + * Pass null for `root_uuid` to omit the root. The returned JSON must be freed + * with `nemo_relay_string_free`. + * + * # Safety + * When non-null, `root_uuid` must point to a valid NUL-terminated C string; + * `out` must be a valid, writable pointer to a C-string output slot. + */ +NemoRelayStatus nemo_relay_capture_propagation_context_with_root_json(const char *root_uuid, + char **out); + +/** + * Create an isolated scope stack from propagation-context JSON. + * + * # Safety + * `context_json` must point to a valid NUL-terminated C string and `out` must + * be a valid, writable pointer to a scope-stack output slot. + */ +NemoRelayStatus nemo_relay_scope_stack_create_from_propagation_json(const char *context_json, + struct FfiScopeStack **out); + /** * Bind an isolated scope stack to the current OS thread. * diff --git a/crates/ffi/src/api/scope_stack.rs b/crates/ffi/src/api/scope_stack.rs index abaa761b1..3115ee74b 100644 --- a/crates/ffi/src/api/scope_stack.rs +++ b/crates/ffi/src/api/scope_stack.rs @@ -2,10 +2,15 @@ // SPDX-License-Identifier: Apache-2.0 use super::{ - FfiScopeStack, FfiThreadScopeStackBinding, NemoRelayStatus, capture_thread_scope_stack, - clear_last_error, create_scope_stack, restore_thread_scope_stack, scope_stack_active, - set_last_error, set_thread_scope_stack, + FfiScopeStack, FfiThreadScopeStackBinding, NemoRelayStatus, c_char, c_str_to_string, + capture_thread_scope_stack, clear_last_error, create_scope_stack, json_to_c_string, + restore_thread_scope_stack, scope_stack_active, set_last_error, set_thread_scope_stack, }; +use nemo_relay::api::runtime::{ + PropagationContext, capture_propagation_context, capture_propagation_context_with_root, + create_scope_stack_from_propagation, +}; +use uuid::Uuid; // --------------------------------------------------------------------------- // Scope stack isolation @@ -42,6 +47,118 @@ pub unsafe extern "C" fn nemo_relay_scope_stack_create( NemoRelayStatus::Ok } +/// Serialize the current causal parent as a versioned propagation context. +/// +/// The returned JSON must be freed with `nemo_relay_string_free`. +/// +/// # Safety +/// `out` must be a valid, writable pointer to a C-string output slot. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_capture_propagation_context_json( + out: *mut *mut c_char, +) -> NemoRelayStatus { + clear_last_error(); + if out.is_null() { + set_last_error("out pointer is null"); + return NemoRelayStatus::NullPointer; + } + match capture_propagation_context().and_then(|context| { + serde_json::to_value(context) + .map_err(|error| nemo_relay::error::FlowError::Internal(error.to_string())) + }) { + Ok(context) => { + unsafe { *out = json_to_c_string(&context) }; + NemoRelayStatus::Ok + } + Err(error) => { + set_last_error(&error.to_string()); + NemoRelayStatus::from(&error) + } + } +} + +/// Serialize the current causal parent with an application-supplied root UUID. +/// +/// Pass null for `root_uuid` to omit the root. The returned JSON must be freed +/// with `nemo_relay_string_free`. +/// +/// # Safety +/// When non-null, `root_uuid` must point to a valid NUL-terminated C string; +/// `out` must be a valid, writable pointer to a C-string output slot. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_capture_propagation_context_with_root_json( + root_uuid: *const c_char, + out: *mut *mut c_char, +) -> NemoRelayStatus { + clear_last_error(); + if out.is_null() { + set_last_error("out pointer is null"); + return NemoRelayStatus::NullPointer; + } + let root_uuid = if root_uuid.is_null() { + Ok(None) + } else { + c_str_to_string(root_uuid) + .map_err(|_| ()) + .and_then(|value| Uuid::parse_str(&value).map_err(|_| ())) + .map(Some) + }; + let Ok(root_uuid) = root_uuid else { + set_last_error("root_uuid must be a valid UUID"); + return NemoRelayStatus::InvalidArg; + }; + match capture_propagation_context_with_root(root_uuid).and_then(|context| { + serde_json::to_value(context) + .map_err(|error| nemo_relay::error::FlowError::Internal(error.to_string())) + }) { + Ok(context) => { + unsafe { *out = json_to_c_string(&context) }; + NemoRelayStatus::Ok + } + Err(error) => { + set_last_error(&error.to_string()); + NemoRelayStatus::from(&error) + } + } +} + +/// Create an isolated scope stack from propagation-context JSON. +/// +/// # Safety +/// `context_json` must point to a valid NUL-terminated C string and `out` must +/// be a valid, writable pointer to a scope-stack output slot. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_scope_stack_create_from_propagation_json( + context_json: *const c_char, + out: *mut *mut FfiScopeStack, +) -> NemoRelayStatus { + clear_last_error(); + if out.is_null() { + set_last_error("out pointer is null"); + return NemoRelayStatus::NullPointer; + } + let context = match c_str_to_string(context_json) { + Ok(value) => match serde_json::from_str::(&value) { + Ok(context) => context, + Err(error) => { + set_last_error(&format!("invalid propagation context JSON: {error}")); + return NemoRelayStatus::InvalidJson; + } + }, + Err(status) => return status, + }; + match create_scope_stack_from_propagation(&context) { + Ok(stack) => { + unsafe { *out = Box::into_raw(Box::new(FfiScopeStack(stack))) }; + NemoRelayStatus::Ok + } + Err(error) => { + set_last_error(&error.to_string()); + NemoRelayStatus::from(&error) + } + } +} + /// Bind an isolated scope stack to the current OS thread. /// /// After this call, all NeMo Relay scope operations on the current thread diff --git a/crates/node/src/api/mod.rs b/crates/node/src/api/mod.rs index 1f8384fae..58dd3f672 100644 --- a/crates/node/src/api/mod.rs +++ b/crates/node/src/api/mod.rs @@ -36,9 +36,13 @@ use nemo_relay::api::runtime::{ ToolExecutionNextFn, }; use nemo_relay::api::runtime::{ - TASK_SCOPE_STACK, create_scope_stack as create_scope_stack_handle, + TASK_SCOPE_STACK, capture_propagation_context as capture_propagation_context_handle, + capture_propagation_context_with_root as capture_propagation_context_with_root_handle, + create_scope_stack as create_scope_stack_handle, + create_scope_stack_from_propagation as create_scope_stack_from_propagation_handle, current_scope_stack as current_scope_stack_handle, scope_stack_active as scope_stack_is_active, set_thread_scope_stack as bind_thread_scope_stack, task_scope_top, + with_scope_stack as with_scope_stack_handle, }; use nemo_relay::api::scope as core_scope_api; use nemo_relay::api::scope::ScopeAttributes; @@ -1726,6 +1730,48 @@ impl Plugin for NodePlugin { // Scope stack isolation // --------------------------------------------------------------------------- +/// Transport-neutral Relay causal context for application-managed transport. +#[napi(object)] +pub struct PropagationContext { + pub version: u32, + pub root_uuid: Option, + pub parent_uuid: String, +} + +fn propagation_context_from_napi( + context: PropagationContext, +) -> napi::Result { + let root_uuid = context + .root_uuid + .as_deref() + .map(uuid::Uuid::parse_str) + .transpose() + .map_err(|error| napi::Error::from_reason(format!("invalid root UUID: {error}")))?; + let parent_uuid = uuid::Uuid::parse_str(&context.parent_uuid) + .map_err(|error| napi::Error::from_reason(format!("invalid parent UUID: {error}")))?; + let version = u16::try_from(context.version) + .map_err(|_| napi::Error::from_reason("propagation context version is out of range"))?; + let context = nemo_relay::api::runtime::PropagationContext { + version, + root_uuid, + parent_uuid, + }; + context + .validate() + .map_err(|error| napi::Error::from_reason(error.to_string()))?; + Ok(context) +} + +fn propagation_context_to_napi( + context: nemo_relay::api::runtime::PropagationContext, +) -> PropagationContext { + PropagationContext { + version: u32::from(context.version), + root_uuid: context.root_uuid.map(|uuid| uuid.to_string()), + parent_uuid: context.parent_uuid.to_string(), + } +} + /// Creates a new isolated scope stack. #[napi] pub fn create_scope_stack() -> ScopeStack { @@ -1734,6 +1780,50 @@ pub fn create_scope_stack() -> ScopeStack { } } +/// Capture the current Relay causal parent for application-managed transport. +#[napi] +pub fn capture_propagation_context() -> napi::Result { + capture_propagation_context_handle() + .map(propagation_context_to_napi) + .map_err(|error| napi::Error::from_reason(error.to_string())) +} + +/// Capture the current parent with an optional stable application session root. +#[napi] +pub fn capture_propagation_context_with_root( + root_uuid: Option, +) -> napi::Result { + let root_uuid = root_uuid + .as_deref() + .map(uuid::Uuid::parse_str) + .transpose() + .map_err(|error| napi::Error::from_reason(format!("invalid root UUID: {error}")))?; + capture_propagation_context_with_root_handle(root_uuid) + .map(propagation_context_to_napi) + .map_err(|error| napi::Error::from_reason(error.to_string())) +} + +/// Create an isolated scope stack seeded from a received propagation context. +#[napi] +pub fn create_scope_stack_from_propagation( + context: PropagationContext, +) -> napi::Result { + create_scope_stack_from_propagation_handle(&propagation_context_from_napi(context)?) + .map(ScopeStack::from) + .map_err(|error| napi::Error::from_reason(error.to_string())) +} + +/// Run a synchronous callback with an isolated scope stack installed. +/// +/// For asynchronous JavaScript request handlers, keep the stack installed with +/// `setThreadScopeStack` for the handler's lifetime instead. +#[napi] +pub fn with_scope_stack(stack: &ScopeStack, callback: JsFunction) -> napi::Result { + with_scope_stack_handle(stack.inner.clone(), || { + callback.call::(None, &[]) + }) +} + /// Returns the current execution context's scope stack handle. #[napi] pub fn current_scope_stack() -> ScopeStack { diff --git a/crates/node/tests/context_tests.mjs b/crates/node/tests/context_tests.mjs index 466145545..5b26280da 100644 --- a/crates/node/tests/context_tests.mjs +++ b/crates/node/tests/context_tests.mjs @@ -18,6 +18,8 @@ const { popScope, ScopeType, ScopeStack, + createScopeStackFromPropagation, + withScopeStack, } = lib; // =========================================================================== @@ -31,6 +33,27 @@ describe('Context isolation', () => { assert.ok(stack instanceof ScopeStack, 'Expected instance of ScopeStack'); }); + it('creates an imported stack with the propagated parent on top', () => { + const original = currentScopeStack(); + const rootUuid = '018f13f0-7c1a-7a80-8000-000000000001'; + const parentUuid = '018f13f0-7c1a-7a80-8000-000000000002'; + const stack = createScopeStackFromPropagation({ version: 1, rootUuid, parentUuid }); + setThreadScopeStack(stack); + assert.equal(getHandle().uuid, parentUuid); + setThreadScopeStack(original); + }); + + it('restores the surrounding stack after withScopeStack', () => { + const original = currentScopeStack(); + const stack = createScopeStack(); + withScopeStack(stack, () => { + pushScope('temporary-with-scope-stack', ScopeType.Agent, null, null); + assert.equal(getHandle().name, 'temporary-with-scope-stack'); + }); + assert.notEqual(getHandle().name, 'temporary-with-scope-stack'); + setThreadScopeStack(original); + }); + it('currentScopeStack returns same in same context', () => { const s1 = currentScopeStack(); const s2 = currentScopeStack(); diff --git a/crates/python/src/py_api/mod.rs b/crates/python/src/py_api/mod.rs index 64e2341e3..4983b777e 100644 --- a/crates/python/src/py_api/mod.rs +++ b/crates/python/src/py_api/mod.rs @@ -17,7 +17,10 @@ use nemo_relay::api::runtime::{ LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn, ToolExecutionNextFn, }; use nemo_relay::api::runtime::{ - TASK_SCOPE_STACK, create_scope_stack as create_scope_stack_handle, + TASK_SCOPE_STACK, capture_propagation_context as capture_propagation_context_handle, + capture_propagation_context_with_root as capture_propagation_context_with_root_handle, + create_scope_stack as create_scope_stack_handle, + create_scope_stack_from_propagation as create_scope_stack_from_propagation_handle, current_scope_stack as current_scope_stack_handle, scope_stack_active as scope_stack_is_active, set_thread_scope_stack as bind_thread_scope_stack, sync_thread_scope_stack as sync_bound_thread_scope_stack, task_scope_top, @@ -38,8 +41,8 @@ use crate::convert::{json_to_py, opt_py_to_json, opt_py_to_timestamp, py_to_json use crate::py_callable; use crate::py_types::{ PyAnnotatedLLMResponse, PyAnthropicMessagesCodec, PyLLMAttributes, PyLLMHandle, PyLLMRequest, - PyLlmStream, PyOpenAIChatCodec, PyOpenAIResponsesCodec, PyScopeAttributes, PyScopeHandle, - PyScopeStack, PyScopeType, PyToolAttributes, PyToolHandle, + PyLlmStream, PyOpenAIChatCodec, PyOpenAIResponsesCodec, PyPropagationContext, + PyScopeAttributes, PyScopeHandle, PyScopeStack, PyScopeType, PyToolAttributes, PyToolHandle, }; pub(crate) type RustJsonStream = LlmJsonStream; @@ -162,6 +165,38 @@ pub fn create_scope_stack() -> PyScopeStack { PyScopeStack(create_scope_stack_handle()) } +/// Capture a transport-neutral context from the current Relay scope stack. +#[pyfunction] +pub fn capture_propagation_context() -> PyResult { + capture_propagation_context_handle() + .map(|inner| PyPropagationContext { inner }) + .map_err(to_py_err) +} + +/// Capture a context with an application-supplied stable session root UUID. +#[pyfunction] +pub fn capture_propagation_context_with_root( + root_uuid: Option<&str>, +) -> PyResult { + let root_uuid = root_uuid + .map(Uuid::parse_str) + .transpose() + .map_err(|error| PyErr::new::(error.to_string()))?; + capture_propagation_context_with_root_handle(root_uuid) + .map(|inner| PyPropagationContext { inner }) + .map_err(to_py_err) +} + +/// Create an isolated scope stack seeded from a received propagation context. +#[pyfunction] +pub fn create_scope_stack_from_propagation( + context: &PyPropagationContext, +) -> PyResult { + create_scope_stack_from_propagation_handle(&context.inner) + .map(PyScopeStack) + .map_err(to_py_err) +} + /// Bind a ``ScopeStack`` to the current thread's thread-local storage. /// /// This ensures that subsequent NeMo Relay API calls on this thread use the given @@ -1740,6 +1775,9 @@ fn scope_deregister_subscriber(scope_uuid: &str, name: &str) -> PyResult { pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> { // Scope stack creation / binding / query m.add_function(wrap_pyfunction!(create_scope_stack, m)?)?; + m.add_function(wrap_pyfunction!(capture_propagation_context, m)?)?; + m.add_function(wrap_pyfunction!(capture_propagation_context_with_root, m)?)?; + m.add_function(wrap_pyfunction!(create_scope_stack_from_propagation, m)?)?; m.add_function(wrap_pyfunction!(set_thread_scope_stack, m)?)?; m.add_function(wrap_pyfunction!(sync_thread_scope_stack, m)?)?; m.add_function(wrap_pyfunction!(py_scope_stack_active, m)?)?; diff --git a/crates/python/src/py_types/core.rs b/crates/python/src/py_types/core.rs index d0f122311..3648e629d 100644 --- a/crates/python/src/py_types/core.rs +++ b/crates/python/src/py_types/core.rs @@ -14,7 +14,9 @@ use super::{ }; use nemo_relay::api::event::{CategoryProfile, EventCategory, PendingMarkSpec}; use nemo_relay::api::llm::LlmRequestInterceptOutcome; -use nemo_relay::api::runtime::{LlmSanitizeRequestContext, LlmSanitizeResponseContext}; +use nemo_relay::api::runtime::{ + LlmSanitizeRequestContext, LlmSanitizeResponseContext, PropagationContext, +}; use nemo_relay::api::tool::ToolExecutionInterceptOutcome; /// Structured identity of the codec active during LLM sanitization. @@ -188,6 +190,49 @@ impl PyScopeStack { } } +/// Transport-neutral causal context used to continue Relay work remotely. +#[pyclass(name = "PropagationContext", skip_from_py_object)] +#[derive(Clone)] +pub struct PyPropagationContext { + pub(crate) inner: PropagationContext, +} + +#[pymethods] +impl PyPropagationContext { + #[new] + #[pyo3(signature = (parent_uuid, root_uuid=None, version=1))] + fn new(parent_uuid: &str, root_uuid: Option<&str>, version: u16) -> PyResult { + let context = PropagationContext { + version, + root_uuid: root_uuid + .map(uuid::Uuid::parse_str) + .transpose() + .map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string()))?, + parent_uuid: uuid::Uuid::parse_str(parent_uuid) + .map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string()))?, + }; + context + .validate() + .map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string()))?; + Ok(Self { inner: context }) + } + + #[getter] + fn version(&self) -> u16 { + self.inner.version + } + + #[getter] + fn root_uuid(&self) -> Option { + self.inner.root_uuid.map(|uuid| uuid.to_string()) + } + + #[getter] + fn parent_uuid(&self) -> String { + self.inner.parent_uuid.to_string() + } +} + // --------------------------------------------------------------------------- // ScopeAttributes (bitflag wrapper) // --------------------------------------------------------------------------- diff --git a/crates/python/src/py_types/mod.rs b/crates/python/src/py_types/mod.rs index eb60374d9..9a13bef16 100644 --- a/crates/python/src/py_types/mod.rs +++ b/crates/python/src/py_types/mod.rs @@ -137,6 +137,7 @@ pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> { fn register_runtime_types(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/docs/about-nemo-relay/concepts/scopes.mdx b/docs/about-nemo-relay/concepts/scopes.mdx index 84b293f96..49ac3bd6d 100644 --- a/docs/about-nemo-relay/concepts/scopes.mdx +++ b/docs/about-nemo-relay/concepts/scopes.mdx @@ -142,6 +142,24 @@ Use this when: - The boundary cannot safely carry a native stack handle - You want a clean root scope with isolated scope-local registrations +## Cross-Process Propagation + +When work crosses a process or remote-workflow boundary, applications can carry +the versioned Relay propagation context instead of a native stack handle. The +context contains an immediate `parent_uuid` and, when the application knows a +stable session root, an optional `root_uuid`. + +The receiver creates a fresh isolated stack from that context and installs it +only for request handling. Its first local event becomes a child of +`parent_uuid`; scope-local middleware and subscribers are never transferred. +The transport is application-owned: authenticate and authorize inbound context +before importing it. Relay does not send headers, make IPC connections, or +trust remote identifiers automatically. + +Relay context is distinct from W3C propagation. An integration may carry +`traceparent` and `tracestate` alongside Relay's JSON context when it needs to +preserve OpenTelemetry sampling or vendor state. + ## Practical Guidance Use these practices when applying the concept in application or integration code. diff --git a/docs/configure-plugins/observability/opentelemetry.mdx b/docs/configure-plugins/observability/opentelemetry.mdx index 48ed44383..7cd249516 100644 --- a/docs/configure-plugins/observability/opentelemetry.mdx +++ b/docs/configure-plugins/observability/opentelemetry.mdx @@ -74,6 +74,12 @@ The collector should receive OTLP trace export requests. The tracing backend should show spans for NeMo Relay scopes, tools, LLM calls, and marks grouped by root scope. +For an imported Relay propagation context with both a root and an external +parent, Relay uses the root UUID as the OTLP trace ID and the lower 64 bits of +the parent UUID as a remote parent span ID. A rootless import begins a new +native trace. Relay does not derive W3C sampling flags or `tracestate`; carry +those separately in an integration when they are required. + The default `inherit` projection follows exporter-native handling: a mark with an active parent span is a span event, while an orphan mark is a standalone zero-duration `mark:` span. `mark_projection = "event"` explicitly selects diff --git a/go/nemo_relay/context_test.go b/go/nemo_relay/context_test.go index 81ab523a4..d89891c33 100644 --- a/go/nemo_relay/context_test.go +++ b/go/nemo_relay/context_test.go @@ -259,6 +259,30 @@ func TestCreateScopeStackCreatesFreshStack(t *testing.T) { } } +func TestNewScopeStackFromPropagationUsesParentAsCurrentHandle(t *testing.T) { + rootUUID := "018f13f0-7c1a-7a80-8000-000000000001" + parentUUID := "018f13f0-7c1a-7a80-8000-000000000002" + stack, err := NewScopeStackFromPropagation(PropagationContext{ + Version: 1, + RootUUID: &rootUUID, + ParentUUID: parentUUID, + }) + if err != nil { + t.Fatalf("NewScopeStackFromPropagation failed: %v", err) + } + defer stack.Close() + + stack.Run(func() { + handle, err := GetHandle() + if err != nil { + t.Fatalf("GetHandle failed: %v", err) + } + if handle.UUID() != parentUUID { + t.Fatalf("expected parent UUID %s, got %s", parentUUID, handle.UUID()) + } + }) +} + func TestConcurrentScopeStacksWithToolCalls(t *testing.T) { const goroutines = 5 var wg sync.WaitGroup diff --git a/go/nemo_relay/nemo_relay.go b/go/nemo_relay/nemo_relay.go index 937427e5d..60057979e 100644 --- a/go/nemo_relay/nemo_relay.go +++ b/go/nemo_relay/nemo_relay.go @@ -229,6 +229,9 @@ extern void nemo_relay_string_free(char* ptr); // Scope stack isolation extern int32_t nemo_relay_scope_stack_create(FfiScopeStack** out); +extern int32_t nemo_relay_capture_propagation_context_json(char** out); +extern int32_t nemo_relay_capture_propagation_context_with_root_json(const char* root_uuid, char** out); +extern int32_t nemo_relay_scope_stack_create_from_propagation_json(const char* context_json, FfiScopeStack** out); extern int32_t nemo_relay_scope_stack_set_thread(const FfiScopeStack* stack); extern int32_t nemo_relay_scope_stack_capture_thread(FfiThreadScopeStackBinding** out); extern int32_t nemo_relay_scope_stack_restore_thread(FfiThreadScopeStackBinding* binding); @@ -1567,12 +1570,68 @@ type ScopeStack struct { ptr *C.FfiScopeStack } +// PropagationContext is the versioned, transport-neutral causal context used +// to continue Relay work in another process. +type PropagationContext struct { + Version uint16 `json:"version"` + RootUUID *string `json:"root_uuid,omitempty"` + ParentUUID string `json:"parent_uuid"` +} + +// CapturePropagationContext captures the current Relay causal parent. +func CapturePropagationContext() (PropagationContext, error) { + var out *C.char + if err := checkStatus(C.nemo_relay_capture_propagation_context_json(&out)); err != nil { + return PropagationContext{}, err + } + defer C.nemo_relay_string_free(out) + var context PropagationContext + if err := json.Unmarshal([]byte(C.GoString(out)), &context); err != nil { + return PropagationContext{}, err + } + return context, nil +} + +// CapturePropagationContextWithRoot captures the current parent with an +// application-supplied stable session root. Pass nil when no root is known. +func CapturePropagationContextWithRoot(rootUUID *string) (PropagationContext, error) { + var cRoot *C.char + if rootUUID != nil { + cRoot = C.CString(*rootUUID) + defer C.free(unsafe.Pointer(cRoot)) + } + var out *C.char + if err := checkStatus(C.nemo_relay_capture_propagation_context_with_root_json(cRoot, &out)); err != nil { + return PropagationContext{}, err + } + defer C.nemo_relay_string_free(out) + var context PropagationContext + if err := json.Unmarshal([]byte(C.GoString(out)), &context); err != nil { + return PropagationContext{}, err + } + return context, nil +} + // NewScopeStack creates a new isolated scope stack. // The caller must call Close() when done. func NewScopeStack() (*ScopeStack, error) { return newScopeStackFunc() } +// NewScopeStackFromPropagation creates an isolated stack seeded from a +// received propagation context. The caller must call Close when done. +func NewScopeStackFromPropagation(context PropagationContext) (*ScopeStack, error) { + payload, err := json.Marshal(context) + if err != nil { + return nil, err + } + cPayload := C.CString(string(payload)) + defer C.free(unsafe.Pointer(cPayload)) + var ptr *C.FfiScopeStack + status := C.nemo_relay_scope_stack_create_from_propagation_json(cPayload, &ptr) + return checkedValue(int32(status), &ScopeStack{ptr: ptr}) +} + // Close frees the scope stack. After calling Close, the ScopeStack must not be used. func (s *ScopeStack) Close() { if s.ptr != nil { diff --git a/python/nemo_relay/__init__.py b/python/nemo_relay/__init__.py index fb8c89152..c3071e023 100644 --- a/python/nemo_relay/__init__.py +++ b/python/nemo_relay/__init__.py @@ -79,6 +79,7 @@ async def main(): import contextvars import typing from collections.abc import Callable as AbcCallable +from contextlib import contextmanager from typing import AsyncIterator, Awaitable, Callable, Literal, Optional, TypeAlias, TypedDict # Native bitflag classes exported at the top level for user code. @@ -110,6 +111,7 @@ async def main(): OpenTelemetryConfig, OpenTelemetrySubscriber, PendingMarkSpec, + PropagationContext, ScopeAttributes, ScopeEvent, ScopeHandle, @@ -119,7 +121,16 @@ async def main(): ToolExecutionInterceptOutcome, ToolHandle, ) +from nemo_relay._native import ( + capture_propagation_context as _capture_propagation_context, +) +from nemo_relay._native import ( + capture_propagation_context_with_root as _capture_propagation_context_with_root, +) from nemo_relay._native import create_scope_stack as _create_scope_stack +from nemo_relay._native import ( + create_scope_stack_from_propagation as _create_scope_stack_from_propagation, +) from nemo_relay._native import scope_stack_active as _native_scope_stack_active from nemo_relay._native import set_thread_scope_stack as _set_thread_scope_stack from nemo_relay._native import sync_thread_scope_stack as _sync_thread_scope_stack @@ -393,6 +404,35 @@ def create_scope_stack() -> ScopeStack: return _create_scope_stack() +def capture_propagation_context() -> PropagationContext: + """Capture the current Relay causal parent for application-managed transport.""" + get_scope_stack() + return _capture_propagation_context() + + +def capture_propagation_context_with_root(root_uuid: str | None) -> PropagationContext: + """Capture the current parent with an optional stable application session root.""" + get_scope_stack() + return _capture_propagation_context_with_root(root_uuid) + + +def create_scope_stack_from_propagation(context: PropagationContext) -> ScopeStack: + """Create an isolated stack seeded from a received propagation context.""" + return _create_scope_stack_from_propagation(context) + + +@contextmanager +def use_scope_stack(stack: ScopeStack): + """Temporarily install ``stack`` in the current Python context.""" + token = _scope_stack_var.set(stack) + _sync_thread_scope_stack(stack) + try: + yield stack + finally: + _scope_stack_var.reset(token) + _sync_thread_scope_stack(get_scope_stack()) + + def set_thread_scope_stack(stack: ScopeStack) -> None: """Install a scope stack into the current thread's native runtime context. @@ -461,11 +501,16 @@ def worker() -> None: "model_pricing", # Scope stack isolation "ScopeStack", + "PropagationContext", "create_scope_stack", + "capture_propagation_context", + "capture_propagation_context_with_root", + "create_scope_stack_from_propagation", "get_scope_stack", "scope_stack_active", "propagate_scope_to_thread", "set_thread_scope_stack", + "use_scope_stack", # Types "ScopeAttributes", "ToolAttributes", diff --git a/python/nemo_relay/__init__.pyi b/python/nemo_relay/__init__.pyi index 7c2372a57..951c58592 100644 --- a/python/nemo_relay/__init__.pyi +++ b/python/nemo_relay/__init__.pyi @@ -105,6 +105,9 @@ from nemo_relay._native import ( from nemo_relay._native import ( PendingMarkSpec as PendingMarkSpec, ) +from nemo_relay._native import ( + PropagationContext as PropagationContext, +) from nemo_relay._native import ( ScopeAttributes as ScopeAttributes, ) @@ -367,6 +370,10 @@ def create_scope_stack() -> ScopeStack: """ ... +def capture_propagation_context() -> PropagationContext: ... +def capture_propagation_context_with_root(root_uuid: str | None) -> PropagationContext: ... +def create_scope_stack_from_propagation(context: PropagationContext) -> ScopeStack: ... +def use_scope_stack(stack: ScopeStack): ... def set_thread_scope_stack(stack: ScopeStack) -> None: """Install a scope stack into the current thread's native runtime context. diff --git a/python/nemo_relay/_native.pyi b/python/nemo_relay/_native.pyi index a76ee5240..93f1628fd 100644 --- a/python/nemo_relay/_native.pyi +++ b/python/nemo_relay/_native.pyi @@ -1325,6 +1325,19 @@ def create_scope_stack() -> ScopeStack: """ ... +class PropagationContext: + """Transport-neutral Relay causal context.""" + def __init__(self, parent_uuid: str, root_uuid: str | None = None, version: int = 1) -> None: ... + @property + def version(self) -> int: ... + @property + def root_uuid(self) -> str | None: ... + @property + def parent_uuid(self) -> str: ... + +def capture_propagation_context() -> PropagationContext: ... +def capture_propagation_context_with_root(root_uuid: str | None) -> PropagationContext: ... +def create_scope_stack_from_propagation(context: PropagationContext) -> ScopeStack: ... def set_thread_scope_stack(stack: ScopeStack) -> None: """Install a scope stack into native thread-local storage. diff --git a/python/tests/test_context_isolation.py b/python/tests/test_context_isolation.py index 0c4146a00..252e69256 100644 --- a/python/tests/test_context_isolation.py +++ b/python/tests/test_context_isolation.py @@ -4,6 +4,7 @@ """Tests for per-request scope stack isolation via ContextVar.""" import asyncio +import uuid import nemo_relay @@ -15,6 +16,20 @@ def test_create_scope_stack_returns_scope_stack(): assert repr(stack) == "" +def test_propagation_context_installs_and_restores_a_scoped_stack(): + original = nemo_relay.get_scope_stack() + root_uuid = str(uuid.uuid4()) + parent_uuid = str(uuid.uuid4()) + context = nemo_relay.PropagationContext(parent_uuid, root_uuid) + stack = nemo_relay.create_scope_stack_from_propagation(context) + + with nemo_relay.use_scope_stack(stack): + assert nemo_relay.get_scope_stack() is stack + assert nemo_relay.scope.get_handle().uuid == parent_uuid + + assert nemo_relay.get_scope_stack() is original + + def test_get_scope_stack_returns_same_in_same_context(): """get_scope_stack returns the same instance within the same context.""" s1 = nemo_relay.get_scope_stack()