diff --git a/agents/nemo-agent-local/src/nemo_agent/register.py b/agents/nemo-agent-local/src/nemo_agent/register.py index e5cc5a6a43..3d124964a0 100644 --- a/agents/nemo-agent-local/src/nemo_agent/register.py +++ b/agents/nemo-agent-local/src/nemo_agent/register.py @@ -39,8 +39,14 @@ Use the deployment's active workspace unless the user explicitly names another one. If a required workspace, resource name, target, or consequential parameter -is missing or ambiguous, ask one focused clarification question and stop. Do -not repeat the same failing tool call or guess a destructive target. +is missing or ambiguous, ask one focused clarification and stop. When the answer +is one of a finite set of choices, ask it with the `ask_user_question` tool +(which renders a native options picker) instead of a free-text question; use a +plain-text question only for genuinely open-ended input. Do not repeat the same +failing tool call or guess a destructive target. + +For choosing an agent, model, dataset/fileset, or evaluation config, always use +the matching `select_*` tool rather than asking in plain text. """ SKILLS_DIR = Path(__file__).parent / "skills" DEFAULT_WORKSPACE = "default" @@ -433,8 +439,17 @@ def select_model( description: str = "", default_model: str | None = None, output_key: str = "model", + display_label: str | None = None, + field_label: str | None = None, + placeholder: str | None = None, + required_message: str | None = None, + submit_label: str | None = None, ) -> str: - """Render Studio's model selector and return the user's selected model.""" + """Render Studio's model selector and return the user's selected model. + + The optional labels (display_label/field_label/placeholder/required_message/ + submit_label) customize the picker's copy for the current task. + """ return json.dumps( _call_studio_tool( studio_session_id, @@ -444,6 +459,11 @@ def select_model( "description": description, "default_model": default_model, "output_key": output_key, + "display_label": display_label, + "field_label": field_label, + "placeholder": placeholder, + "required_message": required_message, + "submit_label": submit_label, }, ) ) @@ -476,6 +496,8 @@ def select_eval_config( title: str = "Select evaluation config", description: str = "", agent: str | None = None, + default_agent: str | None = None, + accepted_file_types: list[str] | None = None, ) -> str: """Render Studio's evaluation-config picker and return the selected config.""" return json.dumps( @@ -486,6 +508,8 @@ def select_eval_config( "title": title, "description": description, "agent": agent, + "default_agent": default_agent, + "accepted_file_types": accepted_file_types or [], }, ) ) @@ -499,6 +523,7 @@ def job_progress( source: str | None = None, title: str | None = None, description: str | None = None, + workspace: str | None = None, ) -> str: """Render a Studio progress card for a platform job that was just launched.""" return json.dumps( @@ -511,6 +536,7 @@ def job_progress( "source": source, "title": title, "description": description, + "workspace": workspace, }, ) ) @@ -539,6 +565,51 @@ def studio_link( ) +@tool +def ask_user_question(studio_session_id: str, questions: str) -> str: + """Ask the user one or more multiple-choice questions and return their selections. + + Prefer this over a free-text clarification whenever the choices are finite + (pick an option, yes/no, choose an approach). Studio renders a native options + picker; this blocks until the user answers or dismisses it. Routed through the + same permission channel Claude Code's ``AskUserQuestion`` uses, so the frontend + already renders it. + + Args: + studio_session_id: The current Studio session id (provided in your context). + questions: A JSON array string. Each element is a question object: + { + "question": "", + "header": "", + "multiSelect": false, + "options": [ + {"label": "", "description": ""} + ] + } + Example: + '[{"question": "Which model size?", "header": "Model", + "options": [{"label": "8B", "description": "faster, cheaper"}, + {"label": "70B", "description": "higher quality"}]}]' + + Returns: + JSON string of the user's answers, or a message if they declined. + """ + try: + parsed = json.loads(questions) + except (json.JSONDecodeError, TypeError) as exc: + return f"Error: `questions` must be a JSON array string: {exc}" + if not isinstance(parsed, list) or not parsed or not all(isinstance(q, dict) for q in parsed): + return "Error: `questions` must be a non-empty JSON array of question objects." + approval = _call_studio_tool( + studio_session_id, + "approval_prompt", + {"tool_name": "AskUserQuestion", "input": {"questions": parsed}}, + ) + if approval.get("behavior") != "allow": + return f"User declined to answer: {approval.get('message') or 'no selection made'}" + return json.dumps(approval.get("updatedInput") or {}) + + @tool def check_status(service: str, job_name: str) -> str: """Check the status of a platform job or deployment. @@ -585,6 +656,7 @@ def check_status(service: str, job_name: str) -> str: select_eval_config, job_progress, studio_link, + ask_user_question, ] diff --git a/agents/nemo-agent-local/tests/test_nemo_agent.py b/agents/nemo-agent-local/tests/test_nemo_agent.py index f0060bdd9e..8c1e34bf36 100644 --- a/agents/nemo-agent-local/tests/test_nemo_agent.py +++ b/agents/nemo-agent-local/tests/test_nemo_agent.py @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import json from pathlib import Path from unittest.mock import MagicMock, patch @@ -25,6 +26,7 @@ _serialize, _StreamSafeGraph, _studio_callback_url, + ask_user_question, check_status, create_nemo_agent, nemo_api, @@ -37,6 +39,7 @@ # through the @tool decorator, so we call the functions directly. _nemo_api = nemo_api.func # ty: ignore[unresolved-attribute] _check_status = check_status.func # ty: ignore[unresolved-attribute] +_ask_user_question = ask_user_question.func # ty: ignore[unresolved-attribute] AGENT_CONFIG = Path(__file__).parents[1] / "src" / "nemo_agent" / "nemo-agent.yml" TRUSTED_SESSION_ID = "00000000-0000-4000-8000-000000000001" @@ -499,6 +502,7 @@ def test_agent_has_expected_tools(self, mock_graph): "select_eval_config", "job_progress", "studio_link", + "ask_user_question", } def test_skills_dir_exists(self): @@ -520,6 +524,73 @@ def test_create_agent_passes_backend_visible_skills(self, mock_graph): assert len(skills) > 0 +class TestAskUserQuestion: + _QUESTIONS = '[{"question": "Which size?", "header": "Model", "options": [{"label": "8B"}, {"label": "70B"}]}]' + + def test_routes_through_approval_prompt_and_returns_answers(self): + with patch( + "nemo_agent.register._call_studio_tool", + return_value={"behavior": "allow", "updatedInput": {"Model": "70B"}}, + ) as studio_tool: + result = _ask_user_question( + studio_session_id=TRUSTED_SESSION_ID, + questions=self._QUESTIONS, + ) + + assert json.loads(result) == {"Model": "70B"} + studio_tool.assert_called_once_with( + TRUSTED_SESSION_ID, + "approval_prompt", + { + "tool_name": "AskUserQuestion", + "input": { + "questions": [ + { + "question": "Which size?", + "header": "Model", + "options": [{"label": "8B"}, {"label": "70B"}], + } + ] + }, + }, + ) + + def test_declined_returns_message(self): + with patch( + "nemo_agent.register._call_studio_tool", + return_value={"behavior": "deny", "message": "dismissed"}, + ): + result = _ask_user_question(studio_session_id=TRUSTED_SESSION_ID, questions=self._QUESTIONS) + assert result == "User declined to answer: dismissed" + + def test_invalid_json_is_reported_not_raised(self): + with patch("nemo_agent.register._call_studio_tool") as studio_tool: + result = _ask_user_question(studio_session_id=TRUSTED_SESSION_ID, questions="not json") + assert result.startswith("Error: `questions` must be a JSON array string") + studio_tool.assert_not_called() + + @pytest.mark.parametrize( + "questions", + [ + '{"question": "one?"}', # object, not a list + '"just a string"', # scalar + "5", # scalar + "[]", # empty array + "[1, 2]", # list of non-objects + '["a", {"question": "ok"}]', # mixed / invalid element + ], + ) + def test_non_question_array_is_rejected_without_calling_studio(self, questions): + with patch("nemo_agent.register._call_studio_tool") as studio_tool: + result = _ask_user_question(studio_session_id=TRUSTED_SESSION_ID, questions=questions) + assert result == "Error: `questions` must be a non-empty JSON array of question objects." + studio_tool.assert_not_called() + + def test_tool_schema_hides_approval_context(self): + # The picker takes only session id + questions; no config/approval leakage. + assert set(ask_user_question.args_schema.model_fields) == {"studio_session_id", "questions"} + + class TestDirectListFastPath: @pytest.mark.parametrize( ("prompt", "resource_path"), diff --git a/services/studio/src/nmp/studio/coding_agent_mcp_tools.py b/services/studio/src/nmp/studio/coding_agent_mcp_tools.py index 4ccd7c9ff2..0c93ce5afd 100644 --- a/services/studio/src/nmp/studio/coding_agent_mcp_tools.py +++ b/services/studio/src/nmp/studio/coding_agent_mcp_tools.py @@ -225,8 +225,11 @@ "call instead." ), ( - "For clarification, multiple-choice, yes/no, or freeform questions that do NOT map to one of " - "the select_* tools, ask one concise plain-text question." + "For any question with a finite set of choices that does NOT map to a select_* tool — including " + "yes/no, multiple-choice, 'pick one of these', or whenever you would offer the user options to " + "choose from — you MUST call AskUserQuestion to render a selectable options picker instead of " + "listing the choices in plain text. Only ask a concise plain-text question for genuinely " + "open-ended, free-form input that has no discrete options." ), ( "Only fall back to plain chat questions when no suitable UI tool exists, the user already " diff --git a/services/studio/src/nmp/studio/coding_agents.py b/services/studio/src/nmp/studio/coding_agents.py index e7e67beb0e..a26afed727 100644 --- a/services/studio/src/nmp/studio/coding_agents.py +++ b/services/studio/src/nmp/studio/coding_agents.py @@ -358,7 +358,7 @@ def _build_studio_system_prompt( "A timeout, disconnect, or other interactive-tool error is not permission to continue or repeat the question in plain text. Leave the input unresolved and tell the user the interactive request must be retried.", "A message that needs user input is not complete until you call the matching Studio input tool. Never end a message with only a plain-text question when an interactive tool applies.", "In particular, if you need an agent, model, dataset file, or evaluation config, call the matching select_* tool before completing the message; mentioning the needed selection in prose is not a substitute for the tool call.", - "For finite choices that have no dedicated Studio picker (for example deployments, jobs, or next actions) and for yes/no or multiple-choice clarifications, ask one concise plain-text question.", + "For any finite set of choices without a dedicated select_* picker — including yes/no, multiple-choice, 'pick one of these' (for example deployments, jobs, or next actions), or whenever you would offer the user options to choose from — you MUST call AskUserQuestion to render a selectable options picker instead of listing the choices in plain text. Only ask a concise plain-text question for genuinely open-ended, free-form input that has no discrete options.", "Conditional message-summary behavior:", "Use a Studio summary block only after substantive work that benefits from collapsing details.", "A summary block is required when you called one or more tools, ran commands, changed files or platform state, performed a multi-step investigation, or produced a long detailed response.", @@ -461,6 +461,11 @@ def _build_nemo_agent_system_prompt( for tool_name in tool_names: context = context.replace(f"mcp__{CLAUDE_MCP_SERVER_NAME}__{tool_name}", tool_name) prompt = prompt.replace(f"mcp__{CLAUDE_MCP_SERVER_NAME}__{tool_name}", tool_name) + # Claude Code's native options picker is AskUserQuestion; the deployed NeMo + # agent exposes it as the ask_user_question tool. Map the name so the + # "use the options picker" directives resolve to the deployed tool. + context = context.replace("AskUserQuestion", "ask_user_question") + prompt = prompt.replace("AskUserQuestion", "ask_user_question") return "\n".join( [ context, @@ -470,7 +475,7 @@ def _build_nemo_agent_system_prompt( "Deployed NeMo agent callback behavior:", ( f"For every select_agent, select_model, select_dataset_file, select_eval_config, " - f"job_progress, and studio_link call, pass studio_session_id='{session_id}'." + f"job_progress, studio_link, and ask_user_question call, pass studio_session_id='{session_id}'." ), "Do not reveal the Studio session id to the user.", ( diff --git a/services/studio/tests/unit/test_coding_agents.py b/services/studio/tests/unit/test_coding_agents.py index 540c5d59d3..df05d53187 100644 --- a/services/studio/tests/unit/test_coding_agents.py +++ b/services/studio/tests/unit/test_coding_agents.py @@ -848,8 +848,8 @@ def test_build_studio_system_prompt_includes_message_summary_contract(): assert "behind a 'worked for