feat(studio): Studio copilot reasoning trace

agents/nemo-studio-copilot/src/nemo_studio_copilot/register.py

@@ -769,13 +769,67 @@ async def astream_events(self, *args, **kwargs):
             yield event
 
 
+def _preserve_reasoning_content(model: Any) -> Any:
+    """Keep the ``reasoning_content`` delta that langchain-openai discards.
+
+    Nemotron (and other vLLM/NIM-served models) stream their chain of thought in a
+    ``reasoning_content`` delta field. ``langchain-openai`` 1.4.x only understands
+    OpenAI's o-series ``reasoning`` block, so it drops the field entirely and the
+    trace never reaches the graph, the NAT stream, or Studio. Re-attach it to the
+    chunk's ``additional_kwargs`` so downstream consumers can surface it.
+    """
+    original = getattr(model, "_convert_chunk_to_generation_chunk", None)
+    if original is None:
+        return model
+
+    def convert_with_reasoning(chunk: Any, *args: Any, **kwargs: Any) -> Any:
+        generation = original(chunk, *args, **kwargs)
+        if generation is None or not isinstance(chunk, dict):
+            return generation
+        reasoning = "".join(
+            (choice.get("delta") or {}).get("reasoning_content") or ""
+            for choice in chunk.get("choices") or []
+            if isinstance(choice, dict)
+        )
+        if reasoning:
+            generation.message.additional_kwargs["reasoning_content"] = reasoning
+        return generation
+
+    original_result = getattr(model, "_create_chat_result", None)
+
+    def make_result_patch(original: Any) -> Any:
+        def create_result_with_reasoning(response: Any, *args: Any, **kwargs: Any) -> Any:
+            result = original(response, *args, **kwargs)
+            # The non-streaming path drops the field too, so patch both: the graph
+            # may invoke the model rather than stream it.
+            try:
+                payload = response if isinstance(response, dict) else response.model_dump()
+                for generation, choice in zip(result.generations, payload.get("choices") or [], strict=False):
+                    reasoning = (choice.get("message") or {}).get("reasoning_content")
+                    if reasoning:
+                        generation.message.additional_kwargs["reasoning_content"] = reasoning
+            except Exception:
+                logger.debug("could not attach reasoning_content to result", exc_info=True)
+            return result
+
+        return create_result_with_reasoning
+
+    try:
+        object.__setattr__(model, "_convert_chunk_to_generation_chunk", convert_with_reasoning)
+        if original_result is not None:
+            object.__setattr__(model, "_create_chat_result", make_result_patch(original_result))
+    except Exception:
+        logger.debug("could not preserve reasoning_content on %s", type(model).__name__, exc_info=True)
+    return model
+
+
 def _get_model():
     """Get the LLM from NAT's builder context (YAML llms: section)."""
     from nat.builder.framework_enum import LLMFrameworkEnum
     from nat.builder.sync_builder import SyncBuilder
 
     model = SyncBuilder.current().get_llm("agent", wrapper_type=LLMFrameworkEnum.LANGCHAIN)
-    return _disable_nat_method_retries(model)
+    return _preserve_reasoning_content(_disable_nat_method_retries(model))
 
 
 def _disable_nat_method_retries(model: Any) -> Any:

agents/nemo-studio-copilot/src/nemo_studio_copilot/wrapper.py

@@ -63,18 +63,25 @@
 import logging
 from collections.abc import AsyncGenerator
 from typing import Any
-from uuid import UUID
+from uuid import UUID, uuid4
 
+from langchain_core.callbacks import BaseCallbackHandler
 from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, MessageLikeRepresentation
 from langchain_core.messages.utils import convert_to_messages
 from langchain_core.prompt_values import PromptValue
 from langchain_core.runnables import RunnableConfig
 from nat.builder.builder import Builder
+from nat.builder.context import Context
 from nat.builder.framework_enum import LLMFrameworkEnum
 from nat.builder.function import Function
 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.data_models.intermediate_step import (
+    IntermediateStepPayload,
+    IntermediateStepType,
+    StreamEventData,
+)
 from nat.plugins.langchain.callback_handler import LangchainProfilerHandler
 from nemo_studio_copilot.register import create_nemo_studio_copilot
 from pydantic import BaseModel, ConfigDict, Field, computed_field, model_validator
@@ -93,6 +100,82 @@
 _MODEL_ERROR_MESSAGE = "The model service returned an error before I could complete that request. Please try again."
 
 
