From 1fef8df9b5168f15afdc91061f5ced3b15d2dded Mon Sep 17 00:00:00 2001 From: Henrique Tolentino Date: Tue, 4 Aug 2026 11:03:26 -0400 Subject: [PATCH 1/4] Re-introduce tool usage information to the Nemo Copilot Signed-off-by: Henrique Tolentino --- .../src/nemo_agent/wrapper.py | 20 ++-- .../nemo-agent-local/tests/test_nemo_agent.py | 8 +- .../studio/src/nmp/studio/coding_agents.py | 108 +++++++++++++++++- .../studio/tests/unit/test_coding_agents.py | 22 +++- 4 files changed, 143 insertions(+), 15 deletions(-) diff --git a/agents/nemo-agent-local/src/nemo_agent/wrapper.py b/agents/nemo-agent-local/src/nemo_agent/wrapper.py index 2e81f0070d..5151c54149 100644 --- a/agents/nemo-agent-local/src/nemo_agent/wrapper.py +++ b/agents/nemo-agent-local/src/nemo_agent/wrapper.py @@ -75,6 +75,7 @@ from nat.cli.register_workflow import register_function from nat.data_models.api_server import ChatRequest, ChatResponse, ChatResponseChunk, Usage from nat.data_models.function import FunctionBaseConfig +from nat.plugins.langchain.callback_handler import LangchainProfilerHandler from nemo_agent.register import create_nemo_agent from pydantic import BaseModel, ConfigDict, Field, computed_field, model_validator @@ -241,13 +242,18 @@ def _build_state(value: NemoAgentWrapperInput) -> dict[str, Any]: @staticmethod def _invocation_config(value: NemoAgentWrapperInput) -> RunnableConfig: - if value.studio_session_id is None: - return {} - return { - "configurable": { - "studio_session_id": str(value.studio_session_id), - } - } + config: RunnableConfig = {} + # Attach NAT's LangChain profiler so tool/LLM calls emit IntermediateSteps, + # which NAT surfaces as ``intermediate_data:`` stream events (and telemetry). + # Construction binds to the request-scoped step manager; guard so calls + # outside a NAT context (e.g. unit tests) degrade gracefully. + try: + config["callbacks"] = [LangchainProfilerHandler()] + except Exception: + logger.debug("LangchainProfilerHandler unavailable; tool-call trace disabled", exc_info=True) + if value.studio_session_id is not None: + config["configurable"] = {"studio_session_id": str(value.studio_session_id)} + return config @staticmethod def _has_tool_calls(message: AIMessage | dict[str, Any]) -> bool: diff --git a/agents/nemo-agent-local/tests/test_nemo_agent.py b/agents/nemo-agent-local/tests/test_nemo_agent.py index 8c1e34bf36..1a92621ffb 100644 --- a/agents/nemo-agent-local/tests/test_nemo_agent.py +++ b/agents/nemo-agent-local/tests/test_nemo_agent.py @@ -866,7 +866,10 @@ def test_chat_request_converts_trusted_session_to_invocation_config(self): value = NemoAgentWrapperFunction.convert_chat_request(request) - assert NemoAgentWrapperFunction._invocation_config(value) == TRUSTED_CONFIG + # ``callbacks`` may also carry the NAT profiler handler; assert the + # session config specifically rather than the whole dict. + config = NemoAgentWrapperFunction._invocation_config(value) + assert config["configurable"] == TRUSTED_CONFIG["configurable"] @pytest.mark.asyncio async def test_wrapper_passes_trusted_session_config_to_graph(self): @@ -884,7 +887,8 @@ async def test_wrapper_passes_trusted_session_config_to_graph(self): result = await wrapper._ainvoke(value) assert result.value == "done" - assert graph.configs == [TRUSTED_CONFIG] + assert len(graph.configs) == 1 + assert graph.configs[0]["configurable"] == TRUSTED_CONFIG["configurable"] @pytest.mark.parametrize( "message", diff --git a/services/studio/src/nmp/studio/coding_agents.py b/services/studio/src/nmp/studio/coding_agents.py index a26afed727..a96c4a8b77 100644 --- a/services/studio/src/nmp/studio/coding_agents.py +++ b/services/studio/src/nmp/studio/coding_agents.py @@ -3,6 +3,7 @@ """Local coding-agent bridge for Studio.""" +import ast import asyncio import json import logging @@ -1298,11 +1299,66 @@ def _nemo_agent_request_payload( ) -> dict[str, Any]: return { "messages": messages, - "stream": False, + # Stream so the agent runs its streaming path and emits NAT + # ``intermediate_data:`` tool steps we relay to the chat as tool-use parts. + "stream": True, "studio_session_id": studio_session_id, } +_TOOL_STEP_PREFIX = "Tool: " + + +def _parse_tool_step_input(payload: Any) -> dict[str, Any]: + """Best-effort extract the tool input dict from a NAT step markdown payload. + + Payloads look like ``**Input:**\\n```json\\n{'resource': 'secrets'}...``; the + dict is a Python repr (single quotes), so parse the first balanced ``{...}`` + with ``ast.literal_eval`` and fall back to an empty dict. + """ + if not isinstance(payload, str): + return {} + start = payload.find("{") + if start == -1: + return {} + depth = 0 + for index in range(start, len(payload)): + if payload[index] == "{": + depth += 1 + elif payload[index] == "}": + depth -= 1 + if depth == 0: + try: + value = ast.literal_eval(payload[start : index + 1]) + except (ValueError, SyntaxError): + return {} + return value if isinstance(value, dict) else {} + return {} + + +def _tool_use_stream_event(tool_name: str, tool_input: dict[str, Any]) -> tuple[str, str]: + return ( + "agent", + json.dumps( + { + "type": "assistant", + "message": { + "id": f"nemo-agent-tool-{uuid.uuid4()}", + "model": _studio_coding_agent_name(), + "content": [ + { + "type": "tool_use", + "id": f"tool-{uuid.uuid4()}", + "name": tool_name, + "input": tool_input, + } + ], + }, + } + ), + ) + + async def _invoke_nemo_agent( agent_url: str, headers: Mapping[str, str], @@ -1315,17 +1371,59 @@ async def _invoke_nemo_agent( write=60.0, pool=10.0, ) + queue = _session_streams.get(studio_session_id) + content_parts: list[str] = [] + model = _studio_coding_agent_name() + seen_tool_ids: set[str] = set() async with httpx.AsyncClient(timeout=timeout) as client: # The origin and agent name are server-configured. The only request-derived URL # component is a NAME_PATTERN-validated, percent-encoded workspace path segment. # codeql[py/partial-ssrf] - response = await client.post( + async with client.stream( + "POST", agent_url, headers=dict(headers), json=_nemo_agent_request_payload(messages, studio_session_id), - ) - response.raise_for_status() - return _nemo_agent_response(response.json()) + ) as response: + response.raise_for_status() + async for line in response.aiter_lines(): + if line.startswith("data:"): + data = line[len("data:") :].strip() + if not data or data == "[DONE]": + continue + try: + chunk = json.loads(data) + except json.JSONDecodeError: + continue + if not isinstance(chunk, dict): + continue + if chunk.get("model"): + model = chunk["model"] + for choice in chunk.get("choices", []): + piece = (choice.get("delta") or {}).get("content") + if isinstance(piece, str) and piece: + content_parts.append(piece) + elif line.startswith("intermediate_data:") and queue is not None: + data = line[len("intermediate_data:") :].strip() + try: + step = json.loads(data) + except json.JSONDecodeError: + continue + name = step.get("name") if isinstance(step, dict) else None + if not isinstance(name, str) or not name.startswith(_TOOL_STEP_PREFIX): + continue + step_id = step.get("id") + if step_id in seen_tool_ids: + continue + if isinstance(step_id, str): + seen_tool_ids.add(step_id) + await queue.put( + _tool_use_stream_event( + name[len(_TOOL_STEP_PREFIX) :].strip(), + _parse_tool_step_input(step.get("payload")), + ) + ) + return "".join(content_parts), model async def _stream_nemo_agent( diff --git a/services/studio/tests/unit/test_coding_agents.py b/services/studio/tests/unit/test_coding_agents.py index df05d53187..3ee68f564d 100644 --- a/services/studio/tests/unit/test_coding_agents.py +++ b/services/studio/tests/unit/test_coding_agents.py @@ -1863,6 +1863,26 @@ def test_nemo_agent_error_detail_does_not_expose_exception_text(): ) +def test_parse_tool_step_input_extracts_python_repr_dict(): + payload = "**Input:**\n```json\n{'action': 'list', 'resource': 'secrets'}\n```\n**Output:** ..." + assert coding_agents._parse_tool_step_input(payload) == {"action": "list", "resource": "secrets"} + + +def test_parse_tool_step_input_returns_empty_on_unparseable(): + assert coding_agents._parse_tool_step_input("no dict here") == {} + assert coding_agents._parse_tool_step_input(None) == {} + assert coding_agents._parse_tool_step_input("**Input:** [1, 2, 3]") == {} + + +def test_tool_use_stream_event_shape(): + event_type, payload = coding_agents._tool_use_stream_event("nemo_api", {"resource": "secrets"}) + assert event_type == "agent" + block = json.loads(payload)["message"]["content"][0] + assert block["type"] == "tool_use" + assert block["name"] == "nemo_api" + assert block["input"] == {"resource": "secrets"} + + def test_nemo_agent_request_payload_keeps_session_outside_model_messages(): messages = [{"role": "user", "content": "hello"}] session_id = str(uuid.uuid4()) @@ -1871,7 +1891,7 @@ def test_nemo_agent_request_payload_keeps_session_outside_model_messages(): assert payload == { "messages": messages, - "stream": False, + "stream": True, "studio_session_id": session_id, } assert session_id not in json.dumps(payload["messages"]) From 8527b4be442a1e2e4f229ad0ce285b3f68abac54 Mon Sep 17 00:00:00 2001 From: Henrique Tolentino Date: Tue, 4 Aug 2026 13:44:10 -0400 Subject: [PATCH 2/4] Harden tool-usage streaming per review - Strip framework-injected studio_session_id from tool-use events before they reach the browser. - Gate tool-step dedup on a validated string id. - Drain queued tool-use events after the invocation completes so late events are not dropped before the final assistant message. - Resolve the CodeQL partial-SSRF finding: validate the requested workspace against the Entity Store (scoped to the caller's forwarded auth) and build the agent URL from the platform's own copy of the name, so no client-derived value reaches the outbound request URL. - Cache confirmed workspace names per session so the membership lookup runs once per session/workspace instead of on every message; only successful resolutions are cached, and the cache is cleared on session eviction. Signed-off-by: Henrique Tolentino --- .../studio/src/nmp/studio/coding_agents.py | 143 +++++++++++++++--- .../studio/tests/unit/test_coding_agents.py | 127 ++++++++++++++++ 2 files changed, 248 insertions(+), 22 deletions(-) diff --git a/services/studio/src/nmp/studio/coding_agents.py b/services/studio/src/nmp/studio/coding_agents.py index a96c4a8b77..5ed76ec05f 100644 --- a/services/studio/src/nmp/studio/coding_agents.py +++ b/services/studio/src/nmp/studio/coding_agents.py @@ -141,6 +141,10 @@ class SessionHistoryResponse(BaseModel): _pending_agent_inputs: dict[str, tuple[str, asyncio.Future[dict[str, Any]]]] = {} _session_conversations: dict[str, list[dict[str, str]]] = {} _session_mtimes: dict[str, float] = {} +# Cache of Entity-Store-confirmed workspace names, keyed by session then by the +# requested workspace, so the membership lookup runs once per session/workspace +# rather than on every message. Cleared when a session is evicted. +_session_workspace_cache: dict[str, dict[str, str]] = {} _AGENT_INPUT_RESPONSE_RESERVED_KEYS = frozenset({"message", "status"}) @@ -158,9 +162,12 @@ def _evict_oldest_sessions(*, protected_session_ids: set[str] | None = None) -> _session_conversations.pop(oldest_session_id, None) _session_mtimes.pop(oldest_session_id, None) _initialized_sessions.discard(oldest_session_id) + _session_workspace_cache.pop(oldest_session_id, None) for session_id in set(_session_mtimes) - set(_session_conversations): _session_mtimes.pop(session_id, None) + for session_id in set(_session_workspace_cache) - set(_session_conversations): + _session_workspace_cache.pop(session_id, None) def _retain_recent_turns(conversation: list[dict[str, str]]) -> None: @@ -1246,6 +1253,66 @@ def _validated_workspace_or_default(value: str | None) -> str: return workspace +# Bounds for the workspace-membership lookup below. The entities list is scoped to +# the caller's own memberships (forwarded auth), so the page count is normally tiny. +_WORKSPACE_LOOKUP_PAGE_SIZE = 100 +_WORKSPACE_LOOKUP_MAX_PAGES = 50 + + +async def _authorized_workspace(workspace: str, headers: Mapping[str, str], session_id: str) -> str: + """Confirm the caller may target ``workspace`` and return the platform's own + spelling of the name. + + The value is looked up in the Entity Store (scoped to the caller's auth) and the + returned name is taken from the *response*, not from client input, so the value + that later becomes part of the agent request URL cannot be attacker-controlled. + The list request itself carries no user-provided value in its URL. + + The confirmed name is cached per (session, requested workspace) so the lookup + runs once per session/workspace instead of on every message; only successful + resolutions are cached. + """ + # ``default`` is our own fallback constant, always addressable; skip the lookup. + if workspace == "default": + return "default" + + cached = _session_workspace_cache.get(session_id, {}).get(workspace) + if cached is not None: + return cached + + base_url = _studio_coding_agent_base_url() + list_url = f"{base_url}/apis/entities/v2/workspaces" + timeout = httpx.Timeout(connect=10.0, read=30.0, write=30.0, pool=10.0) + try: + async with httpx.AsyncClient(timeout=timeout) as client: + for page in range(1, _WORKSPACE_LOOKUP_MAX_PAGES + 1): + response = await client.get( + list_url, + headers=dict(headers), + params={"page": page, "page_size": _WORKSPACE_LOOKUP_PAGE_SIZE}, + ) + response.raise_for_status() + body = response.json() + data = body.get("data") if isinstance(body, dict) else None + if not isinstance(data, list) or not data: + break + for entry in data: + name = entry.get("name") if isinstance(entry, dict) else None + if isinstance(name, str) and name == workspace: + # ``name`` is sourced from the platform response, not the request. + _session_workspace_cache.setdefault(session_id, {})[workspace] = name + return name + pagination = body.get("pagination") if isinstance(body, dict) else None + total_pages = pagination.get("total_pages") if isinstance(pagination, dict) else None + if isinstance(total_pages, int) and page >= total_pages: + break + except httpx.HTTPError as exc: + logger.warning("Workspace authorization lookup failed: %s", type(exc).__name__) + raise HTTPException(status_code=502, detail="Could not verify the requested workspace") from exc + + raise HTTPException(status_code=404, detail="workspace not found or not accessible") + + def _workspace_path_segment(workspace: str) -> str: if WORKSPACE_NAME_RE.fullmatch(workspace) is None: raise ValueError("Invalid workspace name") @@ -1253,8 +1320,8 @@ def _workspace_path_segment(workspace: str) -> str: def _studio_coding_agent_url(workspace: str) -> str: - # MessageRequest validates workspace against the platform's restricted entity-name - # pattern before it can become part of this internal request path. + # ``workspace`` is expected to already be an Entity-Store-confirmed name (see + # _authorized_workspace); it is percent-encoded as a single path segment here. return ( f"{_studio_coding_agent_base_url()}/apis/agents/v2/workspaces/{_workspace_path_segment(workspace)}" f"/agents/{quote(_studio_coding_agent_name(), safe='')}/-/v1/chat/completions" @@ -1336,7 +1403,13 @@ def _parse_tool_step_input(payload: Any) -> dict[str, Any]: return {} +# Framework-injected tool arguments that must never be surfaced in the browser +# tool-use event (they are internal plumbing, not user-facing input). +_TOOL_INPUT_INTERNAL_KEYS = frozenset({"studio_session_id"}) + + def _tool_use_stream_event(tool_name: str, tool_input: dict[str, Any]) -> tuple[str, str]: + safe_input = {key: value for key, value in tool_input.items() if key not in _TOOL_INPUT_INTERNAL_KEYS} return ( "agent", json.dumps( @@ -1350,7 +1423,7 @@ def _tool_use_stream_event(tool_name: str, tool_input: dict[str, Any]) -> tuple[ "type": "tool_use", "id": f"tool-{uuid.uuid4()}", "name": tool_name, - "input": tool_input, + "input": safe_input, } ], }, @@ -1376,9 +1449,8 @@ async def _invoke_nemo_agent( model = _studio_coding_agent_name() seen_tool_ids: set[str] = set() async with httpx.AsyncClient(timeout=timeout) as client: - # The origin and agent name are server-configured. The only request-derived URL - # component is a NAME_PATTERN-validated, percent-encoded workspace path segment. - # codeql[py/partial-ssrf] + # The origin and agent name are server-configured; the workspace path segment is + # an Entity-Store-confirmed name resolved by _authorized_workspace before this call. async with client.stream( "POST", agent_url, @@ -1413,9 +1485,9 @@ async def _invoke_nemo_agent( if not isinstance(name, str) or not name.startswith(_TOOL_STEP_PREFIX): continue step_id = step.get("id") - if step_id in seen_tool_ids: - continue if isinstance(step_id, str): + if step_id in seen_tool_ids: + continue seen_tool_ids.add(step_id) await queue.put( _tool_use_stream_event( @@ -1426,6 +1498,21 @@ async def _invoke_nemo_agent( return "".join(content_parts), model +def _render_session_event(event_type: str, payload: Any) -> str | None: + """Render a queued session event as an SSE frame, or None if it is not relayable.""" + if event_type == "permission_request": + return _sse(payload, event="permission_request") + if event_type == "input_request": + return _sse(payload, event="input_request") + if event_type == "permission_expired": + return _sse(payload, event="permission_expired") + if event_type == "input_expired": + return _sse(payload, event="input_expired") + if event_type == "agent": + return _sse(payload) + return None + + async def _stream_nemo_agent( session_id: str, message: str, @@ -1475,19 +1562,29 @@ async def _stream_nemo_agent( yield ":\n\n" continue if queued_event in done: - event_type, payload = queued_event.result() - if event_type == "permission_request": - yield _sse(payload, event="permission_request") - elif event_type == "input_request": - yield _sse(payload, event="input_request") - elif event_type == "permission_expired": - yield _sse(payload, event="permission_expired") - elif event_type == "input_expired": - yield _sse(payload, event="input_expired") - elif event_type == "agent": - yield _sse(payload) + rendered = _render_session_event(*queued_event.result()) + if rendered is not None: + yield rendered queued_event = asyncio.create_task(queue.get()) + # The invocation finished, but tool-use / prompt events may have been + # queued in the same event-loop turn as completion. Flush them before the + # final assistant message so they are not dropped when the loop exits. + pending: list[tuple[str, Any]] = [] + if queued_event.done() and not queued_event.cancelled(): + pending.append(queued_event.result()) + else: + queued_event.cancel() + while True: + try: + pending.append(queue.get_nowait()) + except asyncio.QueueEmpty: + break + for event_type, payload in pending: + rendered = _render_session_event(event_type, payload) + if rendered is not None: + yield rendered + assistant_text, model = await invocation conversation.extend( [ @@ -1534,12 +1631,14 @@ async def send_message(session_id: str, body: MessageRequest, request: Request) """Send a message to the deployed NeMo Agent and stream Studio events.""" sid = _validate_session_id(session_id) workspace = _validated_workspace_or_default(body.workspace) + agent_headers = _coding_agent_request_headers(request) + canonical_workspace = await _authorized_workspace(workspace, agent_headers, sid) studio_base_url = _studio_base_url_from_request(body, request) studio_pathname = _studio_pathname_from_request(body, request) enabled_destinations = studio_links.enabled_destinations_from_request(request) system_prompt = _build_nemo_agent_system_prompt( sid, - workspace, + canonical_workspace, studio_base_url, studio_pathname, enabled_destinations, @@ -1548,8 +1647,8 @@ async def send_message(session_id: str, body: MessageRequest, request: Request) _stream_nemo_agent( sid, body.message, - _studio_coding_agent_url(workspace), - _coding_agent_request_headers(request), + _studio_coding_agent_url(canonical_workspace), + agent_headers, system_prompt, ), media_type="text/event-stream", diff --git a/services/studio/tests/unit/test_coding_agents.py b/services/studio/tests/unit/test_coding_agents.py index 3ee68f564d..d36402c2df 100644 --- a/services/studio/tests/unit/test_coding_agents.py +++ b/services/studio/tests/unit/test_coding_agents.py @@ -29,6 +29,7 @@ def reset_coding_agent_state(): coding_agents._pending_agent_inputs.clear() coding_agents._session_conversations.clear() coding_agents._session_mtimes.clear() + coding_agents._session_workspace_cache.clear() yield coding_agents._initialized_sessions.clear() coding_agents._session_streams.clear() @@ -36,6 +37,7 @@ def reset_coding_agent_state(): coding_agents._pending_agent_inputs.clear() coding_agents._session_conversations.clear() coding_agents._session_mtimes.clear() + coding_agents._session_workspace_cache.clear() @pytest.fixture @@ -1852,6 +1854,90 @@ def test_platform_route_rejects_workspace_path_injection(service_client: TestCli assert response.status_code == 422 +_WORKSPACES_LIST_URL = "http://127.0.0.1:8080/apis/entities/v2/workspaces" + + +@pytest.mark.asyncio +async def test_authorized_workspace_default_skips_lookup(): + # No HTTP mock configured: the default fallback must not make a network call. + assert await coding_agents._authorized_workspace("default", {}, "sess") == "default" + + +@pytest.mark.asyncio +async def test_authorized_workspace_returns_name_from_entity_store(respx_mock): + respx_mock.get(_WORKSPACES_LIST_URL).mock( + return_value=httpx.Response( + 200, + json={"data": [{"name": "team-a"}, {"name": "my-ws"}], "pagination": {"total_pages": 1}}, + ) + ) + + # The returned value is the platform's own copy of the name (not client input). + assert await coding_agents._authorized_workspace("my-ws", {}, "sess") == "my-ws" + + +@pytest.mark.asyncio +async def test_authorized_workspace_paginates_until_match(respx_mock): + respx_mock.get(_WORKSPACES_LIST_URL).mock( + side_effect=[ + httpx.Response(200, json={"data": [{"name": "a"}], "pagination": {"total_pages": 2}}), + httpx.Response(200, json={"data": [{"name": "target"}], "pagination": {"total_pages": 2}}), + ] + ) + + assert await coding_agents._authorized_workspace("target", {}, "sess") == "target" + + +@pytest.mark.asyncio +async def test_authorized_workspace_caches_per_session(respx_mock): + route = respx_mock.get(_WORKSPACES_LIST_URL).mock( + return_value=httpx.Response(200, json={"data": [{"name": "my-ws"}], "pagination": {"total_pages": 1}}) + ) + + first = await coding_agents._authorized_workspace("my-ws", {}, "sess") + second = await coding_agents._authorized_workspace("my-ws", {}, "sess") + + assert first == second == "my-ws" + # The second resolution is served from the per-session cache, not the network. + assert route.call_count == 1 + # A different session does not share the cache. + await coding_agents._authorized_workspace("my-ws", {}, "other-sess") + assert route.call_count == 2 + + +@pytest.mark.asyncio +async def test_authorized_workspace_rejects_unknown_workspace(respx_mock): + respx_mock.get(_WORKSPACES_LIST_URL).mock( + return_value=httpx.Response(200, json={"data": [{"name": "team-a"}], "pagination": {"total_pages": 1}}) + ) + + with pytest.raises(HTTPException) as excinfo: + await coding_agents._authorized_workspace("not-a-member", {}, "sess") + assert excinfo.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_authorized_workspace_does_not_cache_failures(respx_mock): + route = respx_mock.get(_WORKSPACES_LIST_URL).mock( + return_value=httpx.Response(200, json={"data": [{"name": "team-a"}], "pagination": {"total_pages": 1}}) + ) + + for _ in range(2): + with pytest.raises(HTTPException): + await coding_agents._authorized_workspace("not-a-member", {}, "sess") + # Unresolved workspaces are re-checked every message (no negative caching). + assert route.call_count == 2 + + +@pytest.mark.asyncio +async def test_authorized_workspace_maps_upstream_error_to_502(respx_mock): + respx_mock.get(_WORKSPACES_LIST_URL).mock(side_effect=httpx.ConnectError("boom")) + + with pytest.raises(HTTPException) as excinfo: + await coding_agents._authorized_workspace("my-ws", {}, "sess") + assert excinfo.value.status_code == 502 + + def test_nemo_agent_error_detail_does_not_expose_exception_text(): request = httpx.Request("POST", "https://platform.test/agent") response = httpx.Response(502, request=request) @@ -1883,6 +1969,47 @@ def test_tool_use_stream_event_shape(): assert block["input"] == {"resource": "secrets"} +def test_tool_use_stream_event_strips_internal_session_id(): + _, payload = coding_agents._tool_use_stream_event( + "ask_user_question", + {"studio_session_id": "sess-123", "questions": [{"q": "?"}]}, + ) + block = json.loads(payload)["message"]["content"][0] + assert "studio_session_id" not in block["input"] + assert block["input"] == {"questions": [{"q": "?"}]} + + +@pytest.mark.asyncio +async def test_stream_nemo_agent_flushes_tool_events_before_final_response(monkeypatch: pytest.MonkeyPatch): + session_id = str(uuid.uuid4()) + + async def fake_invoke(agent_url, headers, messages, studio_session_id): + queue = coding_agents._session_streams[studio_session_id] + # Two tool events queued in the same turn the invocation completes: the + # loop can consume at most one, so the drain must flush the remainder. + queue.put_nowait(coding_agents._tool_use_stream_event("nemo_api", {"resource": "secrets"})) + queue.put_nowait(coding_agents._tool_use_stream_event("describe_api", {"path": "secrets"})) + return "final answer", "model-x" + + monkeypatch.setattr(coding_agents, "_invoke_nemo_agent", fake_invoke) + + frames = [ + frame + async for frame in coding_agents._stream_nemo_agent( + session_id, "hello", "https://agent.test/x", {}, "sys prompt" + ) + ] + + body = "".join(frames) + first_tool = body.find("nemo_api") + second_tool = body.find("describe_api") + final = body.find("final answer") + assert first_tool != -1 and second_tool != -1 and final != -1 + # Both tool-use events survive and are emitted before the final assistant message. + assert first_tool < final + assert second_tool < final + + def test_nemo_agent_request_payload_keeps_session_outside_model_messages(): messages = [{"role": "user", "content": "hello"}] session_id = str(uuid.uuid4()) From ecdc75eef16cabc233aa5d78b6557e12e448d920 Mon Sep 17 00:00:00 2001 From: Henrique Tolentino Date: Tue, 4 Aug 2026 13:44:18 -0400 Subject: [PATCH 3/4] Scope cached workspace authorization to the calling credential Session ids carry no caller identity: create_session mints a bare UUID and nothing binds a session to a user. Keying the workspace-membership cache on (session, workspace) alone therefore let a second caller reuse the first caller's authorization decision and skip the Entity Store check. Include a SHA-256 fingerprint of the caller's forwarded credentials in the cache key so a cached decision is never reused across callers. Only the digest is retained, never the raw credential. Signed-off-by: Henrique Tolentino --- .../studio/src/nmp/studio/coding_agents.py | 33 +++++++++++---- .../studio/tests/unit/test_coding_agents.py | 41 +++++++++++++++++++ 2 files changed, 65 insertions(+), 9 deletions(-) diff --git a/services/studio/src/nmp/studio/coding_agents.py b/services/studio/src/nmp/studio/coding_agents.py index 5ed76ec05f..c3428f9d79 100644 --- a/services/studio/src/nmp/studio/coding_agents.py +++ b/services/studio/src/nmp/studio/coding_agents.py @@ -5,6 +5,7 @@ import ast import asyncio +import hashlib import json import logging import os @@ -141,10 +142,12 @@ class SessionHistoryResponse(BaseModel): _pending_agent_inputs: dict[str, tuple[str, asyncio.Future[dict[str, Any]]]] = {} _session_conversations: dict[str, list[dict[str, str]]] = {} _session_mtimes: dict[str, float] = {} -# Cache of Entity-Store-confirmed workspace names, keyed by session then by the -# requested workspace, so the membership lookup runs once per session/workspace -# rather than on every message. Cleared when a session is evicted. -_session_workspace_cache: dict[str, dict[str, str]] = {} +# Cache of Entity-Store-confirmed workspace names, keyed by session then by +# (caller fingerprint, requested workspace), so the membership lookup runs once per +# session/caller/workspace rather than on every message. Session ids are not bound to +# a caller, so the caller's credential participates in the key: a cached authorization +# decision must never be reused for a different caller. Cleared on session eviction. +_session_workspace_cache: dict[str, dict[tuple[str, str], str]] = {} _AGENT_INPUT_RESPONSE_RESERVED_KEYS = frozenset({"message", "status"}) @@ -1259,6 +1262,16 @@ def _validated_workspace_or_default(value: str | None) -> str: _WORKSPACE_LOOKUP_MAX_PAGES = 50 +def _caller_fingerprint(headers: Mapping[str, str]) -> str: + """Return a stable, non-reversible fingerprint of the caller's forwarded credentials. + + Used only to scope cached authorization decisions to the caller they were made + for. The raw credential is never retained -- just the digest. + """ + material = json.dumps({name.lower(): value for name, value in sorted(headers.items())}, sort_keys=True) + return hashlib.sha256(material.encode("utf-8")).hexdigest() + + async def _authorized_workspace(workspace: str, headers: Mapping[str, str], session_id: str) -> str: """Confirm the caller may target ``workspace`` and return the platform's own spelling of the name. @@ -1268,15 +1281,17 @@ async def _authorized_workspace(workspace: str, headers: Mapping[str, str], sess that later becomes part of the agent request URL cannot be attacker-controlled. The list request itself carries no user-provided value in its URL. - The confirmed name is cached per (session, requested workspace) so the lookup - runs once per session/workspace instead of on every message; only successful - resolutions are cached. + The confirmed name is cached per (session, caller, requested workspace) so the + lookup runs once per session/caller/workspace instead of on every message; only + successful resolutions are cached. The caller fingerprint is part of the key so + one caller's authorization decision is never reused for another. """ # ``default`` is our own fallback constant, always addressable; skip the lookup. if workspace == "default": return "default" - cached = _session_workspace_cache.get(session_id, {}).get(workspace) + cache_key = (_caller_fingerprint(headers), workspace) + cached = _session_workspace_cache.get(session_id, {}).get(cache_key) if cached is not None: return cached @@ -1300,7 +1315,7 @@ async def _authorized_workspace(workspace: str, headers: Mapping[str, str], sess name = entry.get("name") if isinstance(entry, dict) else None if isinstance(name, str) and name == workspace: # ``name`` is sourced from the platform response, not the request. - _session_workspace_cache.setdefault(session_id, {})[workspace] = name + _session_workspace_cache.setdefault(session_id, {})[cache_key] = name return name pagination = body.get("pagination") if isinstance(body, dict) else None total_pages = pagination.get("total_pages") if isinstance(pagination, dict) else None diff --git a/services/studio/tests/unit/test_coding_agents.py b/services/studio/tests/unit/test_coding_agents.py index d36402c2df..f76e8cf6da 100644 --- a/services/studio/tests/unit/test_coding_agents.py +++ b/services/studio/tests/unit/test_coding_agents.py @@ -1905,6 +1905,47 @@ async def test_authorized_workspace_caches_per_session(respx_mock): assert route.call_count == 2 +@pytest.mark.asyncio +async def test_authorized_workspace_cache_is_not_shared_across_callers(respx_mock): + """A cached authorization decision must never be reused for a different caller.""" + route = respx_mock.get(_WORKSPACES_LIST_URL).mock( + return_value=httpx.Response(200, json={"data": [{"name": "my-ws"}], "pagination": {"total_pages": 1}}) + ) + session_id = "shared-session" + caller_a = {"authorization": "Bearer token-a"} + caller_b = {"authorization": "Bearer token-b"} + + await coding_agents._authorized_workspace("my-ws", caller_a, session_id) + assert route.call_count == 1 + # Same session id, different credential: must re-verify against the Entity Store. + await coding_agents._authorized_workspace("my-ws", caller_b, session_id) + assert route.call_count == 2 + # Each caller still gets its own cache hit on repeat. + await coding_agents._authorized_workspace("my-ws", caller_a, session_id) + assert route.call_count == 2 + # The raw credential is never retained in the cache key material. + assert "token-a" not in str(coding_agents._session_workspace_cache) + + +@pytest.mark.asyncio +async def test_authorized_workspace_unauthorized_caller_is_rejected_on_cached_session(respx_mock): + """An unauthorized caller cannot ride a session that already resolved the workspace.""" + session_id = "shared-session" + respx_mock.get(_WORKSPACES_LIST_URL).mock( + side_effect=[ + httpx.Response(200, json={"data": [{"name": "my-ws"}], "pagination": {"total_pages": 1}}), + # The second caller is not a member of that workspace. + httpx.Response(200, json={"data": [{"name": "other-ws"}], "pagination": {"total_pages": 1}}), + ] + ) + + assert await coding_agents._authorized_workspace("my-ws", {"authorization": "a"}, session_id) == "my-ws" + + with pytest.raises(HTTPException) as excinfo: + await coding_agents._authorized_workspace("my-ws", {"authorization": "b"}, session_id) + assert excinfo.value.status_code == 404 + + @pytest.mark.asyncio async def test_authorized_workspace_rejects_unknown_workspace(respx_mock): respx_mock.get(_WORKSPACES_LIST_URL).mock( From af11921901780d471246f00726dee48db1064c2a Mon Sep 17 00:00:00 2001 From: Henrique Tolentino Date: Wed, 5 Aug 2026 10:13:55 -0400 Subject: [PATCH 4/4] Format test_copilot.py after the copilot rename The merge renames (coding_agents -> copilot, _stream_nemo_agent -> _stream_copilot) shortened a comprehension enough to fit on one line, which ruff format flags. Cosmetic only; lint-python-style was the sole failing CI lint. Signed-off-by: Henrique Tolentino --- services/studio/tests/unit/test_copilot.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/services/studio/tests/unit/test_copilot.py b/services/studio/tests/unit/test_copilot.py index 2b9146664c..247ba7642e 100644 --- a/services/studio/tests/unit/test_copilot.py +++ b/services/studio/tests/unit/test_copilot.py @@ -2070,10 +2070,7 @@ async def fake_invoke(agent_url, headers, messages, studio_session_id): monkeypatch.setattr(copilot, "_invoke_copilot", fake_invoke) frames = [ - frame - async for frame in copilot._stream_copilot( - session_id, "hello", "https://agent.test/x", {}, "sys prompt" - ) + frame async for frame in copilot._stream_copilot(session_id, "hello", "https://agent.test/x", {}, "sys prompt") ] body = "".join(frames)