feat(studio): Add multiple choice UI for the Nemo Agent

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": "<the question text>",
+              "header": "<short chip label, optional>",
+              "multiSelect": false,
+              "options": [
+                {"label": "<choice>", "description": "<what it means, optional>"}
+              ]
+            }
+            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,
 ]
 
 

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"),

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 "

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.",
             (

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 <time>' accordion" in prompt
     assert "Never end a message with only a plain-text question" in prompt
     assert "call the matching select_* tool before completing the message" in prompt
-    assert "ask one concise plain-text question" in prompt
-    assert "AskUserQuestion" not in prompt
+    assert "you MUST call AskUserQuestion to render a selectable options picker" in prompt
+    assert "Only ask a concise plain-text question for genuinely open-ended" in prompt
     assert "A timeout, disconnect, or other interactive-tool error is not permission to continue" in prompt
     assert "summary's final sentence MUST state the exact unresolved selection or action" in prompt
     assert "Never show only the investigation result" in prompt
@@ -1734,7 +1734,9 @@ async def fake_stream(
     assert "Current Studio route path: /workspaces/default/dashboard/code-agent" in captured["studio_system_prompt"]
     assert "you MUST call select_agent" in captured["studio_system_prompt"]
     assert "you MUST call select_model" in captured["studio_system_prompt"]
-    assert "no dedicated Studio picker" in captured["studio_system_prompt"]
+    # The options-picker directive is present and rewritten to the deployed agent's tool name.
+    assert "you MUST call ask_user_question to render a selectable options picker" in captured["studio_system_prompt"]
+    assert "AskUserQuestion" not in captured["studio_system_prompt"]
     assert "Prefer NeMo Studio MCP tools and Studio views over CLI commands" in captured["studio_system_prompt"]
     assert "Do not tell the user to run nemo CLI commands" in captured["studio_system_prompt"]
     assert "when a Studio view, Studio link, or Studio progress card is available" in captured["studio_system_prompt"]