fix: snapshot ATOF header environment at activation#463
Conversation
Signed-off-by: Will Killian <wkillian@nvidia.com>
WalkthroughATOF endpoints now validate and resolve headers during worker startup. Activated endpoint state is reused across HTTP POST, WebSocket, and NDJSON transports, with Rust, component, and Python tests covering header environment snapshotting. ChangesATOF endpoint activation
Estimated code review effort: 3 (Moderate) | ~25 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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/src/observability/atof.rs`:
- Around line 647-653: Update ActivatedAtofEndpoint by removing its derived
Debug implementation, or replace it with a manual redacted Debug implementation
that never exposes header values such as Authorization. Keep the existing config
and headers fields unchanged.
In `@crates/core/tests/unit/observability/atof_tests.rs`:
- Around line 1466-1490: Extend the #[cfg(feature = "atof-streaming")]
observability tests beyond endpoint.headers by removing the environment variable
after activation, then exercising actual WebSocket handshake/reconnect and
NDJSON delivery requests. Assert each request retains the expected Authorization
header across the initial connection, reconnect, and upload, using request-level
test hooks or mock servers and preserving the existing test mutex isolation.
🪄 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: 72844c33-ec52-4bbf-ace3-a04b3991a7e9
📒 Files selected for processing (4)
crates/core/src/observability/atof.rscrates/core/tests/unit/observability/atof_tests.rscrates/core/tests/unit/observability/plugin_component_tests.rspython/tests/test_observability_plugin.py
📜 Review details
⏰ Context from checks skipped due to timeout. (2)
- GitHub Check: Check / Run
- GitHub Check: Preview docs
🧰 Additional context used
📓 Path-based instructions (17)
**/*.{rs,py}
📄 CodeRabbit inference engine (AGENTS.md)
Follow binding naming conventions in Rust and Python: use
snake_case.
Files:
python/tests/test_observability_plugin.pycrates/core/tests/unit/observability/plugin_component_tests.rscrates/core/src/observability/atof.rscrates/core/tests/unit/observability/atof_tests.rs
**/*.{rs,py,js,mjs,cjs,ts,tsx}
📄 CodeRabbit inference engine (AGENTS.md)
**/*.{rs,py,js,mjs,cjs,ts,tsx}: UseJson = serde_json::Valuein Rust-facing runtime APIs where the existing code expects JSON payloads.
UseResult<T>withFlowErrorin core runtime paths, and keep errors explicit and binding-appropriate at the wrapper layer.
Keep async behavior on the existing tokio-based model; bindings should preserve callback and future lifetimes rather than blocking or hiding async work unexpectedly.
Files:
python/tests/test_observability_plugin.pycrates/core/tests/unit/observability/plugin_component_tests.rscrates/core/src/observability/atof.rscrates/core/tests/unit/observability/atof_tests.rs
**/*.py
📄 CodeRabbit inference engine (CONTRIBUTING.md)
**/*.py: When changing the Python wrapper package, tests, or docs tooling, lint with Ruff (E,F,W,I), format with Ruff formatter (120-character lines, double quotes), and passtytype checking.
Add the SPDX license header to all Python source files using the#comment form.
Files:
python/tests/test_observability_plugin.py
**/*.{rs,py,go,js,ts,c,h}
📄 CodeRabbit inference engine (CONTRIBUTING.md)
Use language-appropriate naming conventions: Rust
snake_case, C FFI exports prefixednemo_relay_, GoPascalCase, Node.jscamelCase, and Pythonsnake_case.
Files:
python/tests/test_observability_plugin.pycrates/core/tests/unit/observability/plugin_component_tests.rscrates/core/src/observability/atof.rscrates/core/tests/unit/observability/atof_tests.rs
{crates/**/src/**/*.rs,python/**/*.py}
📄 CodeRabbit inference engine (.agents/skills/maintain-dynamic-plugins/SKILL.md)
Do not add tests under
src; Rust tests belong in cratetests/trees, and Python SDK tests belong underpython/tests.
Files:
python/tests/test_observability_plugin.pycrates/core/src/observability/atof.rs
**/*.{py,go,js,ts}
📄 CodeRabbit inference engine (.agents/skills/maintain-observability/SKILL.md)
Keep Python, Go, and Node.js config objects and subscriber/exporter methods aligned so all bindings expose the same logical knobs and semantics.
Files:
python/tests/test_observability_plugin.py
python/tests/**/*.py
📄 CodeRabbit inference engine (.agents/skills/test-python-binding/SKILL.md)
python/tests/**/*.py: Pytest is used to run tests.
Do not add@pytest.mark.asyncioto any test; async tests are automatically detected and run by the async runner.
Do not add a-> Nonereturn type annotation to test functions.
When mocking a class, do not define a new class; useunittest.mock.MagicMockorunittest.mock.AsyncMock, with thespecconstructor argument when necessary.
Name mocked classes with themockprefix, notfake.
Prefer pytest fixtures over helper methods.
Do not repeat fixtures; if a fixture is needed in multiple test files, place it in aconftest.pyfile.
When creating a fixture, use@pytest.fixture(name="<fixture_name>"[, scope="<scope>"])and define the fixture function asdef <fixture_name>_fixture() -> <return_type>:; only specifyscopewhen it is notfunction.
Preferpytest.mark.parametrizeover creating individual tests for different input types.
Files:
python/tests/test_observability_plugin.py
**/*
📄 CodeRabbit inference engine (.agents/skills/validate-change/SKILL.md)
**/*: Format changed files with the language-native formatter before the final lint/test pass.
If dynamic plugin behavior changed, usemaintain-dynamic-pluginsand include the native SDK, worker protocol, Python SDK, docs, packaging, and Codecov surfaces in the validation plan.
If code changes alter APIs, bindings, commands, paths, packaging behavior, observability/adaptive semantics, or documented best practices, update any dependent maintainer or consumer skills in the same branch.
During iteration, preferuv run pre-commit run --files <changed files...>.
Before review or handoff, runuv run pre-commit run --all-files.
Files:
python/tests/test_observability_plugin.pycrates/core/tests/unit/observability/plugin_component_tests.rscrates/core/src/observability/atof.rscrates/core/tests/unit/observability/atof_tests.rs
**/*.{rs,py,go,js,ts}
📄 CodeRabbit inference engine (.agents/skills/validate-change/SKILL.md)
If a language surface changed, always run that language's test target even when Rust core did not change.
Files:
python/tests/test_observability_plugin.pycrates/core/tests/unit/observability/plugin_component_tests.rscrates/core/src/observability/atof.rscrates/core/tests/unit/observability/atof_tests.rs
**/*.{rs,py,js,ts,tsx,go,java,kt,swift}
📄 CodeRabbit inference engine (.agents/skills/add-middleware/SKILL.md)
Add tests covering registration and duplicate names, deregistration and missing names, priority ordering, callback failure policy, scope-local inheritance and cleanup, event payload semantics, immutable mark and scope fields, and parity across affected bindings.
Files:
python/tests/test_observability_plugin.pycrates/core/tests/unit/observability/plugin_component_tests.rscrates/core/src/observability/atof.rscrates/core/tests/unit/observability/atof_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:
python/tests/test_observability_plugin.pycrates/core/tests/unit/observability/plugin_component_tests.rscrates/core/tests/unit/observability/atof_tests.rs
**/*.rs
📄 CodeRabbit inference engine (.agents/skills/prepare-pr/SKILL.md)
**/*.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 workWhen changing the core Rust runtime or Rust-facing API surface, format Rust code with
cargo fmt(rustfmt defaults), keepcargo clippy -- -D warningsclean, and satisfycargo deny checkperdeny.toml.
**/*.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.
For Rust changes headed for review, runcargo fmt --allandcargo clippy --workspace --all-targets -- -D warningseven if relying on pre-commit.
Files:
crates/core/tests/unit/observability/plugin_component_tests.rscrates/core/src/observability/atof.rscrates/core/tests/unit/observability/atof_tests.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/tests/unit/observability/plugin_component_tests.rscrates/core/src/observability/atof.rscrates/core/tests/unit/observability/atof_tests.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
Files:
crates/core/tests/unit/observability/plugin_component_tests.rscrates/core/src/observability/atof.rscrates/core/tests/unit/observability/atof_tests.rs
**/*.{rs,go,js,ts}
📄 CodeRabbit inference engine (CONTRIBUTING.md)
Add the SPDX license header to all Rust, Go, JavaScript, and TypeScript source files using the corresponding
//comment form.
Files:
crates/core/tests/unit/observability/plugin_component_tests.rscrates/core/src/observability/atof.rscrates/core/tests/unit/observability/atof_tests.rs
crates/{core,adaptive}/**/*
📄 CodeRabbit inference engine (.agents/skills/validate-change/SKILL.md)
If
crates/coreorcrates/adaptivechanged, run the full validation matrix across Rust, Python, Go, and Node.js.
Files:
crates/core/tests/unit/observability/plugin_component_tests.rscrates/core/src/observability/atof.rscrates/core/tests/unit/observability/atof_tests.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/tests/unit/observability/plugin_component_tests.rscrates/core/src/observability/atof.rscrates/core/tests/unit/observability/atof_tests.rs
🪛 ast-grep (0.44.1)
python/tests/test_observability_plugin.py
[info] 224-224: use jsonify instead of json.dumps for JSON output
Context: json.dumps(report)
Note: [CWE-116] Improper Encoding or Escaping of Output.
(use-jsonify)
🪛 Ruff (0.15.21)
python/tests/test_observability_plugin.py
[error] 49-49: Function argument format is shadowing a Python builtin
(A002)
[warning] 54-54: Remove quotes from type annotation
Remove quotes
(UP037)
🔇 Additional comments (4)
crates/core/src/observability/atof.rs (1)
721-728: LGTM!Also applies to: 756-784, 825-930, 933-1147, 1150-1209
crates/core/tests/unit/observability/atof_tests.rs (1)
1378-1464: LGTM!Also applies to: 1564-1568, 1600-1607, 1626-1645
crates/core/tests/unit/observability/plugin_component_tests.rs (1)
40-84: LGTM!Also applies to: 144-152, 1122-1181
python/tests/test_observability_plugin.py (1)
8-11: LGTM!Also applies to: 35-81, 186-226
mnajafian-nv
left a comment
There was a problem hiding this comment.
LGTM pending the Windows Rust check going green.
|
@mnajafian-nv the rust test failure:
|
|
/merge |
Overview
Resolve RELAY-499 by snapshotting ATOF stream-sink headers sourced from
header_envwhen the endpoint is activated.Details
headersandheader_envonce before the endpoint worker is spawned.Where should the reviewer start?
Start with
crates/core/src/observability/atof.rs, whereActivatedAtofEndpointcarries the validated configuration and snapshotted headers into each transport worker.Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)
Summary by CodeRabbit