Skip to content

fix: snapshot ATOF header environment at activation#463

Merged
rapids-bot[bot] merged 1 commit into
NVIDIA:release/0.6from
willkill07:wkk_fix/relay-499-header-env
Jul 17, 2026
Merged

fix: snapshot ATOF header environment at activation#463
rapids-bot[bot] merged 1 commit into
NVIDIA:release/0.6from
willkill07:wkk_fix/relay-499-header-env

Conversation

@willkill07

@willkill07 willkill07 commented Jul 17, 2026

Copy link
Copy Markdown
Member

Overview

Resolve RELAY-499 by snapshotting ATOF stream-sink headers sourced from header_env when the endpoint is activated.

  • I confirm this contribution is my own work, or I have the right to submit it under this project's license.
  • I searched existing issues and open pull requests, and this does not duplicate existing work.

Details

  • Resolve headers and header_env once before the endpoint worker is spawned.
  • Reuse the activated header map for HTTP POST, WebSocket reconnects, and NDJSON uploads.
  • Add core and Python regressions that remove the environment variable after activation and verify all lifecycle events retain the Authorization header.

Where should the reviewer start?

Start with crates/core/src/observability/atof.rs, where ActivatedAtofEndpoint carries the validated configuration and snapshotted headers into each transport worker.

Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)

  • Relates to: RELAY-499

Summary by CodeRabbit

  • Bug Fixes
    • Improved streaming observability endpoint setup by validating configurations and resolving headers before workers start.
    • Ensured environment-based headers are captured at activation time and remain available for outgoing requests.
    • Applied consistent header handling across HTTP, WebSocket, and NDJSON transports.
  • Tests
    • Added coverage verifying environment-based authorization headers are preserved without exposing credentials in reports, payloads, or logs.

Signed-off-by: Will Killian <wkillian@nvidia.com>
@willkill07
willkill07 requested a review from a team as a code owner July 17, 2026 14:19
@github-actions github-actions Bot added size:M PR is medium Bug issue describes bug; PR fixes bug lang:python PR changes/introduces Python code lang:rust PR changes/introduces Rust code labels Jul 17, 2026
@coderabbitai

coderabbitai Bot commented Jul 17, 2026

Copy link
Copy Markdown

Review Change Stack

Walkthrough

ATOF 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.

Changes

ATOF endpoint activation

Layer / File(s) Summary
Activate endpoints and reuse resolved headers
crates/core/src/observability/atof.rs
Validation now returns activated endpoint state containing resolved headers, and HTTP, WebSocket, and NDJSON workers use that state for request construction and reconnects.
Update Rust activation and worker tests
crates/core/tests/unit/observability/atof_tests.rs
Tests pass activated endpoints into workers and helpers and verify validation, reconnect behavior, and header environment snapshotting.
Capture and verify snapshotted headers
crates/core/tests/unit/observability/plugin_component_tests.rs, python/tests/test_observability_plugin.py
Test capture servers record headers and bodies, and integration tests verify resolved authorization headers remain on emitted requests after environment cleanup.

Estimated code review effort: 3 (Moderate) | ~25 minutes

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Title check ✅ Passed The title follows Conventional Commits and accurately summarizes the header snapshotting fix at activation.
Description check ✅ Passed The description matches the template with Overview, Details, reviewer start, checklist items, and a related issue reference.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Comment @coderabbitai help to get the list of available commands.

@willkill07 willkill07 added this to the 0.6 milestone Jul 17, 2026
@willkill07 willkill07 self-assigned this Jul 17, 2026
@github-actions

Copy link
Copy Markdown

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

📥 Commits

Reviewing files that changed from the base of the PR and between 0325ba2 and 950bb89.

📒 Files selected for processing (4)
  • crates/core/src/observability/atof.rs
  • crates/core/tests/unit/observability/atof_tests.rs
  • crates/core/tests/unit/observability/plugin_component_tests.rs
  • python/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.py
  • crates/core/tests/unit/observability/plugin_component_tests.rs
  • crates/core/src/observability/atof.rs
  • crates/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}: Use Json = serde_json::Value in Rust-facing runtime APIs where the existing code expects JSON payloads.
Use Result<T> with FlowError in 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.py
  • crates/core/tests/unit/observability/plugin_component_tests.rs
  • crates/core/src/observability/atof.rs
  • crates/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 pass ty type 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 prefixed nemo_relay_, Go PascalCase, Node.js camelCase, and Python snake_case.

Files:

  • python/tests/test_observability_plugin.py
  • crates/core/tests/unit/observability/plugin_component_tests.rs
  • crates/core/src/observability/atof.rs
  • crates/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 crate tests/ trees, and Python SDK tests belong under python/tests.

Files:

  • python/tests/test_observability_plugin.py
  • crates/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.asyncio to any test; async tests are automatically detected and run by the async runner.
