feat: add opaque tool result annotations - #575
Draft
bbednarski9 wants to merge 1 commit into
Draft
Conversation
WalkthroughThis change adds annotation-aware tool execution frames, shared legacy/frame interception chains, lifecycle and observability propagation, worker and native plugin transport support, and public C, Go, Node.js, and Python APIs with serialization and integration coverage. ChangesTool execution frame support
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant Relay
participant Intercepts
participant Plugin
participant Observability
Caller->>Relay: execute frame
Relay->>Intercepts: run shared legacy/frame chain
Intercepts->>Plugin: invoke registered frame middleware
Plugin->>Intercepts: call downstream frame continuation
Intercepts->>Relay: return frame outcome and pending marks
Relay->>Observability: emit result and opaque annotation
Relay->>Caller: return ToolExecutionFrame
Important Pre-merge checks failedPlease resolve all errors before merging. Addressing warnings is optional. ❌ Failed checks (1 error)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
github-actions
Bot
added
Feature
bbednarski9
force-pushed
the
feat/tool-result-annotations-445
branch
from
July 28, 2026 19:45
1b1ebb7 to
6ea55f1
Compare
License DiffCompared against Lockfile license changesLockfile License ChangesRustAdded
Removed
Updated/Changed
NodeAdded
Removed
Updated/Changed
PythonAdded
Removed
Updated/Changed
Status output |
bbednarski9
force-pushed
the
feat/tool-result-annotations-445
branch
from
July 28, 2026 20:40
6ea55f1 to
78931ea
Compare
bbednarski9
force-pushed
the
feat/tool-result-annotations-445
branch
from
July 28, 2026 20:47
78931ea to
09e7362
Compare
willkill07
force-pushed
the
wkk_relay-509-async-ffi-go
branch
11 times, most recently
from
July 28, 2026 23:32
6b8b633 to
4002875
Compare
willkill07
force-pushed
the
wkk_relay-509-async-ffi-go
branch
6 times, most recently
from
July 29, 2026 03:44
514769d to
c4ce063
Compare
bbednarski9
force-pushed
the
feat/tool-result-annotations-445
branch
from
July 30, 2026 21:48
09e7362 to
80efe8f
Compare
Signed-off-by: Bryan Bednarski <bbednarski@nvidia.com>
bbednarski9
force-pushed
the
feat/tool-result-annotations-445
branch
from
July 30, 2026 21:51
80efe8f to
c9e8580
Compare
bbednarski9
marked this pull request as ready for review
July 30, 2026 21:54
Contributor
Author
|
@CodeRabbit review |
✅ Action performedReview finished.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 10
🤖 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/native_plugin_tests.rs`:
- Around line 420-423: Wrap the long indexing expression in the assertion for
tool_frame_end so it conforms to rustfmt’s 100-column formatting, matching the
equivalent assertion in middleware_tests.rs. Run cargo fmt --all and preserve
the assertion’s behavior.
In `@crates/ffi/src/callable.rs`:
- Around line 551-578: Document in the header comment for
NemoRelayToolExecFrameInterceptCb that next_ctx and the associated next_fn are
valid only during the callback invocation and must not be retained or called
after the callback returns. Match the existing lifetime-contract wording used by
other raw intercept bindings, without changing the trampoline implementation.
- Around line 449-472: Update wrap_tool_exec_frame_fn in
crates/ffi/src/callable.rs lines 449-472 to call clear_last_error() before cb,
bind json_result_from_ptr without immediately propagating its error, free
result_ptr, then apply ?. Apply the same clear-last-error and free-before-?
ordering to the second frame wrapper in crates/ffi/src/callable.rs lines
580-594.
In `@crates/node/src/api/mod.rs`:
- Around line 2270-2291: Restore ambient scope and publication-context wiring
for frame-aware tool lifecycle calls. In crates/node/src/api/mod.rs:2270-2291,
wrap tool_call_end_frame’s core call with with_effective_scope_stack; in
crates/node/src/api/mod.rs:2353-2409, use effective_scope_context and
with_publication_callback_context around tool_call_execute_frame. In
crates/python/src/py_api/mod.rs:672-693, wrap tool_call_end_frame with
with_python_publication_context; in crates/python/src/py_api/mod.rs:792-834,
capture the publication context and buffer and apply
with_task_nested_publication_buffer and with_task_publication_context, matching
each raw sibling’s behavior.
In `@crates/node/tests/tools_tests.mjs`:
- Around line 344-361: Extend the test after locating the manual end event to
also find the `node_sync_frame_tool` scope-end event and assert its
`category_profile.tool_result_annotation.producer` is `node-sync`. Keep the
existing `waitForSubscriberCallbacks` condition and manual-event assertion
unchanged while covering the promised annotation parity.
In `@crates/types/tests/serialization_tests.rs`:
- Around line 262-286: Extend the ToolExecutionFrameOutcome test coverage
alongside its existing round-trip test: deserialize an object with a frame but
omitted pending_marks and an unknown field, assert pending_marks defaults to
empty, and assert deserialization fails when frame is omitted. Use the existing
ToolExecutionFrameOutcome symbols and preserve the raw-outcome contract pattern.
In `@docs/reference/tool-execution-intercept-outcomes.mdx`:
- Around line 112-120: Update the binding contract section to document
frame-aware return shapes for Python, Rust, Go, and Node.js, including
ToolExecutionFrameOutcome and Node.js’s { frame, pendingMarks? } form. Name the
grpc-v1 frame envelope schema explicitly alongside the existing raw schema, and
clarify that frame.result and optional frame.annotation are preserved opaquely
across these bindings.
In `@go/nemo_relay/callbacks.go`:
- Line 678: Replace the direct json.Marshal call in the frame trampoline with
the shared jsonMarshal helper, preserving its existing error handling and
ensuring the test hook is used consistently with the other serialization paths.
In `@go/nemo_relay/tools_test.go`:
- Around line 263-291: Protect the retainedNext access in the intercept test
with the same synchronization mechanism used for the native callback, guarding
both the assignment inside the intercept closure and the subsequent nil check
and invocation. Ensure the lock establishes a race-detector-visible
synchronization edge between the callback thread and the test goroutine while
preserving the existing continuation behavior.
In `@python/plugin/src/nemo_relay_plugin/_api.py`:
- Around line 1191-1205: Update _push_registration and the registration paths
for TOOL_EXECUTION_INTERCEPT and TOOL_EXECUTION_FRAME_INTERCEPT so both surfaces
share one deduplication namespace for tool-execution names. Reject cross-surface
duplicates immediately with the existing local WorkerSdkError behavior, while
preserving separate deduplication for unrelated surfaces; apply the same
handling to register_tool_execution_frame_intercept and the corresponding legacy
registration flow.
🪄 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: 1b949d06-abd7-493a-823f-ef28957759af
📒 Files selected for processing (70)
crates/core/src/api/registry.rscrates/core/src/api/runtime.rscrates/core/src/api/runtime/callbacks.rscrates/core/src/api/runtime/state.rscrates/core/src/api/runtime/subscriber_dispatcher.rscrates/core/src/api/tool.rscrates/core/src/context/registries.rscrates/core/src/observability/atif.rscrates/core/src/observability/mod.rscrates/core/src/observability/openinference.rscrates/core/src/observability/otel.rscrates/core/src/plugin.rscrates/core/src/plugin/dynamic/native.rscrates/core/src/plugin/dynamic/worker.rscrates/core/tests/fixtures/native_plugin/src/lib.rscrates/core/tests/fixtures/worker_plugin/src/main.rscrates/core/tests/integration/middleware_tests.rscrates/core/tests/integration/native_plugin_tests.rscrates/core/tests/integration/worker_plugin_tests.rscrates/core/tests/unit/atif_tests.rscrates/core/tests/unit/observability/openinference_tests.rscrates/core/tests/unit/observability/otel_tests.rscrates/ffi/nemo_relay.hcrates/ffi/src/api/mod.rscrates/ffi/src/api/plugin.rscrates/ffi/src/api/scope_registry.rscrates/ffi/src/api/tool_lifecycle.rscrates/ffi/src/api/tool_registry.rscrates/ffi/src/callable.rscrates/ffi/tests/unit/api/execution_tests.rscrates/ffi/tests/unit/api_tests.rscrates/node/plugin.d.tscrates/node/src/api/mod.rscrates/node/src/callable.rscrates/node/src/types/mod.rscrates/node/tests/scope_local_tests.mjscrates/node/tests/tools_tests.mjscrates/plugin/src/lib.rscrates/plugin/tests/typed_callbacks.rscrates/python/src/py_api/mod.rscrates/python/src/py_callable.rscrates/python/src/py_plugin.rscrates/python/src/py_types/core.rscrates/python/src/py_types/mod.rscrates/types/src/api/tool.rscrates/types/tests/serialization_tests.rscrates/worker-proto/proto/nemo/relay/worker/v1/plugin_worker.protocrates/worker-proto/tests/proto_tests.rscrates/worker/src/lib.rscrates/worker/tests/worker_sdk_tests.rsdocs/reference/tool-execution-intercept-outcomes.mdxgo/nemo_relay/callbacks.gogo/nemo_relay/deregister_test.gogo/nemo_relay/nemo_relay.gogo/nemo_relay/plugin.gogo/nemo_relay/tools_test.gopython/nemo_relay/__init__.pypython/nemo_relay/__init__.pyipython/nemo_relay/_native.pyipython/nemo_relay/intercepts.pypython/nemo_relay/plugin.pypython/nemo_relay/plugin.pyipython/nemo_relay/scope_local.pypython/nemo_relay/tools.pypython/plugin/src/nemo_relay_plugin/__init__.pypython/plugin/src/nemo_relay_plugin/_api.pypython/tests/plugin/test_public_api_docstrings.pypython/tests/plugin/test_worker_sdk.pypython/tests/test_scope_local.pypython/tests/test_tools.py
Comment on lines
+420
to
+423
| assert_eq!( | ||
| tool_frame_end.category_profile().unwrap().extra[TOOL_RESULT_ANNOTATION_PROFILE_KEY]["producer_annotation"], | ||
| true | ||
| ); |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Line 421 exceeds rustfmt's 100-column default and will fail cargo fmt --check.
Break the index like the equivalent assertion in crates/core/tests/integration/middleware_tests.rs (lines 820-824).
As per coding guidelines, "Any Rust change must run cargo fmt --all".
🎨 Proposed rustfmt-compatible wrapping
assert_eq!(
- tool_frame_end.category_profile().unwrap().extra[TOOL_RESULT_ANNOTATION_PROFILE_KEY]["producer_annotation"],
+ tool_frame_end.category_profile().unwrap().extra[TOOL_RESULT_ANNOTATION_PROFILE_KEY]
+ ["producer_annotation"],
true
);📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| assert_eq!( | |
| tool_frame_end.category_profile().unwrap().extra[TOOL_RESULT_ANNOTATION_PROFILE_KEY]["producer_annotation"], | |
| true | |
| ); | |
| assert_eq!( | |
| tool_frame_end.category_profile().unwrap().extra[TOOL_RESULT_ANNOTATION_PROFILE_KEY] | |
| ["producer_annotation"], | |
| true | |
| ); |
🤖 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/tests/integration/native_plugin_tests.rs` around lines 420 - 423,
Wrap the long indexing expression in the assertion for tool_frame_end so it
conforms to rustfmt’s 100-column formatting, matching the equivalent assertion
in middleware_tests.rs. Run cargo fmt --all and preserve the assertion’s
behavior.
Source: Coding guidelines
Comment on lines
+449
to
+472
| /// Wrap a C annotation-aware tool execution callback into an async Rust closure. | ||
| pub fn wrap_tool_exec_frame_fn( | ||
| cb: NemoRelayToolExecFrameCb, | ||
| user_data: *mut libc::c_void, | ||
| free_fn: NemoRelayFreeFn, | ||
| ) -> Box< | ||
| dyn Fn(Json) -> Pin<Box<dyn Future<Output = Result<ToolExecutionFrame>> + Send>> + Send + Sync, | ||
| > { | ||
| let ud = make_user_data(user_data, free_fn); | ||
| Box::new(move |args: Json| { | ||
| let ud = ud.clone(); | ||
| Box::pin(async move { | ||
| let c_args = json_to_c_string(&args); | ||
| let result_ptr = unsafe { cb(ud.ptr, c_args) }; | ||
| unsafe { nemo_relay_string_free_internal(c_args) }; | ||
| let frame_json = | ||
| json_result_from_ptr(result_ptr, "tool execution frame callback failed")?; | ||
| unsafe { nemo_relay_string_free_internal(result_ptr) }; | ||
| serde_json::from_value::<ToolExecutionFrame>(frame_json).map_err(|error| { | ||
| FlowError::Internal(format!("invalid tool execution frame JSON: {error}")) | ||
| }) | ||
| }) | ||
| }) | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Both frame wrappers drop the error-slot and free discipline used by the existing raw wrappers. Neither calls clear_last_error() before invoking the C callback (so json_result_from_ptr's null path can report a stale error), and both propagate the parse error with ? before freeing the callback-returned string, leaking it on invalid JSON.
crates/ffi/src/callable.rs#L449-L472: addclear_last_error()at the top of the async block and bindjson_result_from_ptr(...)without?, freeingresult_ptrbefore applying?.crates/ffi/src/callable.rs#L580-L594: addclear_last_error()before thecbcall and apply the same free-then-?ordering aroundresult_ptr.
📍 Affects 1 file
crates/ffi/src/callable.rs#L449-L472(this comment)crates/ffi/src/callable.rs#L580-L594
🤖 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/ffi/src/callable.rs` around lines 449 - 472, Update
wrap_tool_exec_frame_fn in crates/ffi/src/callable.rs lines 449-472 to call
clear_last_error() before cb, bind json_result_from_ptr without immediately
propagating its error, free result_ptr, then apply ?. Apply the same
clear-last-error and free-before-? ordering to the second frame wrapper in
crates/ffi/src/callable.rs lines 580-594.
Comment on lines
+551
to
+578
| unsafe extern "C" fn tool_frame_next_trampoline( | ||
| args_json: *const c_char, | ||
| next_ctx: *mut libc::c_void, | ||
| ) -> *mut c_char { | ||
| let next_arc = unsafe { &*(next_ctx as *const ToolExecutionFrameNextFn) }; | ||
| let next = next_arc.clone(); | ||
| let args = if args_json.is_null() { | ||
| Json::Null | ||
| } else { | ||
| let value = unsafe { CStr::from_ptr(args_json) }.to_string_lossy(); | ||
| serde_json::from_str(&value).unwrap_or(Json::Null) | ||
| }; | ||
| let handle = tokio::runtime::Handle::current(); | ||
| let result = tokio::task::block_in_place(|| handle.block_on(next(args))); | ||
| match result { | ||
| Ok(frame) => match serde_json::to_value(frame) { | ||
| Ok(value) => json_to_c_string(&value), | ||
| Err(error) => { | ||
| set_last_error(&error.to_string()); | ||
| std::ptr::null_mut() | ||
| } | ||
| }, | ||
| Err(error) => { | ||
| set_last_error(&error.to_string()); | ||
| std::ptr::null_mut() | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
🔒 Security & Privacy | 🔵 Trivial | ⚡ Quick win
Verify the C-side contract that next_ctx must not be retained past callback return.
The boxed ToolExecutionFrameNextFn is dropped immediately after cb returns (Line 583), so any retained next_fn/next_ctx pair invoked later dereferences freed memory. The Go binding guards this with callbackLifetime, but this raw C surface has no such guard — the header doc for NemoRelayToolExecFrameInterceptCb (Lines 123-132) should state the invocation-scoped lifetime explicitly, as the raw intercept type does for other bindings.
🤖 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/ffi/src/callable.rs` around lines 551 - 578, Document in the header
comment for NemoRelayToolExecFrameInterceptCb that next_ctx and the associated
next_fn are valid only during the callback invocation and must not be retained
or called after the callback returns. Match the existing lifetime-contract
wording used by other raw intercept bindings, without changing the trampoline
implementation.
Comment on lines
+2270
to
+2291
| /// End a manual tool span with an optional opaque result annotation. | ||
| #[napi] | ||
| pub fn tool_call_end_frame( | ||
| handle: &ToolHandle, | ||
| frame: ToolExecutionFrame, | ||
| data: Option<Json>, | ||
| metadata: Option<Json>, | ||
| timestamp: Option<f64>, | ||
| ) -> Result<()> { | ||
| let timestamp = parse_timestamp_micros(timestamp)?; | ||
| core_tool_api::tool_call_end_frame( | ||
| core_tool_api::ToolCallEndFrameParams::builder() | ||
| .handle(&handle.inner) | ||
| .frame(frame.into()) | ||
| .data_opt(opt_json(data)) | ||
| .metadata_opt(opt_json(metadata)) | ||
| .timestamp_opt(timestamp) | ||
| .build(), | ||
| ) | ||
| .map_err(to_napi_err) | ||
| } | ||
|
|
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
Frame-aware tool lifecycle entrypoints drop the ambient scope/publication-context wiring their raw siblings establish. In both Node.js and Python bindings, tool_call_end_frame and tool_call_execute_frame invoke the shared core pipeline without the context setup their tool_call_end/tool_call_execute counterparts perform, even though frame-aware and legacy execution intercepts/guardrails now run through the same priority-ordered chain.
crates/node/src/api/mod.rs#L2270-L2291: wrap thecore_tool_api::tool_call_end_framecall inwith_effective_scope_stack(&env, || {...}), matchingtool_call_end(Line 2256-2266).crates/node/src/api/mod.rs#L2353-L2409: resolvepublication_context_id/publication_bufferviaeffective_scope_context(&env)and wrap the future inwith_publication_callback_context(...), matchingtool_call_execute(Lines 2313-2314, 2324-2346).crates/python/src/py_api/mod.rs#L672-L693: wrap thecore_tool_api::tool_call_end_framecall inwith_python_publication_context(|| {...}), matchingtool_call_end(Line 648).crates/python/src/py_api/mod.rs#L792-L834: capturepublication_context/publication_bufferand wrap the future inwith_task_nested_publication_buffer(...)/with_task_publication_context(...), matchingtool_call_execute(Lines 750-751, 753-756).
Please confirm whether this is an intentional simplification for frame-based calls or a regression; if unintentional, wiring all four sites the same way as their raw siblings avoids inconsistent guardrail/subscriber behavior between the raw and frame-aware tool lifecycle paths.
📍 Affects 2 files
crates/node/src/api/mod.rs#L2270-L2291(this comment)crates/node/src/api/mod.rs#L2353-L2409crates/python/src/py_api/mod.rs#L672-L693crates/python/src/py_api/mod.rs#L792-L834
🤖 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/node/src/api/mod.rs` around lines 2270 - 2291, Restore ambient scope
and publication-context wiring for frame-aware tool lifecycle calls. In
crates/node/src/api/mod.rs:2270-2291, wrap tool_call_end_frame’s core call with
with_effective_scope_stack; in crates/node/src/api/mod.rs:2353-2409, use
effective_scope_context and with_publication_callback_context around
tool_call_execute_frame. In crates/python/src/py_api/mod.rs:672-693, wrap
tool_call_end_frame with with_python_publication_context; in
crates/python/src/py_api/mod.rs:792-834, capture the publication context and
buffer and apply with_task_nested_publication_buffer and
with_task_publication_context, matching each raw sibling’s behavior.
Comment on lines
+344
to
+361
| assert.deepEqual(frame.annotation, { | ||
| producer: 'node-sync', | ||
| }); | ||
| await waitForSubscriberCallbacks( | ||
| () => | ||
| events.some( | ||
| (event) => | ||
| event.name === 'node_manual_frame_tool' && event.kind === 'scope' && event.scope_category === 'end', | ||
| ) && | ||
| events.some( | ||
| (event) => | ||
| event.name === 'node_sync_frame_tool' && event.kind === 'scope' && event.scope_category === 'end', | ||
| ), | ||
| ); | ||
| const manualEnd = events.find( | ||
| (event) => event.name === 'node_manual_frame_tool' && event.kind === 'scope' && event.scope_category === 'end', | ||
| ); | ||
| assert.equal(manualEnd.category_profile.tool_result_annotation.producer, 'node-manual'); |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Test waits for the node_sync_frame_tool end event but never asserts it.
The name promises the two APIs "marshal the same frame shape", yet only the manual end event and the sync return value are checked. Assert the sync end-event annotation so the wait at lines 353-356 has a purpose and the parity claim is actually covered.
💚 Suggested assertion
const manualEnd = events.find(
(event) => event.name === 'node_manual_frame_tool' && event.kind === 'scope' && event.scope_category === 'end',
);
assert.equal(manualEnd.category_profile.tool_result_annotation.producer, 'node-manual');
+ const syncEnd = events.find(
+ (event) => event.name === 'node_sync_frame_tool' && event.kind === 'scope' && event.scope_category === 'end',
+ );
+ assert.deepEqual(syncEnd.category_profile.tool_result_annotation, frame.annotation);📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| assert.deepEqual(frame.annotation, { | |
| producer: 'node-sync', | |
| }); | |
| await waitForSubscriberCallbacks( | |
| () => | |
| events.some( | |
| (event) => | |
| event.name === 'node_manual_frame_tool' && event.kind === 'scope' && event.scope_category === 'end', | |
| ) && | |
| events.some( | |
| (event) => | |
| event.name === 'node_sync_frame_tool' && event.kind === 'scope' && event.scope_category === 'end', | |
| ), | |
| ); | |
| const manualEnd = events.find( | |
| (event) => event.name === 'node_manual_frame_tool' && event.kind === 'scope' && event.scope_category === 'end', | |
| ); | |
| assert.equal(manualEnd.category_profile.tool_result_annotation.producer, 'node-manual'); | |
| assert.deepEqual(frame.annotation, { | |
| producer: 'node-sync', | |
| }); | |
| await waitForSubscriberCallbacks( | |
| () => | |
| events.some( | |
| (event) => | |
| event.name === 'node_manual_frame_tool' && event.kind === 'scope' && event.scope_category === 'end', | |
| ) && | |
| events.some( | |
| (event) => | |
| event.name === 'node_sync_frame_tool' && event.kind === 'scope' && event.scope_category === 'end', | |
| ), | |
| ); | |
| const manualEnd = events.find( | |
| (event) => event.name === 'node_manual_frame_tool' && event.kind === 'scope' && event.scope_category === 'end', | |
| ); | |
| assert.equal(manualEnd.category_profile.tool_result_annotation.producer, 'node-manual'); | |
| const syncEnd = events.find( | |
| (event) => event.name === 'node_sync_frame_tool' && event.kind === 'scope' && event.scope_category === 'end', | |
| ); | |
| assert.deepEqual(syncEnd.category_profile.tool_result_annotation, frame.annotation); |
🤖 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/node/tests/tools_tests.mjs` around lines 344 - 361, Extend the test
after locating the manual end event to also find the `node_sync_frame_tool`
scope-end event and assert its
`category_profile.tool_result_annotation.producer` is `node-sync`. Keep the
existing `waitForSubscriberCallbacks` condition and manual-event assertion
unchanged while covering the promised annotation parity.
Comment on lines
+262
to
+286
| let decoded: ToolExecutionFrameOutcome = | ||
| serde_json::from_value(encoded).expect("frame outcome should deserialize"); | ||
| assert_eq!(decoded, outcome); | ||
| } | ||
|
|
||
| #[test] | ||
| fn tool_execution_frame_requires_raw_result_and_defaults_to_no_annotation() { | ||
| let frame: ToolExecutionFrame = serde_json::from_value(json!({ | ||
| "result": ["plain", "json"] | ||
| })) | ||
| .expect("annotation should be optional"); | ||
| assert_eq!(frame.result, json!(["plain", "json"])); | ||
| assert!(frame.annotation.is_none()); | ||
|
|
||
| assert!( | ||
| serde_json::from_value::<ToolExecutionFrame>(json!({ | ||
| "annotation": { | ||
| "version": 1, | ||
| "status": "succeeded" | ||
| } | ||
| })) | ||
| .is_err(), | ||
| "raw result remains required" | ||
| ); | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick win
Add ToolExecutionFrameOutcome decode-contract coverage to match the raw-outcome test.
The raw-outcome test (lines 207-221) pins defaulted pending_marks, unknown-field tolerance, and the required field. The frame outcome only gets a round-trip; frame being required and pending_marks defaulting are unpinned wire guarantees for worker/C envelopes.
♻️ Suggested addition
let defaults: ToolExecutionFrameOutcome = serde_json::from_value(json!({
"frame": {"result": "plain"},
"future_field": true
}))
.expect("omitted pending marks and unknown fields should be accepted");
assert!(defaults.pending_marks.is_empty());
assert!(
serde_json::from_value::<ToolExecutionFrameOutcome>(json!({"pending_marks": []})).is_err(),
"frame is required"
);🤖 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/types/tests/serialization_tests.rs` around lines 262 - 286, Extend the
ToolExecutionFrameOutcome test coverage alongside its existing round-trip test:
deserialize an object with a frame but omitted pending_marks and an unknown
field, assert pending_marks defaults to empty, and assert deserialization fails
when frame is omitted. Use the existing ToolExecutionFrameOutcome symbols and
preserve the raw-outcome contract pattern.
Comment on lines
+112
to
+120
| - Public C callbacks return Relay outcome JSON with `result` and optional | ||
| `pending_marks`. | ||
| - `grpc-v1` worker SDKs return a `ToolExecutionInterceptOutcome` in a | ||
| `JsonEnvelope` with schema `nemo.relay.ToolExecutionInterceptOutcome@1`. | ||
|
|
||
| Canonical JSON uses `pending_marks` and `category_profile` across bindings. | ||
| Relay's outcome JSON uses `pending_marks` and `category_profile` across | ||
| bindings. Frame-aware worker and C envelopes additionally carry the opaque | ||
| `frame.result` and optional `frame.annotation` values without interpreting | ||
| them. |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Binding Contract omits the frame-aware return shapes for Python, Rust, Go, and Node.js.
Lines 117-120 mention only "worker and C envelopes" for frame values, yet this PR exposes ToolExecutionFrameOutcome (Python/Rust/Go) and { frame, pendingMarks? } (Node.js). The grpc-v1 frame envelope schema name is also undocumented while the raw one is spelled out at line 115.
As per coding guidelines, "Update README.md, fern/, package READMEs, and binding-support notes when public behavior, package names, examples, or supported bindings change."
📝 Suggested addition
Relay's outcome JSON uses `pending_marks` and `category_profile` across
bindings. Frame-aware worker and C envelopes additionally carry the opaque
`frame.result` and optional `frame.annotation` values without interpreting
them.
+
+Frame-aware callbacks return the analogous per-binding shape: Python, Rust,
+native plugins, and Go return `ToolExecutionFrameOutcome`; Node.js returns
+`{ frame, pendingMarks? }`; `grpc-v1` worker SDKs return a
+`ToolExecutionFrameOutcome` in a `JsonEnvelope` with schema
+`nemo.relay.ToolExecutionFrameOutcome@1`.📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| - Public C callbacks return Relay outcome JSON with `result` and optional | |
| `pending_marks`. | |
| - `grpc-v1` worker SDKs return a `ToolExecutionInterceptOutcome` in a | |
| `JsonEnvelope` with schema `nemo.relay.ToolExecutionInterceptOutcome@1`. | |
| Canonical JSON uses `pending_marks` and `category_profile` across bindings. | |
| Relay's outcome JSON uses `pending_marks` and `category_profile` across | |
| bindings. Frame-aware worker and C envelopes additionally carry the opaque | |
| `frame.result` and optional `frame.annotation` values without interpreting | |
| them. | |
| - Public C callbacks return Relay outcome JSON with `result` and optional | |
| `pending_marks`. | |
| - `grpc-v1` worker SDKs return a `ToolExecutionInterceptOutcome` in a | |
| `JsonEnvelope` with schema `nemo.relay.ToolExecutionInterceptOutcome@1`. | |
| Relay's outcome JSON uses `pending_marks` and `category_profile` across | |
| bindings. Frame-aware worker and C envelopes additionally carry the opaque | |
| `frame.result` and optional `frame.annotation` values without interpreting | |
| them. | |
| Frame-aware callbacks return the analogous per-binding shape: Python, Rust, | |
| native plugins, and Go return `ToolExecutionFrameOutcome`; Node.js returns | |
| `{ frame, pendingMarks? }`; `grpc-v1` worker SDKs return a | |
| `ToolExecutionFrameOutcome` in a `JsonEnvelope` with schema | |
| `nemo.relay.ToolExecutionFrameOutcome@1`. |
🤖 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 `@docs/reference/tool-execution-intercept-outcomes.mdx` around lines 112 - 120,
Update the binding contract section to document frame-aware return shapes for
Python, Rust, Go, and Node.js, including ToolExecutionFrameOutcome and Node.js’s
{ frame, pendingMarks? } form. Name the grpc-v1 frame envelope schema explicitly
alongside the existing raw schema, and clarify that frame.result and optional
frame.annotation are preserved opaquely across these bindings.
Source: Coding guidelines
| setLastErrorMessage(err.Error()) | ||
| return nil | ||
| } | ||
| value, err := json.Marshal(frame) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
rg -nP --type=go -C6 'func jsonMarshal\('Repository: NVIDIA/NeMo-Relay
Length of output: 155
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '## jsonMarshal references\n'
rg -n --hidden --glob '!**/node_modules/**' --glob '!**/dist/**' --glob '!**/build/**' '\bjsonMarshal\b' .
printf '\n## callbacks.go slices\n'
nl -ba go/nemo_relay/callbacks.go | sed -n '650,720p'
printf '\n---\n'
nl -ba go/nemo_relay/callbacks.go | sed -n '840,930p'Repository: NVIDIA/NeMo-Relay
Length of output: 2740
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '## test_hooks.go\n'
sed -n '1,80p' go/nemo_relay/test_hooks.go
printf '\n## callbacks.go around the cited call sites\n'
awk 'NR>=660 && NR<=700 {printf "%d:%s\n", NR, $0}' go/nemo_relay/callbacks.go
printf '\n---\n'
awk 'NR>=850 && NR<=890 {printf "%d:%s\n", NR, $0}' go/nemo_relay/callbacks.go
printf '\n---\n'
awk 'NR>=970 && NR<=995 {printf "%d:%s\n", NR, $0}' go/nemo_relay/callbacks.goRepository: NVIDIA/NeMo-Relay
Length of output: 4313
Use the shared marshal helper here too. This keeps the frame trampoline aligned with the other paths and preserves the jsonMarshal test hook.
🤖 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 `@go/nemo_relay/callbacks.go` at line 678, Replace the direct json.Marshal call
in the frame trampoline with the shared jsonMarshal helper, preserving its
existing error handling and ensuring the test hook is used consistently with the
other serialization paths.
Comment on lines
+263
to
+291
| var retainedNext func(json.RawMessage) (ToolExecutionFrame, error) | ||
| err := RegisterToolExecutionFrameIntercept( | ||
| interceptName, | ||
| 1, | ||
| func(_ json.RawMessage, next func(json.RawMessage) (ToolExecutionFrame, error)) (ToolExecutionFrameOutcome, error) { | ||
| retainedNext = next | ||
| return ToolExecutionFrameOutcome{ | ||
| Frame: ToolExecutionFrame{Result: json.RawMessage(`{"short_circuit":true}`)}, | ||
| }, nil | ||
| }, | ||
| ) | ||
| if err != nil { | ||
| t.Fatalf("RegisterToolExecutionFrameIntercept failed: %v", err) | ||
| } | ||
| defer DeregisterToolExecutionIntercept(interceptName) | ||
|
|
||
| _, err = ToolCallExecuteFrame( | ||
| "go_frame_retained_next_tool", | ||
| json.RawMessage(`{}`), | ||
| func(args json.RawMessage) (ToolExecutionFrame, error) { | ||
| return ToolExecutionFrame{Result: args}, nil | ||
| }, | ||
| ) | ||
| if err != nil { | ||
| t.Fatalf("ToolCallExecuteFrame failed: %v", err) | ||
| } | ||
| if retainedNext == nil { | ||
| t.Fatal("frame intercept did not retain the continuation") | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Guard retainedNext — the callback may run on a non-Go-owned thread, so go test -race can flag this write/read pair.
The intercept closure is invoked from the native runtime; unlike events above (mutex-protected), retainedNext is written there and read in the test goroutine with no synchronization edge the race detector can see.
🔒 Proposed fix
+ var mu sync.Mutex
var retainedNext func(json.RawMessage) (ToolExecutionFrame, error)
err := RegisterToolExecutionFrameIntercept(
interceptName,
1,
func(_ json.RawMessage, next func(json.RawMessage) (ToolExecutionFrame, error)) (ToolExecutionFrameOutcome, error) {
+ mu.Lock()
retainedNext = next
+ mu.Unlock()
return ToolExecutionFrameOutcome{
Frame: ToolExecutionFrame{Result: json.RawMessage(`{"short_circuit":true}`)},
}, nil
},
)and take the same lock before the retainedNext == nil check and the call.
🤖 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 `@go/nemo_relay/tools_test.go` around lines 263 - 291, Protect the retainedNext
access in the intercept test with the same synchronization mechanism used for
the native callback, guarding both the assignment inside the intercept closure
and the subsequent nil check and invocation. Ensure the lock establishes a
race-detector-visible synchronization edge between the callback thread and the
test goroutine while preserving the existing continuation behavior.
Source: Path instructions
Comment on lines
+1191
to
+1205
| def register_tool_execution_frame_intercept( | ||
| self, | ||
| name: str, | ||
| callback: ToolExecutionFrameCallback, | ||
| *, | ||
| priority: int = 0, | ||
| ) -> None: | ||
| """Register middleware that carries an opaque annotation with the tool result. | ||
|
|
||
| The callback participates in the same host registry and priority order | ||
| as legacy tool execution intercepts. | ||
| """ | ||
| self._push_registration(name, pb.TOOL_EXECUTION_FRAME_INTERCEPT, priority, False) | ||
| self._handlers.tool_execution_frames[name] = callback | ||
|
|
There was a problem hiding this comment.
🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win
_push_registration dedup doesn't cover the shared tool-execution name namespace.
_push_registration only rejects duplicates within the same surface. But nemo_relay_deregister_tool_execution_intercept in the FFI tests shows raw (TOOL_EXECUTION_INTERCEPT) and frame-aware (TOOL_EXECUTION_FRAME_INTERCEPT) intercepts share one host-side name registry. A plugin registering the same name via register_tool_execution_intercept and register_tool_execution_frame_intercept will pass this client-side check but still collide at Register() time on the host, surfacing a less-specific, delayed error instead of an immediate local WorkerSdkError.
♻️ Proposed fix: treat both surfaces as one namespace for dedup
def _push_registration(self, name: str, surface: int, priority: int, break_chain: bool) -> None:
+ conflicting_surfaces = (
+ {pb.TOOL_EXECUTION_INTERCEPT, pb.TOOL_EXECUTION_FRAME_INTERCEPT}
+ if surface in (pb.TOOL_EXECUTION_INTERCEPT, pb.TOOL_EXECUTION_FRAME_INTERCEPT)
+ else {surface}
+ )
if any(
- registration.local_name == name and registration.surface == surface
+ registration.local_name == name and registration.surface in conflicting_surfaces
for registration in self._handlers.registrations
):
raise WorkerSdkError(f"handler {name!r} is already registered for surface {surface}")Please confirm the host's core registry indeed enforces one shared namespace across raw/frame tool-execution intercepts (per crates/core/src/api/registry.rs) to validate this fix is warranted.
Also applies to: 1335-1348
🤖 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 `@python/plugin/src/nemo_relay_plugin/_api.py` around lines 1191 - 1205, Update
_push_registration and the registration paths for TOOL_EXECUTION_INTERCEPT and
TOOL_EXECUTION_FRAME_INTERCEPT so both surfaces share one deduplication
namespace for tool-execution names. Reject cross-surface duplicates immediately
with the existing local WorkerSdkError behavior, while preserving separate
deduplication for unrelated surfaces; apply the same handling to
register_tool_execution_frame_intercept and the corresponding legacy
registration flow.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Base
Rebased directly onto
mainafter #570 and #571 merged. This PR no longer depends on the closed #572 branch.Summary
This is a breaking tool API change. It replaces the parallel tool-result APIs previously proposed in this PR with one canonical, protocol-neutral contract used by every managed and manual tool completion:
ToolExecutionResultToolExecutionInterceptOutcomenextOwnership boundary
The harness owns tool dispatch, receipt, and the schema of
result. Relay transportsresultandannotationbut interprets neither. This is intentionally different from the normalizedAnnotatedLLMRequestandAnnotatedLLMResponseAPIs.Relay does not introduce a canonical tool-result schema, another registry, a dispatcher, a service, or a sidecar. The existing post-tool-call intercept chain remains intact.
MCP
MCP calls use the same
ToolExecutionResult; Relay has no MCP-specific result type.CallToolResultremains opaque inresult, includingcontent,structuredContent,_meta, andisErrorannotationremains separate from MCP_metaCallToolResultwithisError: trueremains successful application-visible result dataThe existing
nemo-relay mcplifecycle/bootstrap behavior is unchanged.Breaking compatibility gates
compat.native_api = "2"grpc-v2nemo.relay.worker.v1because the control-flow service layout is unchangedToolNextcarriesnemo.relay.ToolExecutionResult@1nemo.relay.ToolExecutionInterceptOutcome@2>=0.7Older native binaries and
grpc-v1workers are rejected during activation. There are no deprecated aliases, dual callback registries, or public raw-result adapters.Semantics
nullannotations normalize to absence and are omitted during serializationnextcalls retain Relay-owned pending marks internallyToolExecutionResult.resultcategory_profile.tool_result_annotationValidation
cargo check --workspace --all-targetscase was added; both legacy-result rejection tests pass independently
(610 current tests)
case was added; both legacy-result rejection tests pass independently
(342 current tests)
cargo fmt --all -- --checkgit diff --checkCloses #445