fix: resolve sanitized LLM request from annotations#241
Conversation
Signed-off-by: Minh Vu <vuhoangminh97@gmail.com>
github-actions
Bot
added
size:M
WalkthroughThis PR extends LLM lifecycle event emission to derive codec annotations from sanitized request and response payloads, introduces configurable end-event behavior for response codec handling, and wires request codecs through managed execution paths to enable sanitization-aware start events. ChangesCodec Annotations from Sanitized Payloads
Sequence DiagramsequenceDiagram
participant Client
participant LLMExecution as LLM Execution
participant RequestSanitize as Request Sanitization
participant EmitStart as emit_llm_start
participant ResponseSanitize as Response Sanitization
participant EmitEnd as emit_llm_end
participant Events as LlmStart/LlmEnd Events
Client->>RequestSanitize: Original Request
RequestSanitize->>EmitStart: Sanitized Request + request_codec
EmitStart->>EmitStart: Decode Sanitized Request with codec
EmitStart->>Events: LlmStart(annotated_request from Sanitized)
Events->>LLMExecution: Execute with Sanitized
LLMExecution->>ResponseSanitize: Response
ResponseSanitize->>EmitEnd: Sanitized Response + response_codec
EmitEnd->>EmitEnd: Decode Sanitized Response with codec
EmitEnd->>EmitEnd: Attach Estimated Cost (if enabled)
EmitEnd->>Events: LlmEnd(annotated_response from Sanitized)
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
crates/core/src/api/llm.rs (1)
433-455:⚠️ Potential issue | 🟠 Major | ⚡ Quick winDo not decode the fallback
datapayload as an LLM response.Lines 434-455 collapse two different payloads into
data: the sanitized provider response and the caller-supplied compatibility payload. When a sanitizer returnsnull,response_codecwill decode that fallback payload instead of an LLM response. Inllm_call_execute, that means arbitrary application data can produce bogusannotated_responsecontent, and in manualllm_call_endit can turn a full-redaction sanitizer into a fatal decode error. Keep a separate decode source and only rundecode_responsewhen the sanitized response itself is non-null.Suggested fix
- let sanitized_response = state.llm_sanitize_response_chain(response, &scope_locals); - let data = if sanitized_response.is_null() { - data - } else { - Some(sanitized_response) - }; + let sanitized_response = state.llm_sanitize_response_chain(response, &scope_locals); + let codec_response = if sanitized_response.is_null() { + None + } else { + Some(sanitized_response) + }; + let data = codec_response.clone().or(data); let annotated_response = match annotated_response { Some(annotated_response) => Some(annotated_response), - None => match (response_codec.as_ref(), data.as_ref()) { + None => match (response_codec.as_ref(), codec_response.as_ref()) { (Some(codec), Some(response)) => match codec.decode_response(response) {🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/core/src/api/llm.rs` around lines 433 - 455, The code currently collapses the sanitized provider response and the caller-supplied fallback into one `data` variable and then decodes `data`, which lets non-LLM fallback payloads be treated as LLM responses; instead, keep the sanitized response value separate (the return of state.llm_sanitize_response_chain) and only call response_codec.decode_response(...) when that sanitized response is non-null. Concretely: do not reuse `data` as the decode source in the annotated_response construction—use the sanitized_response directly for decoding and only fall back to using the caller-supplied compatibility payload as the final `data` value (but never pass that fallback into response_codec). Keep the existing cost attachment behavior (attach_estimated_cost_for_provider(&mut decoded, Some(&handle.name))) when decoding succeeds and still surface decode errors into `decode_error`, but ensure decode_response is invoked exclusively on the non-null sanitized response (affecting llm_call_execute and llm_call_end code paths).
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@crates/core/tests/integration/pipeline_tests.rs`:
- Around line 22-26: Run the repository-required Rust validation set and fix any
issues: execute "just test-rust", then "cargo fmt --all" and commit any
formatting changes, run "cargo clippy --workspace --all-targets -- -D warnings"
and fix all lints (ensure tests/imports referencing symbols like
register_llm_sanitize_request_guardrail, deregister_llm_request_intercept,
register_llm_request_intercept, register_llm_sanitize_response_guardrail,
deregister_llm_sanitize_response_guardrail compile clean), and finally run "uv
run pre-commit run --all-files"; after resolving failures, update the PR with
the command outputs and any code fixes and new commits.
---
Outside diff comments:
In `@crates/core/src/api/llm.rs`:
- Around line 433-455: The code currently collapses the sanitized provider
response and the caller-supplied fallback into one `data` variable and then
decodes `data`, which lets non-LLM fallback payloads be treated as LLM
responses; instead, keep the sanitized response value separate (the return of
state.llm_sanitize_response_chain) and only call
response_codec.decode_response(...) when that sanitized response is non-null.
Concretely: do not reuse `data` as the decode source in the annotated_response
construction—use the sanitized_response directly for decoding and only fall back
to using the caller-supplied compatibility payload as the final `data` value
(but never pass that fallback into response_codec). Keep the existing cost
attachment behavior (attach_estimated_cost_for_provider(&mut decoded,
Some(&handle.name))) when decoding succeeds and still surface decode errors into
`decode_error`, but ensure decode_response is invoked exclusively on the
non-null sanitized response (affecting llm_call_execute and llm_call_end code
paths).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Enterprise
Run ID: 2cd7f213-11fe-4c7f-aa4a-49dd14a8b0a5
📒 Files selected for processing (3)
crates/core/src/api/llm.rscrates/core/src/stream.rscrates/core/tests/integration/pipeline_tests.rs
📜 Review details
🧰 Additional context used
📓 Path-based instructions (16)
**/*.rs
📄 CodeRabbit inference engine (.agents/skills/add-binding-feature/SKILL.md)
Use
snake_casenaming convention for Rust identifiers (e.g.,nemo_relay_tool_call)
**/*.rs: Any Rust change must runjust test-rust
Any Rust change must runcargo fmt --all
Any Rust change must runcargo clippy --workspace --all-targets -- -D warnings
**/*.rs: Runcargo fmt --allfor all FFI work since it is Rust work
Runjust test-rustto validate FFI changes
Runcargo clippy --workspace --all-targets -- -D warningsto enforce strict linting on FFI workWhen Rust files changed as part of Go work, also run
cargo fmt --all,just test-rust, andcargo clippy --workspace --all-targets -- -D warnings
**/*.rs: Runcargo fmt --allwhen Rust files are changed as part of Node work
Runcargo clippy --workspace --all-targets -- -D warningswhen Rust files are changed as part of Node work
Runjust test-rustwhen Rust files are changed as part of Node work
**/*.rs: Runcargo fmt --allto format all Rust code
Runcargo clippy --workspace --all-targets -- -D warningsto enforce all clippy lints as errors
**/*.rs: Runcargo fmt --allwhen Rust files changed as part of WebAssembly work
Runcargo clippy --workspace --all-targets -- -D warningswhen Rust files changed as part of WebAssembly work
**/*.rs: If any Rust code changed, always runjust test-rust
If any Rust code changed, also runcargo fmt --all
If any Rust code changed, also runcargo clippy --workspace --all-targets -- -D warnings
Run Rust formatting withcargo fmt --all
Run Rust linting withcargo clippy --workspace --all-targets -- -D warnings
**/*.rs: Usecargo fmtfor Rust code formatting
Runcargo clippy -- -D warningsto lint Rust code and treat all warnings as errors
Use Rust snake_case naming convention for Rust identifiers
Include SPDX license header in all Rust source files using double-slash comment syntax
Validate Rust code withuv run pre-commit run --all-filesto enforce cargo fmt formatting check, cargo clippy lints, and cargo deny aud...
Files:
crates/core/src/stream.rscrates/core/tests/integration/pipeline_tests.rscrates/core/src/api/llm.rs
**/{Cargo.toml,**/*.rs}
📄 CodeRabbit inference engine (.agents/skills/maintain-packaging/SKILL.md)
Maintain consistency between Rust package names in
Cargo.tomland their actual usage across the codebase
Files:
crates/core/src/stream.rscrates/core/tests/integration/pipeline_tests.rscrates/core/src/api/llm.rs
**/*.{h,hpp,c,cpp,rs}
📄 CodeRabbit inference engine (.agents/skills/maintain-packaging/SKILL.md)
Ensure FFI header and library naming follows consistent conventions across platform-specific builds
Files:
crates/core/src/stream.rscrates/core/tests/integration/pipeline_tests.rscrates/core/src/api/llm.rs
{crates/core,crates/adaptive}/**/*
📄 CodeRabbit inference engine (.agents/skills/prepare-pr/SKILL.md)
Changes to
crates/coreorcrates/adaptivemust run the full language matrix
Files:
crates/core/src/stream.rscrates/core/tests/integration/pipeline_tests.rscrates/core/src/api/llm.rs
**/*.{rs,toml}
📄 CodeRabbit inference engine (.agents/skills/rename-surfaces/SKILL.md)
Update Rust crate names and module prefixes during coordinated rename operations
Files:
crates/core/src/stream.rscrates/core/tests/integration/pipeline_tests.rscrates/core/src/api/llm.rs
crates/core/**/*.rs
📄 CodeRabbit inference engine (.agents/skills/test-go-binding/SKILL.md)
If the change touched
crates/coreor shared runtime semantics, also usevalidate-changefor broader validation
crates/core/**/*.rs: UseJson = serde_json::Valuein Rust-facing runtime APIs where the existing code expects JSON payloads.
UseResult<T>withFlowErrorin core runtime paths. Keep errors explicit and binding-appropriate at the wrapper layer.
Files:
crates/core/src/stream.rscrates/core/tests/integration/pipeline_tests.rscrates/core/src/api/llm.rs
crates/{core,adaptive}/**
📄 CodeRabbit inference engine (.agents/skills/validate-change/SKILL.md)
If
crates/coreorcrates/adaptivechanged, run the full matrix across Rust, Python, Go, Node.js, and WebAssembly
Files:
crates/core/src/stream.rscrates/core/tests/integration/pipeline_tests.rscrates/core/src/api/llm.rs
**/*.{rs,py,js,ts,tsx,jsx,go,sh,toml,yaml,yml,md}
📄 CodeRabbit inference engine (AGENTS.md)
Keep SPDX headers on source, docs, scripts, and configuration files. The project is Apache-2.0.
Files:
crates/core/src/stream.rscrates/core/tests/integration/pipeline_tests.rscrates/core/src/api/llm.rs
**/*.{rs,py,go,js,ts,tsx}
📄 CodeRabbit inference engine (AGENTS.md)
Follow binding naming conventions: Rust and Python use
snake_case, C FFI exports prefixednemo_relay_, Go usesPascalCasefor public APIs, Node.js usescamelCase.
Files:
crates/core/src/stream.rscrates/core/tests/integration/pipeline_tests.rscrates/core/src/api/llm.rs
crates/**/*.rs
📄 CodeRabbit inference engine (AGENTS.md)
crates/**/*.rs: Keep async behavior on the existing tokio-based model. Bindings should preserve callback and future lifetimes rather than blocking or hiding async work unexpectedly.
UseJson = serde_json::Valuein Rust-facing runtime APIs for JSON payload handling.
Files:
crates/core/src/stream.rscrates/core/tests/integration/pipeline_tests.rscrates/core/src/api/llm.rs
crates/{core,adaptive}/**/*.rs
⚙️ CodeRabbit configuration file
crates/{core,adaptive}/**/*.rs: Review the Rust runtime for async correctness, scope isolation, middleware ordering, and event lifecycle regressions.
Pay close attention to task-local/thread-local scope propagation, callback lifetimes, stream finalization, and root_uuid isolation.
Public API changes should preserve existing behavior unless tests and docs show the intended migration path.
Files:
crates/core/src/stream.rscrates/core/tests/integration/pipeline_tests.rscrates/core/src/api/llm.rs
**
⚙️ CodeRabbit configuration file
**:AGENTS.md
This file provides guidance to agents, including Claude Code and OpenAI Codex, when working in this repository.
Project Overview
NeMo Relay is a multi-language agent runtime framework for execution scopes, lifecycle events, middleware, plugins, and observability around tool and LLM calls. The core runtime is Rust. Primary supported bindings are Rust, Python, and Node.js. Go, WebAssembly, and the raw C FFI are experimental and source-first.
The shared runtime model is:
- Scope stacks decide where work belongs and which scope-local behavior is visible.
- Middleware registries decide what guardrails and intercepts run around managed calls.
- Plugins install reusable runtime behavior from configuration.
- Events record runtime behavior in ATOF form.
- Subscribers and exporters consume events in-process or export them to ATIF, OpenTelemetry, OpenInference, or other backends.
Repository Structure
The repository layout separates the Rust runtime, language bindings, documentation,
integration patches, and agent-facing skills.crates/ core/ # Rust core runtime crate, published as nemo-relay adaptive/ # Adaptive runtime primitives and plugin components python/ # PyO3 native extension for the Python package ffi/ # Raw C ABI layer used by downstream bindings such as Go node/ # NAPI Node.js binding and JavaScript/TypeScript entry points wasm/ # wasm-bindgen WebAssembly binding and JS wrappers python/ nemo_relay/ # Python wrapper package: scopes, tools, LLM, middleware, typed helpers, plugins, adaptive helpers tests/ # Python tests go/ nemo_relay/ # Experimental Go CGo binding and tests fern/ # Fern documentation site scripts/ # Stable wrappers and helper scripts; build/test/docs entry points live in justfile third_party/ # P...
Files:
crates/core/src/stream.rscrates/core/tests/integration/pipeline_tests.rscrates/core/src/api/llm.rs
{crates/adaptive/**/*.rs,**/*test*.{rs,py,go,ts,js},**/*adaptive*test*.{rs,py,go,ts,js},docs/plugins/adaptive/**}
📄 CodeRabbit inference engine (.agents/skills/maintain-optimizer/SKILL.md)
Maintain documented and tested validation and report behavior for adaptive surfaces
Files:
crates/core/tests/integration/pipeline_tests.rs
{crates/**/tests/**,python/tests/**,go/nemo_relay/**/*_test.go}
⚙️ CodeRabbit configuration file
{crates/**/tests/**,python/tests/**,go/nemo_relay/**/*_test.go}: Tests should cover the behavior promised by the changed API surface, including error paths and cross-request isolation where relevant.
Prefer assertions on lifecycle events, scope stacks, middleware ordering, and binding parity over shallow smoke tests.
Files:
crates/core/tests/integration/pipeline_tests.rs
crates/core/src/api/**/*.rs
📄 CodeRabbit inference engine (.agents/skills/add-binding-feature/SKILL.md)
Implement behavior first in Rust core API modules:
crates/core/src/api/and related core modules such ascrates/core/src/api/runtime/,crates/core/src/codec/, orcrates/core/src/json.rs
Files:
crates/core/src/api/llm.rs
crates/core/src/api/{tool,llm}.rs
📄 CodeRabbit inference engine (.agents/skills/add-middleware/SKILL.md)
Wire the new middleware chain into the execute path in
crates/core/src/api/tool.rsorcrates/core/src/api/llm.rsat the appropriate pipeline stage
Files:
crates/core/src/api/llm.rs
🔇 Additional comments (3)
crates/core/src/stream.rs (1)
159-170: LGTM!crates/core/tests/integration/pipeline_tests.rs (1)
33-35: LGTM!Also applies to: 233-275, 1052-1110, 1252-1311, 1376-1438
crates/core/src/api/llm.rs (1)
593-640: Please attach the requiredcrates/corevalidation results before merge.The PR summary only shows crate-scoped Rust checks. For a
crates/coreruntime-semantics change, repository guidance is stricter: this should be validated withjust test-rust,cargo fmt --all,cargo clippy --workspace --all-targets -- -D warnings,validate-change, and the full Rust/Python/Go/Node/WebAssembly matrix. As per coding guidelines, Rust changes must runjust test-rust,cargo fmt --all, andcargo clippy --workspace --all-targets -- -D warnings, andcrates/corechanges must also “usevalidate-changefor broader validation” and “run the full matrix across Rust, Python, Go, Node.js, and WebAssembly.”Source: Coding guidelines
willkill07
changed the title
|
/ok to test 278e4af |
|
@fallintoplace thank you for your contribution! Once CI passes I will merge. |
|
/merge |
2 participants
Summary
Why
Sanitize guardrails changed the event payload, but codec annotations could still reflect the raw request or response. OpenInference prefers annotations for flattened LLM attributes, so managed calls could export raw content even when event data was sanitized.
Validation
cargo test -p nemo-relay --test pipeline_integration sanitized -- --nocapturecargo test -p nemo-relay --test pipeline_integration test_response_codec_failure_non_fatal -- --nocapturecargo test -p nemo-relay --test pipeline_integration -- --nocapturecargo test -p nemo-relay openinference -- --nocapturecargo test -p nemo-relaycargo fmt --all --checkcargo clippy -p nemo-relay --all-targets -- -D warningsSummary by CodeRabbit
Bug Fixes
Tests