From aa259202f3bb6e675268c041ac3cbecc7a73ffd1 Mon Sep 17 00:00:00 2001 From: Danielle Ali <44468613+dmariali@users.noreply.github.com> Date: Fri, 29 May 2026 13:34:04 -0400 Subject: [PATCH 1/8] chore(studio): Copy BE wiring and setup FF for coding agents in Studio Signed-off-by: Danielle Ali <44468613+dmariali@users.noreply.github.com> --- .../studio/src/nmp/studio/coding_agents.py | 530 ++++++++++++++++++ services/studio/src/nmp/studio/config.py | 17 +- .../studio/src/nmp/studio/env_mappings.py | 5 + services/studio/src/nmp/studio/service.py | 26 +- .../studio/tests/unit/test_coding_agents.py | 240 ++++++++ services/studio/tests/unit/test_service.py | 13 +- web/packages/studio/env/.env.dev.local.sample | 1 + web/packages/studio/env/.env.fastapi | 1 + .../studio/src/constants/environment.ts | 1 + .../constants/featureFlags/featureFlags.ts | 1 + 10 files changed, 828 insertions(+), 7 deletions(-) create mode 100644 services/studio/src/nmp/studio/coding_agents.py create mode 100644 services/studio/tests/unit/test_coding_agents.py diff --git a/services/studio/src/nmp/studio/coding_agents.py b/services/studio/src/nmp/studio/coding_agents.py new file mode 100644 index 00000000000..bdfe4586f34 --- /dev/null +++ b/services/studio/src/nmp/studio/coding_agents.py @@ -0,0 +1,530 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Local coding-agent bridge for Studio.""" + +import asyncio +import json +import logging +import os +import shutil +import uuid +from collections.abc import AsyncIterator +from pathlib import Path +from typing import Any + +from fastapi import APIRouter, FastAPI, HTTPException, Request +from fastapi.responses import JSONResponse, Response, StreamingResponse +from pydantic import BaseModel, Field +from starlette.routing import NoMatchFound + +logger = logging.getLogger(__name__) + +router = APIRouter(prefix="/v2/coding-agents") + +MCP_ROUTE_NAME = "studio_coding_agent_mcp" +PUBLIC_MCP_ROUTE_NAME = "studio_coding_agent_public_mcp" +PUBLIC_MCP_PATH = "/studio/api/coding-agents/mcp/{session_id}" +CLAUDE_MCP_SERVER_NAME = "nemo_studio" + +CLAUDE_PROJECTS_DIR = Path.home() / ".claude" / "projects" +SERVER_CWD = Path(os.getcwd()).resolve() + + +class NewSessionResponse(BaseModel): + """Response returned when Studio starts a new coding-agent session.""" + + session_id: str + + +class MessageRequest(BaseModel): + """A user message to send to the local coding agent.""" + + message: str = Field(min_length=1) + + +class PermissionDecision(BaseModel): + """Studio's decision for a pending local-agent tool permission request.""" + + approved: bool + reason: str | None = None + updated_input: dict[str, Any] | None = None + + +class HistorySessionResponse(BaseModel): + """Summary of a Claude session stored on disk.""" + + session_id: str + mtime: float + first_prompt: str + message_count: int + + +class SessionHistoryResponse(BaseModel): + """Claude session history normalized for Studio chat replay.""" + + session_id: str + items: list[dict[str, Any]] + + +_initialized_sessions: set[str] = set() +_session_streams: dict[str, asyncio.Queue[tuple[str, Any]]] = {} +_pending_permissions: dict[str, asyncio.Future[dict[str, Any]]] = {} + + +_APPROVAL_TOOL = { + "name": "approval_prompt", + "description": "Ask the human operator whether a tool call should be allowed.", + "inputSchema": { + "type": "object", + "properties": { + "tool_name": {"type": "string"}, + "input": {"type": "object"}, + "tool_use_id": {"type": "string"}, + }, + "required": ["tool_name", "input"], + }, +} + + +def mount_public_mcp_route(app: FastAPI) -> None: + """Mount the MCP callback under /studio so the local Claude CLI can call it.""" + app.add_api_route( + PUBLIC_MCP_PATH, + mcp_endpoint, + methods=["POST"], + name=PUBLIC_MCP_ROUTE_NAME, + include_in_schema=False, + ) + + +def _validate_session_id(session_id: str) -> str: + try: + return str(uuid.UUID(session_id)) + except ValueError as exc: + raise HTTPException(status_code=400, detail="session_id must be a UUID") from exc + + +def _project_history_dir() -> Path: + encoded = str(SERVER_CWD).replace("/", "-") + return CLAUDE_PROJECTS_DIR / encoded + + +def _first_user_prompt_and_count(path: Path) -> tuple[str | None, int]: + first_prompt: str | None = None + count = 0 + try: + with path.open("r", encoding="utf-8", errors="replace") as fh: + for line in fh: + line = line.strip() + if not line: + continue + try: + entry = json.loads(line) + except json.JSONDecodeError: + continue + if entry.get("type") != "user" or entry.get("isSidechain"): + continue + content = (entry.get("message") or {}).get("content") + if not isinstance(content, str): + continue + count += 1 + if first_prompt is None: + first_prompt = content + except OSError: + return None, 0 + return first_prompt, count + + +def _extract_assistant_parts(content: Any) -> list[dict[str, Any]]: + if not isinstance(content, list): + return [] + + parts: list[dict[str, Any]] = [] + for part in content: + if not isinstance(part, dict): + continue + part_type = part.get("type") + if part_type == "text": + text = part.get("text") + if isinstance(text, str) and text: + parts.append({"type": "text", "text": text}) + elif part_type == "thinking": + thinking = part.get("thinking") + if isinstance(thinking, str) and thinking: + parts.append({"type": "thinking", "thinking": thinking}) + elif part_type == "tool_use": + parts.append( + { + "type": "tool_use", + "name": part.get("name") or "tool", + "input": part.get("input") or {}, + } + ) + return parts + + +@router.post("/sessions", response_model=NewSessionResponse) +def create_session() -> NewSessionResponse: + """Create a new local coding-agent session.""" + return NewSessionResponse(session_id=str(uuid.uuid4())) + + +@router.get("/history/sessions", response_model=list[HistorySessionResponse]) +def list_history_sessions() -> list[HistorySessionResponse]: + """List Claude session histories for the Studio service working directory.""" + project_dir = _project_history_dir() + if not project_dir.is_dir(): + return [] + + sessions: list[HistorySessionResponse] = [] + for history_file in project_dir.glob("*.jsonl"): + try: + uuid.UUID(history_file.stem) + except ValueError: + continue + + first_prompt, count = _first_user_prompt_and_count(history_file) + if count == 0: + continue + + try: + mtime = history_file.stat().st_mtime + except OSError: + continue + + sessions.append( + HistorySessionResponse( + session_id=history_file.stem, + mtime=mtime, + first_prompt=first_prompt or "", + message_count=count, + ) + ) + sessions.sort(key=lambda session: session.mtime, reverse=True) + return sessions + + +@router.get("/history/sessions/{session_id}", response_model=SessionHistoryResponse) +def get_session_history(session_id: str) -> SessionHistoryResponse: + """Load Claude session history for chat replay.""" + sid = _validate_session_id(session_id) + path = _project_history_dir() / f"{sid}.jsonl" + if not path.is_file(): + raise HTTPException(status_code=404, detail="no such session history") + + items: list[dict[str, Any]] = [] + try: + with path.open("r", encoding="utf-8", errors="replace") as fh: + for line in fh: + line = line.strip() + if not line: + continue + try: + entry = json.loads(line) + except json.JSONDecodeError: + continue + if entry.get("isSidechain"): + continue + + entry_type = entry.get("type") + message = entry.get("message") + if entry_type == "user" and isinstance(message, dict): + content = message.get("content") + if isinstance(content, str) and content: + items.append({"kind": "user", "text": content}) + elif entry_type == "assistant" and isinstance(message, dict): + parts = _extract_assistant_parts(message.get("content")) + if parts: + items.append({"kind": "assistant", "parts": parts}) + except OSError as exc: + raise HTTPException(status_code=500, detail=str(exc)) from exc + + _initialized_sessions.add(sid) + return SessionHistoryResponse(session_id=sid, items=items) + + +def _mcp_url(request: Request, session_id: str) -> str: + for route_name in (PUBLIC_MCP_ROUTE_NAME, MCP_ROUTE_NAME): + try: + return str(request.url_for(route_name, session_id=session_id)) + except NoMatchFound: + continue + raise RuntimeError("Studio coding-agent MCP route is not mounted") + + +def _build_claude_argv(session_id: str, message: str, mcp_url: str) -> list[str]: + mcp_config = json.dumps( + { + "mcpServers": { + CLAUDE_MCP_SERVER_NAME: { + "type": "http", + "url": mcp_url, + } + } + } + ) + session_flag = "-r" if session_id in _initialized_sessions else "--session-id" + return [ + "claude", + "-p", + message, + "--output-format", + "stream-json", + "--verbose", + "--mcp-config", + mcp_config, + "--permission-prompt-tool", + f"mcp__{CLAUDE_MCP_SERVER_NAME}__approval_prompt", + session_flag, + session_id, + ] + + +def _claude_env() -> dict[str, str]: + """Build a clean environment so Claude Code uses its own local auth.""" + return { + key: value + for key, value in os.environ.items() + if not key.startswith("ANTHROPIC_") and key != "CLAUDECODE" and not key.startswith("CLAUDE_CODE_") + } + + +def _sse(data: str, event: str | None = None) -> str: + prefix = f"event: {event}\n" if event else "" + return f"{prefix}data: {data}\n\n" + + +async def _request_permission(session_id: str, args: dict[str, Any]) -> dict[str, Any]: + queue = _session_streams.get(session_id) + if queue is None: + return {"behavior": "deny", "message": "no active Studio coding-agent session"} + + request_id = str(uuid.uuid4()) + loop = asyncio.get_running_loop() + future: asyncio.Future[dict[str, Any]] = loop.create_future() + _pending_permissions[request_id] = future + + payload = json.dumps( + { + "request_id": request_id, + "tool_name": args.get("tool_name"), + "input": args.get("input") or {}, + "tool_use_id": args.get("tool_use_id"), + } + ) + await queue.put(("permission_request", payload)) + + try: + decision = await asyncio.wait_for(future, timeout=300) + except asyncio.TimeoutError: + return {"behavior": "deny", "message": "permission request timed out"} + finally: + _pending_permissions.pop(request_id, None) + + if decision.get("approved"): + updated = decision.get("updated_input") + if updated is None: + updated = args.get("input") or {} + return {"behavior": "allow", "updatedInput": updated} + return {"behavior": "deny", "message": decision.get("reason") or "denied by user"} + + +async def _pump_stdout( + proc: asyncio.subprocess.Process, + queue: asyncio.Queue[tuple[str, Any]], +) -> None: + if proc.stdout is None: + await queue.put(("end", None)) + return + + while True: + line = await proc.stdout.readline() + if not line: + break + payload = line.decode(errors="replace").rstrip("\n") + if payload: + await queue.put(("claude", payload)) + await queue.put(("end", None)) + + +async def _pump_stderr(proc: asyncio.subprocess.Process, stderr_chunks: list[str]) -> None: + if proc.stderr is None: + return + + while True: + line = await proc.stderr.readline() + if not line: + break + stderr_chunks.append(line.decode(errors="replace")) + + +async def _terminate_process(proc: asyncio.subprocess.Process) -> None: + if proc.returncode is not None: + return + + proc.terminate() + try: + await asyncio.wait_for(proc.wait(), timeout=2) + except asyncio.TimeoutError: + proc.kill() + await proc.wait() + + +async def _stream_claude(session_id: str, message: str, mcp_url: str) -> AsyncIterator[str]: + if shutil.which("claude") is None: + yield _sse( + json.dumps({"exit_code": None, "stderr": "Claude Code CLI not found on PATH"}), + event="error", + ) + return + + if session_id in _session_streams: + yield _sse( + json.dumps({"exit_code": None, "stderr": "session already has an active stream"}), + event="error", + ) + return + + queue: asyncio.Queue[tuple[str, Any]] = asyncio.Queue() + _session_streams[session_id] = queue + argv = _build_claude_argv(session_id, message, mcp_url) + stderr_chunks: list[str] = [] + stdout_task: asyncio.Task[None] | None = None + stderr_task: asyncio.Task[None] | None = None + + try: + proc = await asyncio.create_subprocess_exec( + *argv, + cwd=str(SERVER_CWD), + env=_claude_env(), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + except OSError as exc: + _session_streams.pop(session_id, None) + yield _sse(json.dumps({"exit_code": None, "stderr": str(exc)}), event="error") + return + + stdout_task = asyncio.create_task(_pump_stdout(proc, queue)) + stderr_task = asyncio.create_task(_pump_stderr(proc, stderr_chunks)) + + try: + while True: + event_type, payload = await queue.get() + if event_type == "end": + break + if event_type == "claude": + yield _sse(payload) + elif event_type == "permission_request": + yield _sse(payload, event="permission_request") + + returncode = await proc.wait() + if stderr_task is not None: + await stderr_task + + if returncode == 0: + _initialized_sessions.add(session_id) + yield _sse("", event="done") + else: + yield _sse( + json.dumps({"exit_code": returncode, "stderr": "".join(stderr_chunks)}), + event="error", + ) + except asyncio.CancelledError: + await _terminate_process(proc) + raise + finally: + _session_streams.pop(session_id, None) + for task in (stdout_task, stderr_task): + if task is not None and not task.done(): + task.cancel() + + +@router.post("/sessions/{session_id}/messages") +async def send_message(session_id: str, body: MessageRequest, request: Request) -> StreamingResponse: + """Send a message to Claude and stream JSON events back to Studio.""" + sid = _validate_session_id(session_id) + return StreamingResponse( + _stream_claude(sid, body.message, _mcp_url(request, sid)), + media_type="text/event-stream", + ) + + +@router.post("/sessions/{session_id}/permissions/{request_id}") +async def resolve_permission(session_id: str, request_id: str, body: PermissionDecision) -> dict[str, bool]: + """Resolve a pending Claude tool permission request.""" + _validate_session_id(session_id) + future = _pending_permissions.get(request_id) + if future is None or future.done(): + raise HTTPException(status_code=404, detail="no such pending permission") + future.set_result(body.model_dump()) + return {"ok": True} + + +@router.post("/mcp/{session_id}", name=MCP_ROUTE_NAME, include_in_schema=False) +async def mcp_endpoint(session_id: str, request: Request) -> Response: + """Minimal MCP endpoint used by Claude's permission-prompt tool.""" + sid = _validate_session_id(session_id) + body = await request.json() + request_id = body.get("id") + + if request_id is None: + return Response(status_code=202) + + method = body.get("method") + params = body.get("params") or {} + + if method == "initialize": + client_protocol = params.get("protocolVersion", "2025-06-18") + return JSONResponse( + { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "protocolVersion": client_protocol, + "capabilities": {"tools": {"listChanged": False}}, + "serverInfo": {"name": "nemo-studio-permissions", "version": "0.1.0"}, + }, + } + ) + + if method == "tools/list": + return JSONResponse( + { + "jsonrpc": "2.0", + "id": request_id, + "result": {"tools": [_APPROVAL_TOOL]}, + } + ) + + if method == "tools/call": + name = params.get("name") + args = params.get("arguments") or {} + if name != "approval_prompt": + return JSONResponse( + { + "jsonrpc": "2.0", + "id": request_id, + "error": {"code": -32601, "message": f"unknown tool: {name}"}, + } + ) + + result = await _request_permission(sid, args) + return JSONResponse( + { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "content": [{"type": "text", "text": json.dumps(result)}], + }, + } + ) + + return JSONResponse( + { + "jsonrpc": "2.0", + "id": request_id, + "error": {"code": -32601, "message": f"method not found: {method}"}, + } + ) diff --git a/services/studio/src/nmp/studio/config.py b/services/studio/src/nmp/studio/config.py index 7270a935ae6..431df2b9498 100644 --- a/services/studio/src/nmp/studio/config.py +++ b/services/studio/src/nmp/studio/config.py @@ -11,7 +11,7 @@ from nmp.common.config import Configuration, create_service_config_class from nmp.studio.env_mappings import ENV_MAPPINGS -from pydantic import BaseModel, Field +from pydantic import BaseModel, ConfigDict, Field logger = logging.getLogger(__name__) @@ -47,6 +47,17 @@ def is_origin_allowed(self, origin: str, same_origin: str | None = None) -> bool return any(fnmatchcase(origin, allowed_origin) for allowed_origin in self.allowed_origins) +class StudioFeatureFlagsConfig(BaseModel): + """Feature flags that affect Studio service behavior.""" + + model_config = ConfigDict(extra="allow") + + coding_agents_enabled: bool = Field( + default=False, + description="Enable Studio's local coding-agent backend bridge.", + ) + + class StudioConfig(create_service_config_class("studio")): # type: ignore[misc] """Configuration for the Studio service. @@ -78,6 +89,10 @@ class StudioConfig(create_service_config_class("studio")): # type: ignore[misc] default_factory=StudioOtelConfig, description="Studio UI OpenTelemetry settings.", ) + feature_flags: StudioFeatureFlagsConfig = Field( + default_factory=StudioFeatureFlagsConfig, + description="Studio feature flags.", + ) @cached_property def global_settings(self) -> dict[str, Any]: diff --git a/services/studio/src/nmp/studio/env_mappings.py b/services/studio/src/nmp/studio/env_mappings.py index bf2e1bc09d4..bdd7d161f3b 100644 --- a/services/studio/src/nmp/studio/env_mappings.py +++ b/services/studio/src/nmp/studio/env_mappings.py @@ -64,6 +64,11 @@ class EnvMapping: config_path="studio.feature_flags.base_models_enabled", default="true", ), + EnvMapping( + marker="STUDIO_UI_VITE_FF_CODING_AGENTS_ENABLED", + config_path="studio.feature_flags.coding_agents_enabled", + default="false", + ), EnvMapping( marker="STUDIO_UI_VITE_FF_CUSTOMIZER_ENABLED", config_path="studio.feature_flags.customizer_enabled", diff --git a/services/studio/src/nmp/studio/service.py b/services/studio/src/nmp/studio/service.py index 885a0ea28e4..a303a22891c 100644 --- a/services/studio/src/nmp/studio/service.py +++ b/services/studio/src/nmp/studio/service.py @@ -13,6 +13,7 @@ from fastapi.responses import HTMLResponse from nmp.common.http_clients import shared_async_http_client from nmp.common.service import RouterConfig, Service +from nmp.studio import coding_agents from nmp.studio.config import StudioConfig from nmp.studio.static_files import SPAStaticFiles from starlette.responses import Response @@ -64,14 +65,23 @@ def title(self) -> str: @property def description(self) -> str: """Service description for OpenAPI docs.""" - return "Serves the NeMo Studio web application" + return "Serves the NeMo Studio web application and local coding-agent bridge" def get_routers(self) -> List[RouterConfig]: """Return routers for the studio service. - The studio service doesn't expose API routers - it serves static files. + Studio exposes API routes for local-only UI integrations in addition to + serving static files. """ - return [] + if not self._coding_agents_enabled(): + return [] + return [ + RouterConfig( + coding_agents.router, + tag="Studio Coding Agents", + description="Local coding-agent bridge endpoints", + ) + ] def configure_app(self, app: FastAPI) -> None: """Configure the platform app with static file mounting. @@ -83,8 +93,18 @@ def configure_app(self, app: FastAPI) -> None: app: The platform's FastAPI application """ self._mount_telemetry_proxy(app) + if self._coding_agents_enabled(): + self._mount_coding_agent_mcp(app) self._mount_static_files(app) + def _coding_agents_enabled(self) -> bool: + """Return whether the exploratory local coding-agent bridge is enabled.""" + return self._get_config().feature_flags.coding_agents_enabled + + def _mount_coding_agent_mcp(self, app: FastAPI) -> None: + """Mount the auth-bypassed MCP callback before the /studio static app.""" + coding_agents.mount_public_mcp_route(app) + def _get_config(self) -> StudioConfig: """Get the studio config, creating a default if none is set. diff --git a/services/studio/tests/unit/test_coding_agents.py b/services/studio/tests/unit/test_coding_agents.py new file mode 100644 index 00000000000..bba364b9aee --- /dev/null +++ b/services/studio/tests/unit/test_coding_agents.py @@ -0,0 +1,240 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for the Studio local coding-agent bridge.""" + +import json +import uuid +from pathlib import Path +from typing import Any + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient +from nmp.studio import coding_agents +from nmp.studio.config import StudioConfig +from nmp.studio.service import StudioService + + +@pytest.fixture(autouse=True) +def reset_coding_agent_state(): + """Reset module-level bridge state between tests.""" + coding_agents._initialized_sessions.clear() + coding_agents._session_streams.clear() + coding_agents._pending_permissions.clear() + yield + coding_agents._initialized_sessions.clear() + coding_agents._session_streams.clear() + coding_agents._pending_permissions.clear() + + +@pytest.fixture +def service_client() -> TestClient: + service = StudioService().with_config(coding_agents_enabled_config()) + return TestClient(service.app) + + +def coding_agents_enabled_config() -> StudioConfig: + return StudioConfig(feature_flags={"coding_agents_enabled": True}) + + +def test_create_session_returns_uuid(service_client: TestClient): + response = service_client.post("/v2/coding-agents/sessions") + + assert response.status_code == 200 + uuid.UUID(response.json()["session_id"]) + + +def test_build_claude_argv_uses_new_session_then_resume_flag(): + session_id = str(uuid.uuid4()) + + argv = coding_agents._build_claude_argv(session_id, "hello", "http://test/mcp") + assert argv[:3] == ["claude", "-p", "hello"] + assert "--output-format" in argv + assert "stream-json" in argv + assert "--permission-prompt-tool" in argv + assert f"mcp__{coding_agents.CLAUDE_MCP_SERVER_NAME}__approval_prompt" in argv + assert "--session-id" in argv + assert session_id in argv + + coding_agents._initialized_sessions.add(session_id) + resumed_argv = coding_agents._build_claude_argv(session_id, "again", "http://test/mcp") + assert "-r" in resumed_argv + assert "--session-id" not in resumed_argv + + +def test_list_and_get_history_sessions( + service_client: TestClient, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +): + workdir = tmp_path / "repo" + projects_dir = tmp_path / "claude-projects" + project_dir = projects_dir / str(workdir).replace("/", "-") + project_dir.mkdir(parents=True) + session_id = str(uuid.uuid4()) + history = project_dir / f"{session_id}.jsonl" + history.write_text( + "\n".join( + [ + json.dumps({"type": "user", "message": {"content": "first prompt"}}), + json.dumps( + { + "type": "assistant", + "message": { + "content": [ + {"type": "thinking", "thinking": "checking"}, + {"type": "text", "text": "done"}, + {"type": "tool_use", "name": "Bash", "input": {"command": "pwd"}}, + ] + }, + } + ), + json.dumps({"type": "user", "isSidechain": True, "message": {"content": "ignored"}}), + "not-json", + ] + ) + ) + + monkeypatch.setattr(coding_agents, "SERVER_CWD", workdir) + monkeypatch.setattr(coding_agents, "CLAUDE_PROJECTS_DIR", projects_dir) + + list_response = service_client.get("/v2/coding-agents/history/sessions") + + assert list_response.status_code == 200 + assert list_response.json() == [ + { + "session_id": session_id, + "mtime": history.stat().st_mtime, + "first_prompt": "first prompt", + "message_count": 1, + } + ] + + history_response = service_client.get(f"/v2/coding-agents/history/sessions/{session_id}") + + assert history_response.status_code == 200 + assert history_response.json() == { + "session_id": session_id, + "items": [ + {"kind": "user", "text": "first prompt"}, + { + "kind": "assistant", + "parts": [ + {"type": "thinking", "thinking": "checking"}, + {"type": "text", "text": "done"}, + {"type": "tool_use", "name": "Bash", "input": {"command": "pwd"}}, + ], + }, + ], + } + assert session_id in coding_agents._initialized_sessions + + +def test_invalid_session_id_returns_400(service_client: TestClient): + response = service_client.get("/v2/coding-agents/history/sessions/not-a-uuid") + + assert response.status_code == 400 + assert response.json()["detail"] == "session_id must be a UUID" + + +def test_mcp_initialize_and_tools_list(service_client: TestClient): + session_id = str(uuid.uuid4()) + + initialize_response = service_client.post( + f"/v2/coding-agents/mcp/{session_id}", + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": {"protocolVersion": "2025-06-18"}, + }, + ) + tools_response = service_client.post( + f"/v2/coding-agents/mcp/{session_id}", + json={"jsonrpc": "2.0", "id": 2, "method": "tools/list"}, + ) + + assert initialize_response.status_code == 200 + assert initialize_response.json()["result"]["serverInfo"]["name"] == "nemo-studio-permissions" + assert tools_response.status_code == 200 + assert tools_response.json()["result"]["tools"][0]["name"] == "approval_prompt" + + +def test_mcp_tools_call_denies_without_active_stream(service_client: TestClient): + session_id = str(uuid.uuid4()) + + response = service_client.post( + f"/v2/coding-agents/mcp/{session_id}", + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": { + "name": "approval_prompt", + "arguments": {"tool_name": "Bash", "input": {"command": "pwd"}}, + }, + }, + ) + + assert response.status_code == 200 + result_text = response.json()["result"]["content"][0]["text"] + assert json.loads(result_text) == { + "behavior": "deny", + "message": "no active Studio coding-agent session", + } + + +def test_platform_route_stream_uses_public_mcp_callback(monkeypatch: pytest.MonkeyPatch): + service = StudioService().with_config(coding_agents_enabled_config()) + app = FastAPI() + app.include_router(service.app.router, prefix="/apis/studio") + service.configure_app(app) + client = TestClient(app) + session_id = str(uuid.uuid4()) + captured: dict[str, Any] = {} + + async def fake_stream(session_id: str, message: str, mcp_url: str): + captured.update({"session_id": session_id, "message": message, "mcp_url": mcp_url}) + yield coding_agents._sse(json.dumps({"type": "system", "subtype": "init"})) + yield coding_agents._sse("", event="done") + + monkeypatch.setattr(coding_agents, "_stream_claude", fake_stream) + + response = client.post( + f"/apis/studio/v2/coding-agents/sessions/{session_id}/messages", + json={"message": "hello"}, + ) + + assert response.status_code == 200 + assert response.headers["content-type"].startswith("text/event-stream") + assert "event: done" in response.text + assert captured == { + "session_id": session_id, + "message": "hello", + "mcp_url": f"http://testserver/studio/api/coding-agents/mcp/{session_id}", + } + + +def test_public_mcp_route_is_mounted_before_static_fallback(): + service = StudioService().with_config(coding_agents_enabled_config()) + app = FastAPI() + service.configure_app(app) + client = TestClient(app) + session_id = str(uuid.uuid4()) + + response = client.post( + f"/studio/api/coding-agents/mcp/{session_id}", + json={"jsonrpc": "2.0", "id": 1, "method": "tools/list"}, + ) + + assert response.status_code == 200 + assert response.json()["result"]["tools"][0]["name"] == "approval_prompt" + + +def test_coding_agent_routes_are_disabled_by_default(): + client = TestClient(StudioService().app) + + response = client.post("/v2/coding-agents/sessions") + + assert response.status_code == 404 diff --git a/services/studio/tests/unit/test_service.py b/services/studio/tests/unit/test_service.py index 695f3121b77..bc1450d0cb3 100644 --- a/services/studio/tests/unit/test_service.py +++ b/services/studio/tests/unit/test_service.py @@ -49,14 +49,21 @@ def test_service_title(self): def test_service_description(self): """Test that the service has the correct description.""" service = StudioService() - assert service.description == "Serves the NeMo Studio web application" + assert service.description == "Serves the NeMo Studio web application and local coding-agent bridge" - def test_get_routers_returns_empty_list(self): - """Test that the service returns no API routers (it serves static files).""" + def test_get_routers_returns_empty_list_by_default(self): + """Test that exploratory coding-agent routes are disabled by default.""" service = StudioService() routers = service.get_routers() assert routers == [] + def test_get_routers_returns_coding_agent_router_when_enabled(self): + """Test that the service exposes the local coding-agent API router when enabled.""" + service = StudioService().with_config(StudioConfig(feature_flags={"coding_agents_enabled": True})) + routers = service.get_routers() + assert len(routers) == 1 + assert routers[0].tag == "Studio Coding Agents" + def test_module_name(self): """Test that the service has the correct module name.""" service = StudioService() diff --git a/web/packages/studio/env/.env.dev.local.sample b/web/packages/studio/env/.env.dev.local.sample index c1d719d4783..a2e220f41f3 100644 --- a/web/packages/studio/env/.env.dev.local.sample +++ b/web/packages/studio/env/.env.dev.local.sample @@ -15,6 +15,7 @@ VITE_AUTH_SCOPE_PREFIX="api://" # See src/constants/featureFlags.ts for available flags and documentation. VITE_FF_AGENTS_ENABLED='true' VITE_FF_BASE_MODELS_ENABLED='true' +VITE_FF_CODING_AGENTS_ENABLED='false' VITE_FF_CUSTOMIZER_ENABLED='false' VITE_FF_DASHBOARD_ENABLED='false' VITE_FF_DATA_DESIGNER_ENABLED='false' diff --git a/web/packages/studio/env/.env.fastapi b/web/packages/studio/env/.env.fastapi index e7849414e6c..f8e2c0907c6 100644 --- a/web/packages/studio/env/.env.fastapi +++ b/web/packages/studio/env/.env.fastapi @@ -19,6 +19,7 @@ VITE_TELEMETRY_ENABLED=STUDIO_UI_VITE_TELEMETRY_ENABLED # Feature Flags (VITE_FF_* prefix) VITE_FF_AGENTS_ENABLED=STUDIO_UI_VITE_FF_AGENTS_ENABLED +VITE_FF_CODING_AGENTS_ENABLED=STUDIO_UI_VITE_FF_CODING_AGENTS_ENABLED VITE_FF_CUSTOMIZER_ENABLED=STUDIO_UI_VITE_FF_CUSTOMIZER_ENABLED VITE_FF_DASHBOARD_ENABLED=STUDIO_UI_VITE_FF_DASHBOARD_ENABLED VITE_FF_DATA_DESIGNER_ENABLED=STUDIO_UI_VITE_FF_DATA_DESIGNER_ENABLED diff --git a/web/packages/studio/src/constants/environment.ts b/web/packages/studio/src/constants/environment.ts index e0476dab07d..8a7aeccd92a 100644 --- a/web/packages/studio/src/constants/environment.ts +++ b/web/packages/studio/src/constants/environment.ts @@ -32,6 +32,7 @@ export const PLATFORM_BASE_URL = resolveBrowserBaseUrl(getEnvVar('VITE_PLATFORM_ // distinguish that logic from code that calls the URL itself export const AGENTS_ENABLED = featureFlags.agentsEnabled !== false; export const BASE_MODELS_ENABLED = featureFlags.baseModelsEnabled !== false; +export const CODING_AGENTS_ENABLED = featureFlags.codingAgentsEnabled !== false; export const CUSTOMIZER_ENABLED = featureFlags.customizerEnabled !== false; export const DASHBOARD_ENABLED = featureFlags.dashboardEnabled !== false; export const DATA_DESIGNER_ENABLED = featureFlags.dataDesignerEnabled !== false; diff --git a/web/packages/studio/src/constants/featureFlags/featureFlags.ts b/web/packages/studio/src/constants/featureFlags/featureFlags.ts index e42c3ce155f..19458ea1f76 100644 --- a/web/packages/studio/src/constants/featureFlags/featureFlags.ts +++ b/web/packages/studio/src/constants/featureFlags/featureFlags.ts @@ -56,6 +56,7 @@ import { z } from 'zod'; export const flagDefinitions = { agentsEnabled: previewFlag('VITE_FF_AGENTS_ENABLED', true), baseModelsEnabled: previewFlag('VITE_FF_BASE_MODELS_ENABLED', true), + codingAgentsEnabled: previewFlag('VITE_FF_CODING_AGENTS_ENABLED', false), customizerEnabled: previewFlag('VITE_FF_CUSTOMIZER_ENABLED', false), dashboardEnabled: previewFlag('VITE_FF_DASHBOARD_ENABLED', false), dataDesignerEnabled: previewFlag('VITE_FF_DATA_DESIGNER_ENABLED'), From d023cee4d6e7a42a5f4380f0a9c920b773effe7a Mon Sep 17 00:00:00 2001 From: Danielle Ali <44468613+dmariali@users.noreply.github.com> Date: Fri, 29 May 2026 13:44:58 -0400 Subject: [PATCH 2/8] cleanup Signed-off-by: Danielle Ali <44468613+dmariali@users.noreply.github.com> --- .../studio/src/nmp/studio/coding_agents.py | 24 +++++-- .../studio/tests/unit/test_coding_agents.py | 71 ++++++++++++++++++- 2 files changed, 88 insertions(+), 7 deletions(-) diff --git a/services/studio/src/nmp/studio/coding_agents.py b/services/studio/src/nmp/studio/coding_agents.py index bdfe4586f34..2839dfd8d3c 100644 --- a/services/studio/src/nmp/studio/coding_agents.py +++ b/services/studio/src/nmp/studio/coding_agents.py @@ -69,7 +69,7 @@ class SessionHistoryResponse(BaseModel): _initialized_sessions: set[str] = set() _session_streams: dict[str, asyncio.Queue[tuple[str, Any]]] = {} -_pending_permissions: dict[str, asyncio.Future[dict[str, Any]]] = {} +_pending_permissions: dict[str, tuple[str, asyncio.Future[dict[str, Any]]]] = {} _APPROVAL_TOOL = { @@ -303,7 +303,7 @@ async def _request_permission(session_id: str, args: dict[str, Any]) -> dict[str request_id = str(uuid.uuid4()) loop = asyncio.get_running_loop() future: asyncio.Future[dict[str, Any]] = loop.create_future() - _pending_permissions[request_id] = future + _pending_permissions[request_id] = (session_id, future) payload = json.dumps( { @@ -454,9 +454,12 @@ async def send_message(session_id: str, body: MessageRequest, request: Request) @router.post("/sessions/{session_id}/permissions/{request_id}") async def resolve_permission(session_id: str, request_id: str, body: PermissionDecision) -> dict[str, bool]: """Resolve a pending Claude tool permission request.""" - _validate_session_id(session_id) - future = _pending_permissions.get(request_id) - if future is None or future.done(): + sid = _validate_session_id(session_id) + pending = _pending_permissions.get(request_id) + if pending is None: + raise HTTPException(status_code=404, detail="no such pending permission") + pending_session_id, future = pending + if pending_session_id != sid or future.done(): raise HTTPException(status_code=404, detail="no such pending permission") future.set_result(body.model_dump()) return {"ok": True} @@ -466,13 +469,22 @@ async def resolve_permission(session_id: str, request_id: str, body: PermissionD async def mcp_endpoint(session_id: str, request: Request) -> Response: """Minimal MCP endpoint used by Claude's permission-prompt tool.""" sid = _validate_session_id(session_id) - body = await request.json() + try: + body = await request.json() + except ValueError: + return JSONResponse(status_code=400, content={"detail": "invalid JSON body"}) + if not isinstance(body, dict): + return JSONResponse(status_code=400, content={"detail": "JSON body must be an object"}) + request_id = body.get("id") if request_id is None: return Response(status_code=202) method = body.get("method") + raw_params = body.get("params") + if raw_params is not None and not isinstance(raw_params, dict): + return JSONResponse(status_code=400, content={"detail": "JSON-RPC params must be an object"}) params = body.get("params") or {} if method == "initialize": diff --git a/services/studio/tests/unit/test_coding_agents.py b/services/studio/tests/unit/test_coding_agents.py index bba364b9aee..8a22309f141 100644 --- a/services/studio/tests/unit/test_coding_agents.py +++ b/services/studio/tests/unit/test_coding_agents.py @@ -3,13 +3,14 @@ """Unit tests for the Studio local coding-agent bridge.""" +import asyncio import json import uuid from pathlib import Path from typing import Any import pytest -from fastapi import FastAPI +from fastapi import FastAPI, HTTPException from fastapi.testclient import TestClient from nmp.studio import coding_agents from nmp.studio.config import StudioConfig @@ -161,6 +162,40 @@ def test_mcp_initialize_and_tools_list(service_client: TestClient): assert tools_response.json()["result"]["tools"][0]["name"] == "approval_prompt" +def test_mcp_rejects_malformed_json(service_client: TestClient): + session_id = str(uuid.uuid4()) + + response = service_client.post( + f"/v2/coding-agents/mcp/{session_id}", + content="{", + headers={"content-type": "application/json"}, + ) + + assert response.status_code == 400 + assert response.json()["detail"] == "invalid JSON body" + + +def test_mcp_rejects_non_object_json(service_client: TestClient): + session_id = str(uuid.uuid4()) + + response = service_client.post(f"/v2/coding-agents/mcp/{session_id}", json=[]) + + assert response.status_code == 400 + assert response.json()["detail"] == "JSON body must be an object" + + +def test_mcp_rejects_non_object_params(service_client: TestClient): + session_id = str(uuid.uuid4()) + + response = service_client.post( + f"/v2/coding-agents/mcp/{session_id}", + json={"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": []}, + ) + + assert response.status_code == 400 + assert response.json()["detail"] == "JSON-RPC params must be an object" + + def test_mcp_tools_call_denies_without_active_stream(service_client: TestClient): session_id = str(uuid.uuid4()) @@ -185,6 +220,40 @@ def test_mcp_tools_call_denies_without_active_stream(service_client: TestClient) } +async def test_resolve_permission_rejects_cross_session_request(): + owner_session_id = str(uuid.uuid4()) + other_session_id = str(uuid.uuid4()) + request_id = str(uuid.uuid4()) + future = asyncio.get_running_loop().create_future() + coding_agents._pending_permissions[request_id] = (owner_session_id, future) + + with pytest.raises(HTTPException) as exc_info: + await coding_agents.resolve_permission( + other_session_id, + request_id, + coding_agents.PermissionDecision(approved=True), + ) + + assert exc_info.value.status_code == 404 + assert not future.done() + + +async def test_resolve_permission_sets_result_for_owning_session(): + session_id = str(uuid.uuid4()) + request_id = str(uuid.uuid4()) + future = asyncio.get_running_loop().create_future() + coding_agents._pending_permissions[request_id] = (session_id, future) + + response = await coding_agents.resolve_permission( + session_id, + request_id, + coding_agents.PermissionDecision(approved=True), + ) + + assert response == {"ok": True} + assert future.result() == {"approved": True, "reason": None, "updated_input": None} + + def test_platform_route_stream_uses_public_mcp_callback(monkeypatch: pytest.MonkeyPatch): service = StudioService().with_config(coding_agents_enabled_config()) app = FastAPI() From 841a229b92e7d8fe2cafc7a85842b21141cfe3f5 Mon Sep 17 00:00:00 2001 From: Danielle Ali <44468613+dmariali@users.noreply.github.com> Date: Mon, 1 Jun 2026 12:24:15 -0400 Subject: [PATCH 3/8] codeQL issue Signed-off-by: Danielle Ali <44468613+dmariali@users.noreply.github.com> --- services/studio/src/nmp/studio/coding_agents.py | 8 ++++++-- services/studio/tests/unit/test_coding_agents.py | 16 ++++++++++++++++ 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/services/studio/src/nmp/studio/coding_agents.py b/services/studio/src/nmp/studio/coding_agents.py index 2839dfd8d3c..e38b1504337 100644 --- a/services/studio/src/nmp/studio/coding_agents.py +++ b/services/studio/src/nmp/studio/coding_agents.py @@ -401,9 +401,13 @@ async def _stream_claude(session_id: str, message: str, mcp_url: str) -> AsyncIt stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) - except OSError as exc: + except OSError: + logger.exception("Failed to start Claude Code subprocess for session %s", session_id) _session_streams.pop(session_id, None) - yield _sse(json.dumps({"exit_code": None, "stderr": str(exc)}), event="error") + yield _sse( + json.dumps({"exit_code": None, "stderr": "Failed to start Claude Code process"}), + event="error", + ) return stdout_task = asyncio.create_task(_pump_stdout(proc, queue)) diff --git a/services/studio/tests/unit/test_coding_agents.py b/services/studio/tests/unit/test_coding_agents.py index 9f20f38691b..9a4fbcd5d5d 100644 --- a/services/studio/tests/unit/test_coding_agents.py +++ b/services/studio/tests/unit/test_coding_agents.py @@ -139,6 +139,22 @@ def test_invalid_session_id_returns_400(service_client: TestClient): assert response.json()["detail"] == "session_id must be a UUID" +async def test_stream_claude_hides_startup_oserror(monkeypatch: pytest.MonkeyPatch): + session_id = str(uuid.uuid4()) + + async def fail_start(*args: Any, **kwargs: Any): + raise OSError("secret local path") + + monkeypatch.setattr(coding_agents.shutil, "which", lambda name: "/usr/bin/claude") + monkeypatch.setattr(coding_agents.asyncio, "create_subprocess_exec", fail_start) + + chunks = [chunk async for chunk in coding_agents._stream_claude(session_id, "hello", "http://test/mcp")] + + assert chunks == ['event: error\ndata: {"exit_code": null, "stderr": "Failed to start Claude Code process"}\n\n'] + assert "secret local path" not in chunks[0] + assert session_id not in coding_agents._session_streams + + def test_mcp_initialize_and_tools_list(service_client: TestClient): session_id = str(uuid.uuid4()) From bb1794906c601ccfeae79212de5b6d32a6a62a65 Mon Sep 17 00:00:00 2001 From: Danielle Ali <44468613+dmariali@users.noreply.github.com> Date: Mon, 1 Jun 2026 13:14:57 -0400 Subject: [PATCH 4/8] remove extra flagging logic Signed-off-by: Danielle Ali <44468613+dmariali@users.noreply.github.com> --- docs/set-up/config-reference.md | 2 +- services/studio/src/nmp/studio/config.py | 4 ++-- services/studio/src/nmp/studio/service.py | 9 +-------- services/studio/tests/unit/test_coding_agents.py | 16 ++++++---------- services/studio/tests/unit/test_service.py | 10 ++-------- 5 files changed, 12 insertions(+), 29 deletions(-) diff --git a/docs/set-up/config-reference.md b/docs/set-up/config-reference.md index 5bd6f0e5a3f..d267036b0a5 100644 --- a/docs/set-up/config-reference.md +++ b/docs/set-up/config-reference.md @@ -877,6 +877,6 @@ studio: - http://localhost:* # Studio feature flags. feature_flags: - # Enable Studio's local coding-agent backend bridge. | default: False + # Show Studio's local coding-agent UI. | default: False coding_agent_studio_enabled: false ``` diff --git a/services/studio/src/nmp/studio/config.py b/services/studio/src/nmp/studio/config.py index e2b0a5e4a83..27a14589c58 100644 --- a/services/studio/src/nmp/studio/config.py +++ b/services/studio/src/nmp/studio/config.py @@ -48,13 +48,13 @@ def is_origin_allowed(self, origin: str, same_origin: str | None = None) -> bool class StudioFeatureFlagsConfig(BaseModel): - """Feature flags that affect Studio service behavior.""" + """Feature flags injected into the Studio UI.""" model_config = ConfigDict(extra="allow") coding_agent_studio_enabled: bool = Field( default=False, - description="Enable Studio's local coding-agent backend bridge.", + description="Show Studio's local coding-agent UI.", ) diff --git a/services/studio/src/nmp/studio/service.py b/services/studio/src/nmp/studio/service.py index f1684816e3e..51a18d3a445 100644 --- a/services/studio/src/nmp/studio/service.py +++ b/services/studio/src/nmp/studio/service.py @@ -73,8 +73,6 @@ def get_routers(self) -> List[RouterConfig]: Studio exposes API routes for local-only UI integrations in addition to serving static files. """ - if not self._coding_agents_enabled(): - return [] return [ RouterConfig( coding_agents.router, @@ -93,14 +91,9 @@ def configure_app(self, app: FastAPI) -> None: app: The platform's FastAPI application """ self._mount_telemetry_proxy(app) - if self._coding_agents_enabled(): - self._mount_coding_agent_mcp(app) + self._mount_coding_agent_mcp(app) self._mount_static_files(app) - def _coding_agents_enabled(self) -> bool: - """Return whether the exploratory local coding-agent bridge is enabled.""" - return self._get_config().feature_flags.coding_agent_studio_enabled - def _mount_coding_agent_mcp(self, app: FastAPI) -> None: """Mount the auth-bypassed MCP callback before the /studio static app.""" coding_agents.mount_public_mcp_route(app) diff --git a/services/studio/tests/unit/test_coding_agents.py b/services/studio/tests/unit/test_coding_agents.py index 9a4fbcd5d5d..37007943ad6 100644 --- a/services/studio/tests/unit/test_coding_agents.py +++ b/services/studio/tests/unit/test_coding_agents.py @@ -13,7 +13,6 @@ from fastapi import FastAPI, HTTPException from fastapi.testclient import TestClient from nmp.studio import coding_agents -from nmp.studio.config import StudioConfig from nmp.studio.service import StudioService @@ -31,14 +30,10 @@ def reset_coding_agent_state(): @pytest.fixture def service_client() -> TestClient: - service = StudioService().with_config(coding_agents_enabled_config()) + service = StudioService() return TestClient(service.app) -def coding_agents_enabled_config() -> StudioConfig: - return StudioConfig(feature_flags={"coding_agent_studio_enabled": True}) - - def test_create_session_returns_uuid(service_client: TestClient): response = service_client.post("/v2/coding-agents/sessions") @@ -271,7 +266,7 @@ async def test_resolve_permission_sets_result_for_owning_session(): def test_platform_route_stream_uses_public_mcp_callback(monkeypatch: pytest.MonkeyPatch): - service = StudioService().with_config(coding_agents_enabled_config()) + service = StudioService() app = FastAPI() app.include_router(service.app.router, prefix="/apis/studio") service.configure_app(app) @@ -302,7 +297,7 @@ async def fake_stream(session_id: str, message: str, mcp_url: str): def test_public_mcp_route_is_mounted_before_static_fallback(): - service = StudioService().with_config(coding_agents_enabled_config()) + service = StudioService() app = FastAPI() service.configure_app(app) client = TestClient(app) @@ -317,9 +312,10 @@ def test_public_mcp_route_is_mounted_before_static_fallback(): assert response.json()["result"]["tools"][0]["name"] == "approval_prompt" -def test_coding_agent_routes_are_disabled_by_default(): +def test_coding_agent_routes_are_available_by_default(): client = TestClient(StudioService().app) response = client.post("/v2/coding-agents/sessions") - assert response.status_code == 404 + assert response.status_code == 200 + uuid.UUID(response.json()["session_id"]) diff --git a/services/studio/tests/unit/test_service.py b/services/studio/tests/unit/test_service.py index d4d69dac86a..33c5e5df9cc 100644 --- a/services/studio/tests/unit/test_service.py +++ b/services/studio/tests/unit/test_service.py @@ -51,16 +51,10 @@ def test_service_description(self): service = StudioService() assert service.description == "Serves the NeMo Studio web application and local coding-agent bridge" - def test_get_routers_returns_empty_list_by_default(self): - """Test that exploratory coding-agent routes are disabled by default.""" + def test_get_routers_returns_coding_agent_router(self): + """Test that the service exposes the local coding-agent API router.""" service = StudioService() routers = service.get_routers() - assert routers == [] - - def test_get_routers_returns_coding_agent_router_when_enabled(self): - """Test that the service exposes the local coding-agent API router when enabled.""" - service = StudioService().with_config(StudioConfig(feature_flags={"coding_agent_studio_enabled": True})) - routers = service.get_routers() assert len(routers) == 1 assert routers[0].tag == "Studio Coding Agents" From e01fd5d70a21be07048c7087b4085e95c76caf22 Mon Sep 17 00:00:00 2001 From: Danielle Ali <44468613+dmariali@users.noreply.github.com> Date: Mon, 1 Jun 2026 13:24:05 -0400 Subject: [PATCH 5/8] cleanup flagging logic Signed-off-by: Danielle Ali <44468613+dmariali@users.noreply.github.com> --- docs/set-up/config-reference.md | 4 +--- services/studio/src/nmp/studio/config.py | 5 ----- 2 files changed, 1 insertion(+), 8 deletions(-) diff --git a/docs/set-up/config-reference.md b/docs/set-up/config-reference.md index d267036b0a5..b3dd267a416 100644 --- a/docs/set-up/config-reference.md +++ b/docs/set-up/config-reference.md @@ -876,7 +876,5 @@ studio: - http://localhost - http://localhost:* # Studio feature flags. - feature_flags: - # Show Studio's local coding-agent UI. | default: False - coding_agent_studio_enabled: false + feature_flags: {} ``` diff --git a/services/studio/src/nmp/studio/config.py b/services/studio/src/nmp/studio/config.py index 27a14589c58..a7994d7fbf1 100644 --- a/services/studio/src/nmp/studio/config.py +++ b/services/studio/src/nmp/studio/config.py @@ -52,11 +52,6 @@ class StudioFeatureFlagsConfig(BaseModel): model_config = ConfigDict(extra="allow") - coding_agent_studio_enabled: bool = Field( - default=False, - description="Show Studio's local coding-agent UI.", - ) - class StudioConfig(create_service_config_class("studio")): # type: ignore[misc] """Configuration for the Studio service. From 061c135edef8e4f71bf4cf2106f1d8c068ce76c1 Mon Sep 17 00:00:00 2001 From: Danielle Ali <44468613+dmariali@users.noreply.github.com> Date: Mon, 1 Jun 2026 13:28:26 -0400 Subject: [PATCH 6/8] Update config-reference.md Signed-off-by: Danielle Ali <44468613+dmariali@users.noreply.github.com> --- docs/set-up/config-reference.md | 2 -- 1 file changed, 2 deletions(-) diff --git a/docs/set-up/config-reference.md b/docs/set-up/config-reference.md index b3dd267a416..93f2e2f6a8e 100644 --- a/docs/set-up/config-reference.md +++ b/docs/set-up/config-reference.md @@ -875,6 +875,4 @@ studio: allowed_origins: - http://localhost - http://localhost:* - # Studio feature flags. - feature_flags: {} ``` From bf2d11253432ad3af4d82332e639cd3333cc4d92 Mon Sep 17 00:00:00 2001 From: Danielle Ali <44468613+dmariali@users.noreply.github.com> Date: Mon, 1 Jun 2026 14:02:30 -0400 Subject: [PATCH 7/8] Update config.py Signed-off-by: Danielle Ali <44468613+dmariali@users.noreply.github.com> --- services/studio/src/nmp/studio/config.py | 12 +----------- 1 file changed, 1 insertion(+), 11 deletions(-) diff --git a/services/studio/src/nmp/studio/config.py b/services/studio/src/nmp/studio/config.py index a7994d7fbf1..7270a935ae6 100644 --- a/services/studio/src/nmp/studio/config.py +++ b/services/studio/src/nmp/studio/config.py @@ -11,7 +11,7 @@ from nmp.common.config import Configuration, create_service_config_class from nmp.studio.env_mappings import ENV_MAPPINGS -from pydantic import BaseModel, ConfigDict, Field +from pydantic import BaseModel, Field logger = logging.getLogger(__name__) @@ -47,12 +47,6 @@ def is_origin_allowed(self, origin: str, same_origin: str | None = None) -> bool return any(fnmatchcase(origin, allowed_origin) for allowed_origin in self.allowed_origins) -class StudioFeatureFlagsConfig(BaseModel): - """Feature flags injected into the Studio UI.""" - - model_config = ConfigDict(extra="allow") - - class StudioConfig(create_service_config_class("studio")): # type: ignore[misc] """Configuration for the Studio service. @@ -84,10 +78,6 @@ class StudioConfig(create_service_config_class("studio")): # type: ignore[misc] default_factory=StudioOtelConfig, description="Studio UI OpenTelemetry settings.", ) - feature_flags: StudioFeatureFlagsConfig = Field( - default_factory=StudioFeatureFlagsConfig, - description="Studio feature flags.", - ) @cached_property def global_settings(self) -> dict[str, Any]: From f027dc12dac9d3f173d48b38ee90cc69639bf1ca Mon Sep 17 00:00:00 2001 From: Danielle Ali <44468613+dmariali@users.noreply.github.com> Date: Tue, 2 Jun 2026 11:14:27 -0400 Subject: [PATCH 8/8] Matt's comments Signed-off-by: Danielle Ali <44468613+dmariali@users.noreply.github.com> --- .../studio/src/nmp/studio/coding_agents.py | 125 +++++++++++++++--- .../studio/tests/unit/test_coding_agents.py | 33 ++++- 2 files changed, 142 insertions(+), 16 deletions(-) diff --git a/services/studio/src/nmp/studio/coding_agents.py b/services/studio/src/nmp/studio/coding_agents.py index e38b1504337..9c61826e06b 100644 --- a/services/studio/src/nmp/studio/coding_agents.py +++ b/services/studio/src/nmp/studio/coding_agents.py @@ -10,6 +10,8 @@ import shutil import uuid from collections.abc import AsyncIterator +from dataclasses import dataclass +from dataclasses import field as dataclass_field from pathlib import Path from typing import Any @@ -58,6 +60,9 @@ class HistorySessionResponse(BaseModel): mtime: float first_prompt: str message_count: int + token_count: int + tool_call_count: int + tool_calls: list[str] class SessionHistoryResponse(BaseModel): @@ -72,6 +77,17 @@ class SessionHistoryResponse(BaseModel): _pending_permissions: dict[str, tuple[str, asyncio.Future[dict[str, Any]]]] = {} +@dataclass +class HistorySummary: + """Aggregated metadata from a Claude session history file.""" + + first_prompt: str | None = None + message_count: int = 0 + token_count: int = 0 + tool_call_count: int = 0 + tool_calls: list[str] = dataclass_field(default_factory=list) + + _APPROVAL_TOOL = { "name": "approval_prompt", "description": "Ask the human operator whether a tool call should be allowed.", @@ -110,9 +126,68 @@ def _project_history_dir() -> Path: return CLAUDE_PROJECTS_DIR / encoded -def _first_user_prompt_and_count(path: Path) -> tuple[str | None, int]: - first_prompt: str | None = None - count = 0 +_TOKEN_USAGE_FIELDS = ( + "input_tokens", + "cache_creation_input_tokens", + "cache_read_input_tokens", + "output_tokens", +) + + +def _int_metric(value: Any) -> int: + return value if isinstance(value, int) and not isinstance(value, bool) else 0 + + +def _usage_token_count(usage: Any) -> int: + if not isinstance(usage, dict): + return 0 + return sum(_int_metric(usage.get(field)) for field in _TOKEN_USAGE_FIELDS) + + +def _tool_result_token_count(tool_result: Any) -> int: + if not isinstance(tool_result, dict): + return 0 + total_tokens = _int_metric(tool_result.get("totalTokens")) + if total_tokens: + return total_tokens + return _usage_token_count(tool_result.get("usage")) + + +def _usage_identity(entry: dict[str, Any], message: dict[str, Any]) -> tuple[str, str] | None: + request_id = entry.get("requestId") + message_id = message.get("id") + if not isinstance(request_id, str) and not isinstance(message_id, str): + return None + return (request_id if isinstance(request_id, str) else "", message_id if isinstance(message_id, str) else "") + + +def _append_tool_call(summary: HistorySummary, tool_name: str) -> None: + summary.tool_call_count += 1 + if tool_name not in summary.tool_calls: + summary.tool_calls.append(tool_name) + + +def _record_assistant_tool_calls( + summary: HistorySummary, + message: dict[str, Any], + seen_tool_use_ids: set[str], +) -> None: + for part in message.get("content") or []: + if not isinstance(part, dict) or part.get("type") != "tool_use": + continue + tool_use_id = part.get("id") + if isinstance(tool_use_id, str): + if tool_use_id in seen_tool_use_ids: + continue + seen_tool_use_ids.add(tool_use_id) + tool_name = part.get("name") + _append_tool_call(summary, tool_name if isinstance(tool_name, str) and tool_name else "tool") + + +def _summarize_history_session(path: Path) -> HistorySummary: + summary = HistorySummary() + seen_usage_events: set[tuple[str, str]] = set() + seen_tool_use_ids: set[str] = set() try: with path.open("r", encoding="utf-8", errors="replace") as fh: for line in fh: @@ -123,17 +198,34 @@ def _first_user_prompt_and_count(path: Path) -> tuple[str | None, int]: entry = json.loads(line) except json.JSONDecodeError: continue - if entry.get("type") != "user" or entry.get("isSidechain"): + if entry.get("isSidechain"): continue - content = (entry.get("message") or {}).get("content") - if not isinstance(content, str): + if not isinstance(entry, dict): continue - count += 1 - if first_prompt is None: - first_prompt = content + + message = entry.get("message") + if isinstance(message, dict): + usage_identity = _usage_identity(entry, message) + if usage_identity is None or usage_identity not in seen_usage_events: + summary.token_count += _usage_token_count(message.get("usage")) + if usage_identity is not None: + seen_usage_events.add(usage_identity) + + summary.token_count += _tool_result_token_count(entry.get("toolUseResult")) + + entry_type = entry.get("type") + if entry_type == "assistant" and isinstance(message, dict): + _record_assistant_tool_calls(summary, message, seen_tool_use_ids) + elif entry_type == "user" and isinstance(message, dict): + content = message.get("content") + if not isinstance(content, str): + continue + summary.message_count += 1 + if summary.first_prompt is None: + summary.first_prompt = content except OSError: - return None, 0 - return first_prompt, count + return HistorySummary() + return summary def _extract_assistant_parts(content: Any) -> list[dict[str, Any]]: @@ -184,8 +276,8 @@ def list_history_sessions() -> list[HistorySessionResponse]: except ValueError: continue - first_prompt, count = _first_user_prompt_and_count(history_file) - if count == 0: + summary = _summarize_history_session(history_file) + if summary.message_count == 0: continue try: @@ -197,8 +289,11 @@ def list_history_sessions() -> list[HistorySessionResponse]: HistorySessionResponse( session_id=history_file.stem, mtime=mtime, - first_prompt=first_prompt or "", - message_count=count, + first_prompt=summary.first_prompt or "", + message_count=summary.message_count, + token_count=summary.token_count, + tool_call_count=summary.tool_call_count, + tool_calls=summary.tool_calls, ) ) sessions.sort(key=lambda session: session.mtime, reverse=True) diff --git a/services/studio/tests/unit/test_coding_agents.py b/services/studio/tests/unit/test_coding_agents.py index 37007943ad6..778e2dd6976 100644 --- a/services/studio/tests/unit/test_coding_agents.py +++ b/services/studio/tests/unit/test_coding_agents.py @@ -78,12 +78,40 @@ def test_list_and_get_history_sessions( { "type": "assistant", "message": { + "id": "msg_1", "content": [ {"type": "thinking", "thinking": "checking"}, {"type": "text", "text": "done"}, - {"type": "tool_use", "name": "Bash", "input": {"command": "pwd"}}, + { + "type": "tool_use", + "id": "toolu_1", + "name": "Bash", + "input": {"command": "pwd"}, + }, + ], + "usage": { + "input_tokens": 10, + "cache_creation_input_tokens": 2, + "cache_read_input_tokens": 3, + "output_tokens": 4, + }, + }, + "requestId": "req_1", + } + ), + json.dumps( + { + "type": "user", + "message": { + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_1", + "content": "done", + } ] }, + "toolUseResult": {"totalTokens": 11}, } ), json.dumps({"type": "user", "isSidechain": True, "message": {"content": "ignored"}}), @@ -104,6 +132,9 @@ def test_list_and_get_history_sessions( "mtime": history.stat().st_mtime, "first_prompt": "first prompt", "message_count": 1, + "token_count": 30, + "tool_call_count": 1, + "tool_calls": ["Bash"], } ]