Do not add a -> None return type annotation to test functions.
When mocking a class, do not define a new class; use unittest.mock.MagicMock or unittest.mock.AsyncMock, with the spec constructor argument when necessary.
Name mocked classes with the mock prefix, not fake.
Prefer pytest fixtures over helper methods.
Do not repeat fixtures; if a fixture is needed in multiple test files, place it in a conftest.py file.
When creating a fixture, use @pytest.fixture(name="<fixture_name>"[, scope="<scope>"]) and define the fixture function as def <fixture_name>_fixture() -> <return_type>:; only specify scope when it is not function.
Prefer pytest.mark.parametrize over 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, use maintain-dynamic-plugins and 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, prefer uv run pre-commit run --files <changed files...>.
Before review or handoff, run uv run pre-commit run --all-files.

Files:

  • python/tests/test_observability_plugin.py
  • crates/core/tests/unit/observability/plugin_component_tests.rs
  • crates/core/src/observability/atof.rs
  • crates/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.py
  • crates/core/tests/unit/observability/plugin_component_tests.rs
  • crates/core/src/observability/atof.rs
  • crates/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.py
  • crates/core/tests/unit/observability/plugin_component_tests.rs
  • crates/core/src/observability/atof.rs
  • crates/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.py
  • crates/core/tests/unit/observability/plugin_component_tests.rs
  • crates/core/tests/unit/observability/atof_tests.rs
**/*.rs

📄 CodeRabbit inference engine (.agents/skills/prepare-pr/SKILL.md)

**/*.rs: Any Rust change must run just test-rust
Any Rust change must run cargo fmt --all
Any Rust change must run cargo clippy --workspace --all-targets -- -D warnings

**/*.rs: Run cargo fmt --all for all FFI work since it is Rust work
Run just test-rust to validate FFI changes
Run cargo clippy --workspace --all-targets -- -D warnings to enforce strict linting on FFI work

When Rust files changed as part of Go work, also run cargo fmt --all, just test-rust, and cargo clippy --workspace --all-targets -- -D warnings

**/*.rs: Run cargo fmt --all when Rust files are changed as part of Node work
Run cargo clippy --workspace --all-targets -- -D warnings when Rust files are changed as part of Node work
Run just test-rust when Rust files are changed as part of Node work

When changing the core Rust runtime or Rust-facing API surface, format Rust code with cargo fmt (rustfmt defaults), keep cargo clippy -- -D warnings clean, and satisfy cargo deny check per deny.toml.

**/*.rs: If any Rust code changed, always run just test-rust.
If any Rust code changed, also run cargo fmt --all.
If any Rust code changed, also run cargo clippy --workspace --all-targets -- -D warnings.
For Rust changes headed for review, run cargo fmt --all and cargo clippy --workspace --all-targets -- -D warnings even if relying on pre-commit.

Files:

  • crates/core/tests/unit/observability/plugin_component_tests.rs
  • crates/core/src/observability/atof.rs
  • crates/core/tests/unit/observability/atof_tests.rs
{crates/core,crates/adaptive}/**/*

📄 CodeRabbit inference engine (.agents/skills/prepare-pr/SKILL.md)

Changes to crates/core or crates/adaptive must run the full language matrix

Files:

  • crates/core/tests/unit/observability/plugin_component_tests.rs
  • crates/core/src/observability/atof.rs
  • crates/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/core or shared runtime semantics, also use validate-change for broader validation

Files:

  • crates/core/tests/unit/observability/plugin_component_tests.rs
  • crates/core/src/observability/atof.rs
  • crates/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.rs
  • crates/core/src/observability/atof.rs
  • crates/core/tests/unit/observability/atof_tests.rs
crates/{core,adaptive}/**/*

📄 CodeRabbit inference engine (.agents/skills/validate-change/SKILL.md)

If crates/core or crates/adaptive changed, run the full validation matrix across Rust, Python, Go, and Node.js.

Files:

  • crates/core/tests/unit/observability/plugin_component_tests.rs
  • crates/core/src/observability/atof.rs
  • crates/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.rs
  • crates/core/src/observability/atof.rs
  • crates/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

Comment thread crates/core/src/observability/atof.rs
Comment thread crates/core/tests/unit/observability/atof_tests.rs

@mnajafian-nv mnajafian-nv left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM pending the Windows Rust check going green.

@willkill07

Copy link
Copy Markdown
Member Author

@mnajafian-nv the rust test failure:

  • it's non-blocking and unrelated
  • i have another PR open that fixes it!

@willkill07

Copy link
Copy Markdown
Member Author

/merge

@rapids-bot
rapids-bot Bot merged commit ea22a8e into NVIDIA:release/0.6 Jul 17, 2026
59 of 60 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Bug issue describes bug; PR fixes bug lang:python PR changes/introduces Python code lang:rust PR changes/introduces Rust code size:M PR is medium

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants