feat(studio): Studio agent add tool usage info

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:

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

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 @@
 ) -> 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 @@
         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(

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