fix: restore task-local Python scope stacks#462
Conversation
Signed-off-by: Will Killian <wkillian@nvidia.com>
github-actions
Bot
added
size:M
WalkthroughCentralizes scope-stack synchronization in a shared context helper, applies it across LLM and tool lifecycle APIs, routes scope cleanup through the Python pop wrapper, and adds concurrent task-isolation coverage for scope restoration and event ancestry. ChangesTask-local scope isolation
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant Task
participant LLMOrTools
participant ensure_scope_stack
participant NativeLifecycle
participant ScopeSubscribers
Task->>LLMOrTools: invoke lifecycle helper
LLMOrTools->>ensure_scope_stack: synchronize task scope
ensure_scope_stack->>NativeLifecycle: update native scope state
LLMOrTools->>NativeLifecycle: execute lifecycle operation
NativeLifecycle->>ScopeSubscribers: emit scope events
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ 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: 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 `@python/nemo_relay/_context.py`:
- Around line 20-28: The scope initialization logic around
_native_scope_stack_active must not treat an active thread-local stack as
ownership for a fresh asyncio task; detect an active task and initialize it
through get_scope_stack(), while preserving explicit worker-thread behavior.
Update tests to exercise this through the public API with initially unset
ContextVar values, avoiding direct _scope_stack_var manipulation.
In `@python/tests/test_context_isolation.py`:
- Around line 50-65: Update the test’s run_scope/main coordination to use
asyncio.Event barriers: ensure task B has activated its scope before task A
exits, rather than relying on asyncio.sleep(0) or asserting scheduler order.
Replace the completion-order assertion with an assertion that both tasks
complete, while preserving the scope setup and _scope_stack_var reset behavior.
🪄 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: 4f07593e-93c8-463e-beed-78e5ebcbaeb7
📒 Files selected for processing (5)
python/nemo_relay/_context.pypython/nemo_relay/llm.pypython/nemo_relay/scope.pypython/nemo_relay/tools.pypython/tests/test_context_isolation.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 (15)
**/*.{rs,py}
📄 CodeRabbit inference engine (AGENTS.md)
Follow binding naming conventions in Rust and Python: use
snake_case.
Files:
python/nemo_relay/_context.pypython/nemo_relay/tools.pypython/nemo_relay/llm.pypython/tests/test_context_isolation.pypython/nemo_relay/scope.py
**/*.{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/nemo_relay/_context.pypython/nemo_relay/tools.pypython/nemo_relay/llm.pypython/tests/test_context_isolation.pypython/nemo_relay/scope.py
python/nemo_relay/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Python wrapper modules live under
python/nemo_relay/, and the native extension is built fromcrates/pythonwithmaturin.
Files:
python/nemo_relay/_context.pypython/nemo_relay/tools.pypython/nemo_relay/llm.pypython/nemo_relay/scope.py
**/*.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/nemo_relay/_context.pypython/nemo_relay/tools.pypython/nemo_relay/llm.pypython/tests/test_context_isolation.pypython/nemo_relay/scope.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/nemo_relay/_context.pypython/nemo_relay/tools.pypython/nemo_relay/llm.pypython/tests/test_context_isolation.pypython/nemo_relay/scope.py
{crates/python/src/py_api/mod.rs,python/nemo_relay/**/*.py,python/nemo_relay/**/*.pyi,go/nemo_relay/**/*.go,crates/node/src/api/**/*.rs}
📄 CodeRabbit inference engine (.agents/skills/add-binding-feature/SKILL.md)
Update the language-native bindings for every exposed surface in Python, Go, and Node.js.
Files:
python/nemo_relay/_context.pypython/nemo_relay/tools.pypython/nemo_relay/llm.pypython/nemo_relay/scope.py
{python/nemo_relay/**/*.py,python/nemo_relay/**/*.pyi,go/nemo_relay/**/*.go}
📄 CodeRabbit inference engine (.agents/skills/add-binding-feature/SKILL.md)
Update language wrapper helpers such as Python wrapper modules, Python type stubs, and Go shorthand packages when the new behavior belongs in those helper layers.
Files:
python/nemo_relay/_context.pypython/nemo_relay/tools.pypython/nemo_relay/llm.pypython/nemo_relay/scope.py
{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/nemo_relay/_context.pypython/nemo_relay/tools.pypython/nemo_relay/llm.pypython/tests/test_context_isolation.pypython/nemo_relay/scope.py
**/*.{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/nemo_relay/_context.pypython/nemo_relay/tools.pypython/nemo_relay/llm.pypython/tests/test_context_isolation.pypython/nemo_relay/scope.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/nemo_relay/_context.pypython/nemo_relay/tools.pypython/nemo_relay/llm.pypython/tests/test_context_isolation.pypython/nemo_relay/scope.py
**/*.{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/nemo_relay/_context.pypython/nemo_relay/tools.pypython/nemo_relay/llm.pypython/tests/test_context_isolation.pypython/nemo_relay/scope.py
**/*.{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/nemo_relay/_context.pypython/nemo_relay/tools.pypython/nemo_relay/llm.pypython/tests/test_context_isolation.pypython/nemo_relay/scope.py
python/nemo_relay/**/*
⚙️ CodeRabbit configuration file
python/nemo_relay/**/*: Review Python wrapper changes for typed API consistency, contextvars-based scope isolation, async behavior, and parity with the native extension.
Stubs and runtime implementations should stay aligned.
Files:
python/nemo_relay/_context.pypython/nemo_relay/tools.pypython/nemo_relay/llm.pypython/nemo_relay/scope.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_context_isolation.py
{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_context_isolation.py
🪛 OpenGrep (1.25.0)
python/tests/test_context_isolation.py
[ERROR] 98-102: SQL query built via f-string passed to execute()/executemany(). Use parameterized queries with placeholders instead.
(coderabbit.sql-injection.python-fstring-execute)
[ERROR] 158-162: SQL query built via f-string passed to execute()/executemany(). Use parameterized queries with placeholders instead.
(coderabbit.sql-injection.python-fstring-execute)
🪛 Ruff (0.15.21)
python/tests/test_context_isolation.py
[warning] 52-52: Missing return type annotation for private function run_scope
Add return type annotation: None
(ANN202)
[warning] 61-61: Missing return type annotation for private function main
Add return type annotation: None
(ANN202)
[warning] 71-71: Missing return type annotation for private function run_tool
(ANN202)
[warning] 80-80: Boolean positional value in function call
(FBT003)
[warning] 81-81: Unused lambda argument: name
(ARG005)
[warning] 87-87: Unused lambda argument: name
(ARG005)
[warning] 107-107: Missing return type annotation for private function main
(ANN202)
[warning] 128-128: Missing return type annotation for private function run_llm
(ANN202)
[warning] 134-134: Missing return type annotation for private function intercept
(ANN202)
[warning] 134-134: Unused function argument: name
(ARG001)
[warning] 140-140: Boolean positional value in function call
(FBT003)
[warning] 166-166: Missing return type annotation for private function stream_func
(ANN202)
[warning] 184-184: Missing return type annotation for private function main
(ANN202)
🔇 Additional comments (4)
python/nemo_relay/scope.py (1)
24-24: LGTM!Also applies to: 260-260
python/nemo_relay/llm.py (1)
36-36: LGTM!Also applies to: 123-123, 174-174, 262-262, 360-360, 393-393, 412-412
python/nemo_relay/tools.py (1)
25-25: LGTM!Also applies to: 99-100, 132-133, 185-186, 205-206, 224-225
python/tests/test_context_isolation.py (1)
68-203: LGTM!
|
/merge |
2 participants
Overview
Restore the active task's Python scope stack before every context-sensitive native lifecycle call. This prevents concurrent asyncio tasks from closing scopes or invoking tool/LLM middleware against a sibling task's stack.
Details
pop()wrapper.Where should the reviewer start?
Start with
python/nemo_relay/_context.py, thenpython/tests/test_context_isolation.pyfor the concurrent task coverage.Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)
Summary by CodeRabbit
Bug Fixes
Tests