+REASONING_STEP_PREFIX = "Reasoning: "
+
+
+class ReasoningStreamHandler(BaseCallbackHandler):
+    """Publish the model's chain of thought as NAT intermediate steps.
+
+    ``reasoning_content`` is re-attached to each chunk by
+    ``register._preserve_reasoning_content`` (langchain-openai drops it). NAT's
+    ``ChatResponseChunk`` only carries ``content``, so the trace rides the same
+    ``intermediate_data:`` channel the tool-call trace already uses and Studio
+    already parses.
+
+    Tokens are buffered and flushed per LLM run rather than per token: one step
+    per token would be thousands of SSE frames for a single answer.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._step_manager = Context.get().intermediate_step_manager
+        self._buffers: dict[str, list[str]] = {}
+
+    def on_llm_new_token(self, token: str | list[str | dict[str, Any]], **kwargs: Any) -> None:
+        chunk = kwargs.get("chunk")
+        reasoning = getattr(getattr(chunk, "message", None), "additional_kwargs", {}).get("reasoning_content")
+        if reasoning:
+            self._buffers.setdefault(str(kwargs.get("run_id", "")), []).append(reasoning)
+
+    def on_llm_end(self, response: Any, **kwargs: Any) -> None:
+        run_id = str(kwargs.get("run_id", ""))
+        if run_id not in self._buffers:
+            # The graph may invoke the model instead of streaming it, so no tokens
+            # arrived; take the reasoning off the finished message.
+            for generations in getattr(response, "generations", []) or []:
+                for generation in generations or []:
+                    reasoning = getattr(getattr(generation, "message", None), "additional_kwargs", {}).get(
+                        "reasoning_content"
+                    )
+                    if reasoning:
+                        self._buffers.setdefault(run_id, []).append(reasoning)
+        self._flush(run_id)
+
+    def on_llm_error(self, error: BaseException, **kwargs: Any) -> None:
+        self._flush(str(kwargs.get("run_id", "")))
+
+    def _flush(self, run_id: str) -> None:
+        reasoning = "".join(self._buffers.pop(run_id, []))
+        if not reasoning.strip():
+            return
+        # NAT's StepAdaptor (DEFAULT mode) only forwards LLM, TOOL and FUNCTION
+        # categories, so a CUSTOM step never reaches the stream. Reasoning is LLM
+        # output, so publish it as an LLM step pair -- an END without a matching
+        # START is dropped as "not found in outstanding start steps".
+        step_id = str(uuid4())
+        name = f"{REASONING_STEP_PREFIX}model"
+        try:
+            self._step_manager.push_intermediate_step(
+                IntermediateStepPayload(
+                    event_type=IntermediateStepType.LLM_START,
+                    name=name,
+                    UUID=step_id,
+                    # The adaptor reads start.data.input and drops the pair when it is empty.
+                    data=StreamEventData(input="chain of thought"),
+                )
+            )
+            self._step_manager.push_intermediate_step(
+                IntermediateStepPayload(
+                    event_type=IntermediateStepType.LLM_END,
+                    name=name,
+                    UUID=step_id,
+                    data=StreamEventData(output=reasoning),
+                )
+            )
+        except Exception:
+            logger.debug("could not publish reasoning trace", exc_info=True)
+
+
 class NemoStudioCopilotWrapperConfig(FunctionBaseConfig, name="nemo_studio_copilot_wrapper"):
     """Configuration for the nemo-studio-copilot NAT workflow.
 
@@ -249,10 +332,17 @@ def _invocation_config(value: NemoStudioCopilotWrapperInput) -> RunnableConfig:
         # 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.
+        callbacks: list[Any] = []
         try:
-            config["callbacks"] = [LangchainProfilerHandler()]
+            callbacks.append(LangchainProfilerHandler())
         except Exception:
             logger.debug("LangchainProfilerHandler unavailable; tool-call trace disabled", exc_info=True)
+        try:
+            callbacks.append(ReasoningStreamHandler())
+        except Exception:
+            logger.debug("reasoning trace unavailable outside a NAT context", exc_info=True)
+        if callbacks:
+            config["callbacks"] = callbacks
         if value.studio_session_id is not None:
             config["configurable"] = {"studio_session_id": str(value.studio_session_id)}
         return config

services/studio/src/nmp/studio/copilot.py

@@ -1522,6 +1522,72 @@ def _parse_tool_step_input(payload: Any) -> dict[str, Any]:
 _TOOL_INPUT_INTERNAL_KEYS = frozenset({"studio_session_id"})
 
 
+_REASONING_STEP_PREFIX = "Reasoning: "
+
+
+def _parse_reasoning_step_output(payload: Any) -> str:
+    """Pull the chain of thought out of a NAT step payload.
+
+    NAT renders steps as markdown with an ``**Input:**`` block and, once the step
+    completes, an ``**Output:**`` block holding the reasoning.
+    """
+    if not isinstance(payload, str):
+        return ""
+    marker = "**Output:**"
+    index = payload.find(marker)
+    if index == -1:
+        return ""
+    return payload[index + len(marker) :].strip()
+
+
+def _reasoning_stream_event(reasoning: str) -> tuple[str, str]:
+    return (
+        "agent",
+        json.dumps(
+            {
+                "type": "assistant",
+                "message": {
+                    "id": f"nemo-copilot-reasoning-{uuid.uuid4()}",
+                    "model": _studio_copilot_name(),
+                    "content": [{"type": "reasoning", "text": reasoning}],
+                },
+            }
+        ),
+    )
+
+
+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 {}
+
+
+# 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 (
@@ -1562,6 +1628,7 @@ async def _invoke_copilot(
     content_parts: list[str] = []
     model = _studio_copilot_name()
     seen_tool_ids: set[str] = set()
+    seen_reasoning_ids: set[str] = set()
     async with httpx.AsyncClient(timeout=timeout) as client:
         # 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.
@@ -1596,7 +1663,25 @@ async def _invoke_copilot(
                     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):
+                    if not isinstance(name, str):
+                        continue
+                    if name.startswith(_REASONING_STEP_PREFIX):
+                        # Published as a start/end pair sharing one id, so this cannot use
+                        # the tool dedup below: the start would claim the id and the end,
+                        # which carries the trace, would be dropped. Only the end has an
+                        # Output block, so record the id only once one is parsed -- that
+                        # skips a repeated end without ever marking the start as seen.
+                        reasoning = _parse_reasoning_step_output(step.get("payload"))
+                        if not reasoning:
+                            continue
+                        step_id = step.get("id")
+                        if isinstance(step_id, str):
+                            if step_id in seen_reasoning_ids:
+                                continue
+                            seen_reasoning_ids.add(step_id)
+                        await queue.put(_reasoning_stream_event(reasoning))
+                        continue
+                    if not name.startswith(_TOOL_STEP_PREFIX):
                         continue
                     step_id = step.get("id")
                     if isinstance(step_id, str):