refactor: narrow core runtime registry surface by willkill07 · Pull Request #136 · NVIDIA/NeMo-Flow

willkill07 · 2026-05-21T13:11:16Z

Overview

Narrows the public core crate surface for registry internals while preserving the public runtime helpers used by crate consumers and bindings.

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

Make core registry storage types and scope-local registry internals crate-private.
Keep public tool and LLM request-intercept and conditional-execution helpers available through api::tool and api::llm.
Switch registry-stored callback aliases to Arc so guardrail snapshots can clone callbacks without exposing registry implementation details.
Update Rust, Python, Node, FFI, WASM, adaptive, and CLI call sites/tests to use public behavior rather than raw core registry fields.

Validation:

cargo fmt --all
cargo test -p nemo-flow
cargo test -p nemo-flow-adaptive
cargo check -p nemo-flow-adaptive
cargo check --workspace
cargo clippy --workspace --all-targets -- -D warnings
just test-rust
commit pre-hook checks passed, including cargo fmt, cargo clippy, cargo check, docs linkcheck, SPDX/file hygiene

Where should the reviewer start?

Start with crates/core/src/api/registry.rs and crates/core/src/registry.rs for the narrowed registry model, then crates/core/src/api/runtime/state.rs for the remaining crate-private runtime chains and snapshot evaluation.

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

Relates to: none

Summary by CodeRabbit

Bug Fixes
- Improved callback memory management by switching from boxed closures to atomic reference-counted (Arc) closures for all middleware callbacks, enabling safer concurrent callback usage.
Refactoring
- Consolidated internal registry entry representations into a unified generic structure for cleaner middleware management.
- Restricted several internal APIs from public to crate-private visibility to reduce API surface.
- Updated callback registration APIs to use standardized type aliases for consistency.

Signed-off-by: Will Killian <wkillian@nvidia.com>

coderabbitai · 2026-05-21T13:11:33Z

Walkthrough

This PR refactors the middleware registry system from Box-based to Arc-based callback ownership and replaces the caller-supplied priority-function approach with self-describing registry entries. The core changes propagate through type aliases, registration macros, chain builders, scope-local registries, language bindings, and test suites.

Changes

Core Registry Refactoring

Layer / File(s)	Summary
RegistryEntry trait and SortedRegistry refactoring `crates/core/src/registry.rs`	Introduces `RegistryEntry` trait with `name()` and `priority()` methods; refactors `SortedRegistry<T: RegistryEntry>` to be trait-bound, removes stored `priority_fn`, changes constructor to parameterless, updates `register()` to derive keys from `entry.name()`, restricts `deregister()` and `sorted_values()` to `pub(crate)`, adds `Default` impl.
Registry entry types and macro-generated registration `crates/core/src/api/registry.rs`	Introduces `RegistryRecord<T>` and `RequestIntercept<F>` payload types; replaces public `Intercept`/`ExecutionIntercept`/`GuardrailEntry` structs with `pub(crate)` type aliases; updates macro-generated registration code to call `Guardrail::new()`/`Intercept::new()`/`ExecutionIntercept::new()` constructors.
Runtime callback type aliases and visibility `crates/core/src/api/runtime/callbacks.rs`, `crates/core/src/api/runtime.rs`	Changes `ToolSanitizeFn`, `ToolInterceptFn`, `LlmSanitizeRequestFn`, `LlmSanitizeResponseFn`, and `LlmRequestInterceptFn` from `Box<dyn Fn...>` to `Arc<dyn Fn...>`; restricts `LlmStreamExecutionRegistryRef` and `LlmStreamExecutionRegistryRefs` to `pub(crate)` visibility; removes re-exports from `runtime` module.
NemoFlowContextState chain builders and registry fields `crates/core/src/api/runtime/state.rs`	Updates guardrail registry fields from `SortedRegistry<GuardrailEntry<...>>` to `SortedRegistry<Guardrail<...>>`, changes visibility to `pub(crate)`, updates initialization to `SortedRegistry::new()` without priority closure, refactors chain-builder methods to use `entry.payload(...)` and accept borrowed `&[Guardrail<...>]` snapshot entries instead of owned values, removes `ConditionalGuardrailSnapshot` type.
ScopeStack and scope-local registries `crates/core/src/api/runtime/scope_stack.rs`, `crates/core/src/context/registries.rs`	Updates `ScopeLocalRegistries` and registry fields to use `Guardrail<F>` types, restricts visibility to `pub(crate)`, removes priority-closure configuration from `SortedRegistry::new()`, updates merge helpers to work with new entry types and visibility.
Crate-level visibility `crates/core/src/lib.rs`	Removes `registry` module from public crate exports (`pub mod registry;` → `mod registry;`); updates crate documentation.
API call sites `crates/core/src/api/tool.rs`, `crates/core/src/api/llm.rs`	Updates `tool_call_execute` and `tool_conditional_execution` to pass `&entries` instead of `entries` to snapshot-chain builders; adjusts formatting in LLM execution functions.

Language Binding Wrapper Updates

Layer / File(s)	Summary
FFI, Node.js, Python, and WASM wrappers `crates/ffi/src/callable.rs`, `crates/node/src/callable.rs`, `crates/python/src/py_callable.rs`, `crates/wasm/src/callable.rs`	Switches callback construction from `Box::new(...)` to `Arc::new(...)` across all language bindings; updates return types to use shared callback type aliases (`ToolSanitizeFn`, `LlmSanitizeRequestFn`, `LlmSanitizeResponseFn`, `ToolInterceptFn`) instead of inline `Box<dyn Fn...>` trait objects; extends imports to include callback type aliases.

Test Suite Updates

Layer / File(s)	Summary
Core unit tests `crates/core/tests/unit/context_tests.rs`, `crates/core/tests/unit/plugin_tests.rs`, `crates/core/tests/unit/registry_tests.rs`, `crates/core/tests/unit/shared_tests.rs`	Introduces `PriorityItem` test struct implementing `RegistryEntry`; rewrites registry tests to construct `SortedRegistry::new()` and register via `SortedRegistry::register(item(...))` API; updates plugin and context tests to wrap closures in `Arc::new` and use new `Guardrail`/`Intercept`/`ExecutionIntercept` struct shapes.
Core integration tests `crates/core/tests/integration/api_surface_tests.rs`, `crates/core/tests/integration/middleware_tests.rs`, `crates/core/tests/integration/pipeline_tests.rs`, `crates/core/tests/integration/scope_local_tests.rs`, `crates/core/tests/integration/context_isolation_tests.rs`	Updates all middleware registrations to wrap closures in `Arc::new` instead of `Box::new`; adjusts concurrent-execution verification to use repeated `deregister()` assertions (`== false` for already-deregistered entries) instead of direct global context state inspection.
Adaptive and CLI tests `crates/adaptive/tests/unit/runtime_features_tests.rs`, `crates/adaptive/tests/integration/runtime_integration_tests.rs`, `crates/cli/tests/coverage/server_tests.rs`	Adds registry-level helper assertions for checking registration/absence status; updates feature tests to use new helpers instead of direct global context inspection; wraps closures in `Arc::new`.

Plugin Layer Updates

Layer / File(s)	Summary
Adaptive plugins `crates/adaptive/src/acg_component.rs`, `crates/adaptive/src/adaptive_hints_intercept.rs`	Updates intercept closure construction from `Box::new(...)` to `Arc::new(...)`.
Observability feature gate `crates/core/src/observability/plugin_component.rs`	Gates `std::time::Duration` import behind `#[cfg(any(feature = "otel", feature = "openinference"))]`.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related PRs

NVIDIA/NeMo-Flow#135: Refactors the core conditional-execution guardrail/registry plumbing (Guardrail payloads, snapshot-chain building) that this PR's CLI/OpenClaw toolConditionalExecution guardrail integration relies on; changes are directly connected at guardrail execution code level.
NVIDIA/NeMo-Flow#133: Both PRs touch conditional-guardrail evaluation in runtime/state.rs—feat: Emit guardrail scopes for conditional guardrails #133 adds snapshot-based guardrail scope emission using Box→Arc callbacks, while this PR refactors that conditional guardrail snapshot model (removes ConditionalGuardrailSnapshot, uses Guardrail/new registry payloads) so changes are code-level related.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	Title follows Conventional Commits format (refactor type, concise imperative summary) and is accurate—the PR narrows the public registry surface by making internals crate-private while preserving runtime helpers.
Description check	✅ Passed	Description covers all required template sections: overview explains narrowing the public surface while preserving helpers, details list concrete changes and validation performed, and review start points are specified.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
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.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

coderabbitai

Actionable comments posted: 1

🤖 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/adaptive/tests/unit/runtime_features_tests.rs`:
- Around line 590-591: After calling rollback_registrations(&mut registrations)
assert that all callback collections have been cleared: verify each registration
bucket/field on the registrations value (e.g., on_message/on_response/on_close
or the actual callback collections present in the Registrations struct) is empty
or that a provided helper like has_active_callbacks() returns false; update the
test in runtime_features_tests.rs to explicitly check every callback type the
Registrations struct exposes so rollback_registrations truly removed all
registered callbacks.

🪄 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: 5bb6e79d-aeb2-43c2-8464-73eee330abe5

📥 Commits

Reviewing files that changed from the base of the PR and between 78e83c0 and cba26f6.

📒 Files selected for processing (29)

crates/adaptive/src/acg_component.rs
crates/adaptive/src/adaptive_hints_intercept.rs
crates/adaptive/tests/integration/runtime_integration_tests.rs
crates/adaptive/tests/unit/runtime_features_tests.rs
crates/cli/tests/coverage/server_tests.rs
crates/core/src/api/llm.rs
crates/core/src/api/registry.rs
crates/core/src/api/runtime.rs
crates/core/src/api/runtime/callbacks.rs
crates/core/src/api/runtime/scope_stack.rs
crates/core/src/api/runtime/state.rs
crates/core/src/api/tool.rs
crates/core/src/context/registries.rs
crates/core/src/lib.rs
crates/core/src/observability/plugin_component.rs
crates/core/src/registry.rs
crates/core/tests/integration/api_surface_tests.rs
crates/core/tests/integration/context_isolation_tests.rs
crates/core/tests/integration/middleware_tests.rs
crates/core/tests/integration/pipeline_tests.rs
crates/core/tests/integration/scope_local_tests.rs
crates/core/tests/unit/context_tests.rs
crates/core/tests/unit/plugin_tests.rs
crates/core/tests/unit/registry_tests.rs
crates/core/tests/unit/shared_tests.rs
crates/ffi/src/callable.rs
crates/node/src/callable.rs
crates/python/src/py_callable.rs
crates/wasm/src/callable.rs

💤 Files with no reviewable changes (1)

crates/core/tests/integration/context_isolation_tests.rs

📜 Review details

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)

GitHub Check: Check / Run

🧰 Additional context used

📓 Path-based instructions (29)

**/*.rs

📄 CodeRabbit inference engine (.agents/skills/test-ffi-surface/SKILL.md)

**/*.rs: Run cargo fmt --all for FFI work as it is Rust work
Run just test-rust for FFI validation
Run cargo clippy --workspace --all-targets -- -D warnings to enforce warnings-as-errors linting

**/*.rs: Run cargo fmt --all for Rust code formatting
Run cargo clippy --workspace --all-targets -- -D warnings to enforce Rust linting with no warnings
Run just test-rust as the shared-runtime build/test wrapper for Rust changes

Use Rust snake_case naming convention for Rust code

**/*.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 Rust code formatting when Node changes touch Rust files
Run cargo clippy --workspace --all-targets -- -D warnings to enforce strict linting when Rust files changed as part of Node work

**/*.rs: Always run just test-rust when any Rust code changes
Always run cargo fmt --all when any Rust code changes
Always run cargo clippy --workspace --all-targets -- -D warnings when any Rust code changes

If any Rust files changed as part of the Python work, also run cargo fmt --all, just test-rust, and cargo clippy --workspace --all-targets -- -D warnings

Files:

crates/adaptive/src/acg_component.rs
crates/adaptive/src/adaptive_hints_intercept.rs
crates/core/src/observability/plugin_component.rs
crates/core/src/api/runtime.rs
crates/core/src/api/llm.rs
crates/adaptive/tests/integration/runtime_integration_tests.rs
crates/core/src/api/tool.rs
crates/cli/tests/coverage/server_tests.rs
crates/core/src/api/runtime/scope_stack.rs
crates/core/tests/unit/shared_tests.rs
crates/adaptive/tests/unit/runtime_features_tests.rs
crates/core/tests/unit/plugin_tests.rs
crates/core/src/api/runtime/callbacks.rs
crates/core/tests/integration/scope_local_tests.rs
crates/core/tests/integration/pipeline_tests.rs
crates/core/src/lib.rs
crates/python/src/py_callable.rs
crates/core/tests/unit/context_tests.rs
crates/core/tests/unit/registry_tests.rs
crates/core/src/api/registry.rs
crates/node/src/callable.rs
crates/core/tests/integration/api_surface_tests.rs
crates/core/src/context/registries.rs
crates/ffi/src/callable.rs
crates/core/tests/integration/middleware_tests.rs
crates/wasm/src/callable.rs
crates/core/src/registry.rs
crates/core/src/api/runtime/state.rs

**/*.{rs,go,js,ts,tsx}

📄 CodeRabbit inference engine (CONTRIBUTING.md)

Include SPDX license headers in all Rust, Go, JavaScript, and TypeScript source files using C-style comment syntax

Files:

crates/adaptive/src/acg_component.rs
crates/adaptive/src/adaptive_hints_intercept.rs
crates/core/src/observability/plugin_component.rs
crates/core/src/api/runtime.rs
crates/core/src/api/llm.rs
crates/adaptive/tests/integration/runtime_integration_tests.rs
crates/core/src/api/tool.rs
crates/cli/tests/coverage/server_tests.rs
crates/core/src/api/runtime/scope_stack.rs
crates/core/tests/unit/shared_tests.rs
crates/adaptive/tests/unit/runtime_features_tests.rs
crates/core/tests/unit/plugin_tests.rs
crates/core/src/api/runtime/callbacks.rs
crates/core/tests/integration/scope_local_tests.rs
crates/core/tests/integration/pipeline_tests.rs
crates/core/src/lib.rs
crates/python/src/py_callable.rs
crates/core/tests/unit/context_tests.rs
crates/core/tests/unit/registry_tests.rs
crates/core/src/api/registry.rs
crates/node/src/callable.rs
crates/core/tests/integration/api_surface_tests.rs
crates/core/src/context/registries.rs
crates/ffi/src/callable.rs
crates/core/tests/integration/middleware_tests.rs
crates/wasm/src/callable.rs
crates/core/src/registry.rs
crates/core/src/api/runtime/state.rs

**/*.{rs,py,go,js,ts,tsx}

📄 CodeRabbit inference engine (CONTRIBUTING.md)

Use SONAR_IGNORE_START / SONAR_IGNORE_END markers only for documented false positives that cannot be resolved in code; keep ignored blocks small, add explanatory comments, and require reviewer sign-off

Files:

crates/adaptive/src/acg_component.rs
crates/adaptive/src/adaptive_hints_intercept.rs
crates/core/src/observability/plugin_component.rs
crates/core/src/api/runtime.rs
crates/core/src/api/llm.rs
crates/adaptive/tests/integration/runtime_integration_tests.rs
crates/core/src/api/tool.rs
crates/cli/tests/coverage/server_tests.rs
crates/core/src/api/runtime/scope_stack.rs
crates/core/tests/unit/shared_tests.rs
crates/adaptive/tests/unit/runtime_features_tests.rs
crates/core/tests/unit/plugin_tests.rs
crates/core/src/api/runtime/callbacks.rs
crates/core/tests/integration/scope_local_tests.rs
crates/core/tests/integration/pipeline_tests.rs
crates/core/src/lib.rs
crates/python/src/py_callable.rs
crates/core/tests/unit/context_tests.rs
crates/core/tests/unit/registry_tests.rs
crates/core/src/api/registry.rs
crates/node/src/callable.rs
crates/core/tests/integration/api_surface_tests.rs
crates/core/src/context/registries.rs
crates/ffi/src/callable.rs
crates/core/tests/integration/middleware_tests.rs
crates/wasm/src/callable.rs
crates/core/src/registry.rs
crates/core/src/api/runtime/state.rs

{crates/adaptive/**,python/nemo_flow/{adaptive,plugin}.py,go/nemo_flow/{adaptive,**}/*.go,**/*.{ts,js,wasm}}