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 0000000000..9c61826e06 --- /dev/null +++ b/services/studio/src/nmp/studio/coding_agents.py @@ -0,0 +1,641 @@ +# 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 dataclasses import dataclass +from dataclasses import field as dataclass_field +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 + token_count: int + tool_call_count: int + tool_calls: list[str] + + +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, 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.", + "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 + + +_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: + line = line.strip() + if not line: + continue + try: + entry = json.loads(line) + except json.JSONDecodeError: + continue + if entry.get("isSidechain"): + continue + if not isinstance(entry, dict): + continue + + 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 HistorySummary() + return summary + + +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 + + summary = _summarize_history_session(history_file) + if summary.message_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=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) + 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] = (session_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: + 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": "Failed to start Claude Code process"}), + 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.""" + 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} + + +@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) + 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": + 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/service.py b/services/studio/src/nmp/studio/service.py index 885a0ea28e..51a18d3a44 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,21 @@ 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 [] + 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 +91,13 @@ def configure_app(self, app: FastAPI) -> None: app: The platform's FastAPI application """ self._mount_telemetry_proxy(app) + self._mount_coding_agent_mcp(app) self._mount_static_files(app) + 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 0000000000..778e2dd697 --- /dev/null +++ b/services/studio/tests/unit/test_coding_agents.py @@ -0,0 +1,352 @@ +# 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 asyncio +import json +import uuid +from pathlib import Path +from typing import Any + +import pytest +from fastapi import FastAPI, HTTPException +from fastapi.testclient import TestClient +from nmp.studio import coding_agents +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() + return TestClient(service.app) + + +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": { + "id": "msg_1", + "content": [ + {"type": "thinking", "thinking": "checking"}, + {"type": "text", "text": "done"}, + { + "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"}}), + "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, + "token_count": 30, + "tool_call_count": 1, + "tool_calls": ["Bash"], + } + ] + + 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" + + +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()) + + 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_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()) + + 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", + } + + +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() + 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() + 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_available_by_default(): + client = TestClient(StudioService().app) + + response = client.post("/v2/coding-agents/sessions") + + 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 695f3121b7..33c5e5df9c 100644 --- a/services/studio/tests/unit/test_service.py +++ b/services/studio/tests/unit/test_service.py @@ -49,13 +49,14 @@ 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_coding_agent_router(self): + """Test that the service exposes the local coding-agent API router.""" service = StudioService() routers = service.get_routers() - assert 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."""