feat(studio): Studio agent add tool usage info - #1067
Open
htolentino-nvidia wants to merge 3 commits into
Open
Conversation
dmariali
approved these changes
Aug 4, 2026
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Signed-off-by: Henrique Tolentino <htolentino@nvidia.com>
htolentino-nvidia
force-pushed
the
studio-agent-add-tool-usage-info/htolentino
branch
from
August 4, 2026 15:16
bf1b5c9 to
1fef8df
Compare
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThe NeMo wrapper now attaches profiler callbacks. Studio authorizes workspaces through Entity Store, uses canonical names, streams agent responses, and emits sanitized tool-use events before final responses. ChangesInteractive NeMo agent flow
Sequence Diagram(s)sequenceDiagram
participant StudioCodingAgent
participant EntityStore
participant NeMoAgent
participant SessionStream
StudioCodingAgent->>EntityStore: authorize requested workspace
EntityStore-->>StudioCodingAgent: canonical workspace name
StudioCodingAgent->>NeMoAgent: streaming request with canonical workspace
NeMoAgent-->>StudioCodingAgent: assistant content and tool-step data
StudioCodingAgent->>StudioCodingAgent: parse and deduplicate tool steps
StudioCodingAgent->>SessionStream: emit queued tool-use events
StudioCodingAgent->>SessionStream: emit final assistant response
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (1)
agents/nemo-agent-local/src/nemo_agent/wrapper.py (1)
245-256: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winRemove the broad profiler fallback. NAT 1.8.0 does not require an active request context to construct
LangchainProfilerHandler. CatchingExceptionstill hides genuine regressions and disables streamed tool events. Construct the handler directly, and test callback attachment.🤖 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 `@agents/nemo-agent-local/src/nemo_agent/wrapper.py` around lines 245 - 256, Remove the try-except block surrounding the LangchainProfilerHandler construction since NAT 1.8.0 no longer requires an active request context. Assign the LangchainProfilerHandler instance directly to config["callbacks"] without the exception handler, eliminating the debug fallback path. This ensures genuine errors are not masked and tool-call tracing is reliably enabled.
🤖 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 `@agents/nemo-agent-local/src/nemo_agent/register.py`:
- Around line 597-606: Strengthen validation in the AskUserQuestion handling
around parsed and before _call_studio_tool: require each question object to
contain a non-empty string question, a non-empty options list with valid option
objects and non-empty string labels, and validate the types of any optional
fields according to the picker contract. Return the existing error format for
every rejected shape, and add tests covering each invalid question and option
case.
In `@services/studio/src/nmp/studio/coding_agents.py`:
- Around line 1420-1424: Update _stream_nemo_agent so invocation completion does
not end queue processing before pending tool-use events are emitted: after
invocation.done(), emit the completed queued_event, drain all remaining events
from queue, then emit the final assistant event. Add a test covering a tool step
and completion arriving in the same event-loop turn, verifying the tool event is
preserved before the final response.
- Around line 1415-1419: Update the deduplication logic around step_id so type
validation occurs before checking membership in seen_tool_ids. Skip non-string
IDs, and only test and add valid string IDs to the set while preserving the
existing duplicate-skipping behavior.
- Around line 1339-1354: Update _tool_use_stream_event and its
_invoke_nemo_agent call path to sanitize tool_input before json serialization,
removing studio_session_id and redacting other sensitive fields; preferably
expose only an explicit safe display-field allowlist per tool while preserving
the browser event structure.
---
Nitpick comments:
In `@agents/nemo-agent-local/src/nemo_agent/wrapper.py`:
- Around line 245-256: Remove the try-except block surrounding the
LangchainProfilerHandler construction since NAT 1.8.0 no longer requires an
active request context. Assign the LangchainProfilerHandler instance directly to
config["callbacks"] without the exception handler, eliminating the debug
fallback path. This ensures genuine errors are not masked and tool-call tracing
is reliably enabled.
🪄 Autofix
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: CHILL
Plan: Enterprise
Run ID: b2cba84c-657b-4f02-8afd-e23f90ca9bec
📒 Files selected for processing (6)
agents/nemo-agent-local/src/nemo_agent/register.pyagents/nemo-agent-local/src/nemo_agent/wrapper.pyagents/nemo-agent-local/tests/test_nemo_agent.pyservices/studio/src/nmp/studio/coding_agent_mcp_tools.pyservices/studio/src/nmp/studio/coding_agents.pyservices/studio/tests/unit/test_coding_agents.py
Contributor
There was a problem hiding this comment.
Caution
Inline review comments failed to post. This is likely due to GitHub's internal server error or limits when posting large numbers of comments. If you are seeing this consistently it is likely a permissions issue. Please check "Moderation" -> "Code review limits" under your organization settings.
Actionable comments posted: 4
🧹 Nitpick comments (1)
agents/nemo-agent-local/src/nemo_agent/wrapper.py (1)
245-256: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winRemove the broad profiler fallback. NAT 1.8.0 does not require an active request context to construct
LangchainProfilerHandler. CatchingExceptionstill hides genuine regressions and disables streamed tool events. Construct the handler directly, and test callback attachment.🤖 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 `@agents/nemo-agent-local/src/nemo_agent/wrapper.py` around lines 245 - 256, Remove the try-except block surrounding the LangchainProfilerHandler construction since NAT 1.8.0 no longer requires an active request context. Assign the LangchainProfilerHandler instance directly to config["callbacks"] without the exception handler, eliminating the debug fallback path. This ensures genuine errors are not masked and tool-call tracing is reliably enabled.
🤖 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 `@agents/nemo-agent-local/src/nemo_agent/register.py`:
- Around line 597-606: Strengthen validation in the AskUserQuestion handling
around parsed and before _call_studio_tool: require each question object to
contain a non-empty string question, a non-empty options list with valid option
objects and non-empty string labels, and validate the types of any optional
fields according to the picker contract. Return the existing error format for
every rejected shape, and add tests covering each invalid question and option
case.
In `@services/studio/src/nmp/studio/coding_agents.py`:
- Around line 1420-1424: Update _stream_nemo_agent so invocation completion does
not end queue processing before pending tool-use events are emitted: after
invocation.done(), emit the completed queued_event, drain all remaining events
from queue, then emit the final assistant event. Add a test covering a tool step
and completion arriving in the same event-loop turn, verifying the tool event is
preserved before the final response.
- Around line 1415-1419: Update the deduplication logic around step_id so type
validation occurs before checking membership in seen_tool_ids. Skip non-string
IDs, and only test and add valid string IDs to the set while preserving the
existing duplicate-skipping behavior.
- Around line 1339-1354: Update _tool_use_stream_event and its
_invoke_nemo_agent call path to sanitize tool_input before json serialization,
removing studio_session_id and redacting other sensitive fields; preferably
expose only an explicit safe display-field allowlist per tool while preserving
the browser event structure.
---
Nitpick comments:
In `@agents/nemo-agent-local/src/nemo_agent/wrapper.py`:
- Around line 245-256: Remove the try-except block surrounding the
LangchainProfilerHandler construction since NAT 1.8.0 no longer requires an
active request context. Assign the LangchainProfilerHandler instance directly to
config["callbacks"] without the exception handler, eliminating the debug
fallback path. This ensures genuine errors are not masked and tool-call tracing
is reliably enabled.
🪄 Autofix
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: CHILL
Plan: Enterprise
Run ID: b2cba84c-657b-4f02-8afd-e23f90ca9bec
📒 Files selected for processing (6)
agents/nemo-agent-local/src/nemo_agent/register.pyagents/nemo-agent-local/src/nemo_agent/wrapper.pyagents/nemo-agent-local/tests/test_nemo_agent.pyservices/studio/src/nmp/studio/coding_agent_mcp_tools.pyservices/studio/src/nmp/studio/coding_agents.pyservices/studio/tests/unit/test_coding_agents.py
🛑 Comments failed to post (1)
agents/nemo-agent-local/src/nemo_agent/register.py (1)
597-606: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Validate each question object before calling Studio.
Line 601 accepts
[{}], missingquestion, emptyoptions, and malformed option entries. These values violate the documented picker contract and can block the interactive flow.Validate required strings, a non-empty
optionslist, option labels, and optional field types. Add tests for each rejected shape.Proposed validation
- if not isinstance(parsed, list) or not parsed or not all(isinstance(q, dict) for q in parsed): + if ( + not isinstance(parsed, list) + or not parsed + or any( + not isinstance(question, dict) + or not isinstance(question.get("question"), str) + or not question["question"].strip() + or not isinstance(question.get("options"), list) + or not question["options"] + or any( + not isinstance(option, dict) + or not isinstance(option.get("label"), str) + or not option["label"].strip() + for option in question["options"] + ) + for question in parsed + ) + ): return "Error: `questions` must be a non-empty JSON array of question objects."🤖 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 `@agents/nemo-agent-local/src/nemo_agent/register.py` around lines 597 - 606, Strengthen validation in the AskUserQuestion handling around parsed and before _call_studio_tool: require each question object to contain a non-empty string question, a non-empty options list with valid option objects and non-empty string labels, and validate the types of any optional fields according to the picker contract. Return the existing error format for every rejected shape, and add tests covering each invalid question and option case.
Contributor
Author
|
@CodeRabbit re-review now |
Contributor
|
✅ Action performedReview finished.
|
Contributor
|
htolentino-nvidia
force-pushed
the
studio-agent-add-tool-usage-info/htolentino
branch
3 times, most recently
from
August 4, 2026 17:23
e123125 to
308fe3c
Compare
Contributor
Author
|
One more review please @CodeRabbit |
Contributor
|
✅ Action performedReview finished.
|
Contributor
There was a problem hiding this comment.
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 `@services/studio/src/nmp/studio/coding_agents.py`:
- Around line 144-147: Update _session_workspace_cache and its access paths to
include the resolved caller identity, such as user ID, in the cache key
alongside session_id and requested workspace. Ensure workspace authorization
lookups use this caller-scoped key so cached membership results cannot be reused
by another caller.
🪄 Autofix
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: CHILL
Plan: Enterprise
Run ID: d18a0803-f816-46ba-a186-e47b13266ab2
📒 Files selected for processing (2)
services/studio/src/nmp/studio/coding_agents.pyservices/studio/tests/unit/test_coding_agents.py
- Strip framework-injected studio_session_id from tool-use events before they reach the browser. - Gate tool-step dedup on a validated string id. - Drain queued tool-use events after the invocation completes so late events are not dropped before the final assistant message. - Resolve the CodeQL partial-SSRF finding: validate the requested workspace against the Entity Store (scoped to the caller's forwarded auth) and build the agent URL from the platform's own copy of the name, so no client-derived value reaches the outbound request URL. - Cache confirmed workspace names per session so the membership lookup runs once per session/workspace instead of on every message; only successful resolutions are cached, and the cache is cleared on session eviction. Signed-off-by: Henrique Tolentino <htolentino@nvidia.com>
Session ids carry no caller identity: create_session mints a bare UUID and nothing binds a session to a user. Keying the workspace-membership cache on (session, workspace) alone therefore let a second caller reuse the first caller's authorization decision and skip the Entity Store check. Include a SHA-256 fingerprint of the caller's forwarded credentials in the cache key so a cached decision is never reused across callers. Only the digest is retained, never the raw credential. Signed-off-by: Henrique Tolentino <htolentino@nvidia.com>
htolentino-nvidia
force-pushed
the
studio-agent-add-tool-usage-info/htolentino
branch
from
August 4, 2026 17:45
308fe3c to
ecdc75e
Compare
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
3 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.
Summary by CodeRabbit
New Features
Bug Fixes