diff --git a/agents/nemo-studio-copilot/src/nemo_studio_copilot/register.py b/agents/nemo-studio-copilot/src/nemo_studio_copilot/register.py index d54aa3fc92..e36755d379 100644 --- a/agents/nemo-studio-copilot/src/nemo_studio_copilot/register.py +++ b/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: diff --git a/agents/nemo-studio-copilot/src/nemo_studio_copilot/wrapper.py b/agents/nemo-studio-copilot/src/nemo_studio_copilot/wrapper.py index f8932a2e50..eb07200b8d 100644 --- a/agents/nemo-studio-copilot/src/nemo_studio_copilot/wrapper.py +++ b/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 diff --git a/services/studio/src/nmp/studio/copilot.py b/services/studio/src/nmp/studio/copilot.py index 86ec839d82..1bc13c32fa 100644 --- a/services/studio/src/nmp/studio/copilot.py +++ b/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): diff --git a/services/studio/tests/unit/test_copilot.py b/services/studio/tests/unit/test_copilot.py index 5ddf03ece6..8caadec377 100644 --- a/services/studio/tests/unit/test_copilot.py +++ b/services/studio/tests/unit/test_copilot.py @@ -2452,3 +2452,148 @@ def test_copilot_routes_are_available_by_default(): assert response.status_code == 200 uuid.UUID(response.json()["session_id"]) + + +def test_parse_reasoning_step_output_extracts_the_trace(): + payload = "**Input:**\n```python\nchain of thought\n```\n\n**Output:** The user asked a math question." + assert copilot._parse_reasoning_step_output(payload) == "The user asked a math question." + + +def test_parse_reasoning_step_output_ignores_a_start_step(): + # The paired start step has no Output block yet. + assert copilot._parse_reasoning_step_output("**Input:**\n```python\nchain of thought\n```") == "" + assert copilot._parse_reasoning_step_output(None) == "" + + +def test_reasoning_stream_event_shape(): + event_type, payload = copilot._reasoning_stream_event("thinking out loud") + assert event_type == "agent" + block = json.loads(payload)["message"]["content"][0] + assert block == {"type": "reasoning", "text": "thinking out loud"} + + +@pytest.mark.asyncio +async def test_invoke_copilot_relays_reasoning_despite_shared_step_id(monkeypatch: pytest.MonkeyPatch): + """A reasoning step is a start/end pair sharing one id. The tool dedup must not + claim that id on the start, or the end -- which carries the trace -- is dropped. + """ + session_id = str(uuid.uuid4()) + queue: asyncio.Queue = asyncio.Queue() + copilot._session_streams[session_id] = queue + + start = json.dumps({"id": "step-1", "name": "Reasoning: model", "payload": "**Input:**\n```python\nx\n```"}) + end = json.dumps( + { + "id": "step-1", + "name": "Reasoning: model", + "payload": "**Input:**\n```python\nx\n```\n\n**Output:** I thought about it.", + } + ) + lines = [ + f"intermediate_data: {start}", + f"intermediate_data: {end}", + 'data: {"choices":[{"delta":{"content":"done"}}]}', + "data: [DONE]", + ] + + class _Response: + status_code = 200 + + def raise_for_status(self): + return None + + async def aiter_lines(self): + for line in lines: + yield line + + class _Stream: + async def __aenter__(self): + return _Response() + + async def __aexit__(self, *args): + return False + + class _Client: + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + return False + + def stream(self, *args, **kwargs): + return _Stream() + + monkeypatch.setattr(copilot.httpx, "AsyncClient", lambda **kwargs: _Client()) + + text, _ = await copilot._invoke_copilot("https://a.test", {}, [], session_id) + + assert text == "done" + events = [] + while not queue.empty(): + events.append(queue.get_nowait()) + reasoning = [json.loads(payload) for _, payload in events] + blocks = [b for message in reasoning for b in message["message"]["content"]] + assert {"type": "reasoning", "text": "I thought about it."} in blocks + + +@pytest.mark.asyncio +async def test_invoke_copilot_emits_a_repeated_reasoning_step_once(monkeypatch: pytest.MonkeyPatch): + """A repeated completed step must not render the same trace twice, while the + start of the pair -- which shares its id -- must never claim that id.""" + session_id = str(uuid.uuid4()) + queue: asyncio.Queue = asyncio.Queue() + copilot._session_streams[session_id] = queue + + start = json.dumps({"id": "step-1", "name": "Reasoning: model", "payload": "**Input:**\n```python\nx\n```"}) + end = json.dumps( + { + "id": "step-1", + "name": "Reasoning: model", + "payload": "**Input:**\n```python\nx\n```\n\n**Output:** I thought about it.", + } + ) + lines = [ + f"intermediate_data: {start}", + f"intermediate_data: {end}", + f"intermediate_data: {end}", + 'data: {"choices":[{"delta":{"content":"done"}}]}', + "data: [DONE]", + ] + + class _Response: + status_code = 200 + + def raise_for_status(self): + return None + + async def aiter_lines(self): + for line in lines: + yield line + + class _Stream: + async def __aenter__(self): + return _Response() + + async def __aexit__(self, *args): + return False + + class _Client: + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + return False + + def stream(self, *args, **kwargs): + return _Stream() + + monkeypatch.setattr(copilot.httpx, "AsyncClient", lambda **kwargs: _Client()) + + await copilot._invoke_copilot("https://a.test", {}, [], session_id) + + blocks = [] + while not queue.empty(): + _, payload = queue.get_nowait() + blocks.extend(json.loads(payload)["message"]["content"]) + reasoning_blocks = [b for b in blocks if b.get("type") == "reasoning"] + assert reasoning_blocks == [{"type": "reasoning", "text": "I thought about it."}] diff --git a/web/packages/studio/src/routes/agents/CopilotChatRoute/stream.test.ts b/web/packages/studio/src/routes/agents/CopilotChatRoute/stream.test.ts index b0df1a8258..0679bc9816 100644 --- a/web/packages/studio/src/routes/agents/CopilotChatRoute/stream.test.ts +++ b/web/packages/studio/src/routes/agents/CopilotChatRoute/stream.test.ts @@ -253,4 +253,25 @@ describe('Copilot stream utilities', () => { loggerSpy.mockRestore(); }); + it('maps a reasoning part so the UI can render the chain of thought', () => { + const parts = getAssistantPartsFromCopilotEvent({ + type: 'assistant', + message: { + id: 'msg-reasoning', + content: [{ type: 'reasoning', text: 'The user asked a math question.' }], + }, + }); + + // Reads as ordinary narration between tool calls, like Claude's thoughts. + expect(parts).toEqual([{ type: 'text', text: 'The user asked a math question.' }]); + }); + + it('drops an empty reasoning part', () => { + const parts = getAssistantPartsFromCopilotEvent({ + type: 'assistant', + message: { id: 'msg-empty', content: [{ type: 'reasoning', text: '' }] }, + }); + + expect(parts).toEqual([]); + }); }); diff --git a/web/packages/studio/src/routes/agents/CopilotChatRoute/stream.ts b/web/packages/studio/src/routes/agents/CopilotChatRoute/stream.ts index 57ebac04da..e897804080 100644 --- a/web/packages/studio/src/routes/agents/CopilotChatRoute/stream.ts +++ b/web/packages/studio/src/routes/agents/CopilotChatRoute/stream.ts @@ -86,6 +86,12 @@ export const getAssistantPartsFromCopilotEvent = ( if (part.type === 'text' && typeof part.text === 'string') { return part.text ? { type: 'text', text: part.text } : undefined; } + // Render the model's chain of thought as ordinary assistant text, the way + // Claude's narration reads between tool calls, rather than tucking it into a + // collapsed block. + if (part.type === 'reasoning' && typeof part.text === 'string') { + return part.text ? { type: 'text', text: part.text } : undefined; + } if (part.type === 'tool_use') { const toolName = typeof part.name === 'string' ? part.name : 'tool'; diff --git a/web/packages/studio/src/routes/agents/CopilotChatRoute/useCustomAssistantChatRuntime.test.ts b/web/packages/studio/src/routes/agents/CopilotChatRoute/useCustomAssistantChatRuntime.test.ts index 50cdea1ec1..3877be8877 100644 --- a/web/packages/studio/src/routes/agents/CopilotChatRoute/useCustomAssistantChatRuntime.test.ts +++ b/web/packages/studio/src/routes/agents/CopilotChatRoute/useCustomAssistantChatRuntime.test.ts @@ -13,6 +13,7 @@ import { import { type CustomAssistantBeforeRunContext, type CustomAssistantRunContext, + type CustomAssistantRunResult, useCustomAssistantChatRuntime, } from '@studio/routes/agents/CopilotChatRoute/useCustomAssistantChatRuntime'; import { act, renderHook, waitFor } from '@testing-library/react'; @@ -44,6 +45,9 @@ const getAssistantContent = (messages: readonly ThreadMessageLike[]) => { return content; }; +const runContextOf = (onRun: ReturnType): CustomAssistantRunContext | undefined => + onRun.mock.calls[0]?.[0] as CustomAssistantRunContext | undefined; + describe('useCustomAssistantChatRuntime', () => { beforeEach(() => { mocks.useExternalStoreRuntime.mockClear(); @@ -475,4 +479,112 @@ describe('useCustomAssistantChatRuntime', () => { }); }); }); + it('starts a new assistant message for tool calls that follow a picker answer', async () => { + let runContext: CustomAssistantRunContext | undefined; + const onRun = vi.fn(async (context: CustomAssistantRunContext) => { + runContext = context; + await new Promise((resolve) => { + context.signal.addEventListener('abort', () => resolve(), { once: true }); + }); + }); + const { result } = renderHook(() => useCustomAssistantChatRuntime({ onRun })); + + act(() => { + void result.current.submitPrompt('Anonymize a dataset'); + }); + + await waitFor(() => { + expect(getMockRuntime(result.current.runtime).messages).toHaveLength(2); + }); + + const pickerPart: ThreadAssistantMessagePart = { + type: 'tool-call', + toolCallId: 'toolu_select', + toolName: 'select_dataset_file', + args: {}, + argsText: '{}', + }; + const afterPart: ThreadAssistantMessagePart = { + type: 'tool-call', + toolCallId: 'toolu_after', + toolName: 'nemo_api', + args: { resource: 'files' }, + argsText: '{"resource":"files"}', + }; + + act(() => { + runContext?.appendAssistantParts([pickerPart]); + // The picker is shown, which closes the active assistant message... + runContext?.prepareForUserInput(); + // ...but a tool step still streams in before the user answers, which + // re-attaches the run to that now-closed message. + runContext?.appendAssistantParts([pickerPart]); + }); + + act(() => { + result.current.appendUserMessage('Selected dataset: test-pi/titanic.parquet'); + runContext?.appendAssistantParts([afterPart]); + }); + + await waitFor(() => { + const messages = getMockRuntime(result.current.runtime).messages; + // The tool call that followed the answer must render below it, not merged + // back into the assistant message that came before. + expect(messages.map((message) => message.role)).toEqual([ + 'user', + 'assistant', + 'user', + 'assistant', + ]); + expect(getMessageText(messages[2]!)).toBe('Selected dataset: test-pi/titanic.parquet'); + expect(JSON.stringify(messages[3]?.content)).toContain('toolu_after'); + expect(JSON.stringify(messages[1]?.content)).not.toContain('toolu_after'); + }); + + await act(async () => { + await getMockRuntime(result.current.runtime).onCancel(); + }); + }); + it('keeps a returned result below a picker answer instead of dropping it', async () => { + let resolveRun: ((result: CustomAssistantRunResult) => void) | undefined; + const onRun = vi.fn( + () => + new Promise((resolve) => { + resolveRun = resolve; + }) + ); + const { result } = renderHook(() => useCustomAssistantChatRuntime({ onRun })); + + act(() => { + void result.current.submitPrompt('Anonymize a dataset'); + }); + + await waitFor(() => { + expect(getMockRuntime(result.current.runtime).messages).toHaveLength(2); + }); + + act(() => { + runContextOf(onRun)?.appendAssistantText('Picking a dataset.'); + // Showing the picker completes the active assistant message... + runContextOf(onRun)?.prepareForUserInput(); + // ...and the answer is appended after it. + result.current.appendUserMessage('Selected dataset: titanic.parquet'); + }); + + act(() => { + resolveRun?.({ text: 'Anonymized 100 rows.' }); + }); + + await waitFor(() => { + const messages = getMockRuntime(result.current.runtime).messages; + // The result must land in a new message below the answer, not be dropped. + expect(messages.map((message) => message.role)).toEqual([ + 'user', + 'assistant', + 'user', + 'assistant', + ]); + expect(getMessageText(messages[3]!)).toBe('Anonymized 100 rows.'); + }); + }); }); diff --git a/web/packages/studio/src/routes/agents/CopilotChatRoute/useCustomAssistantChatRuntime.ts b/web/packages/studio/src/routes/agents/CopilotChatRoute/useCustomAssistantChatRuntime.ts index 997844e44a..ed11f7cf2e 100644 --- a/web/packages/studio/src/routes/agents/CopilotChatRoute/useCustomAssistantChatRuntime.ts +++ b/web/packages/studio/src/routes/agents/CopilotChatRoute/useCustomAssistantChatRuntime.ts @@ -232,7 +232,16 @@ export const useCustomAssistantChatRuntime = ({ }; const ensureAssistantMessage = () => { - if (assistantMessageId) return; + if (assistantMessageId) { + if (messagesRef.current.at(-1)?.id === assistantMessageId) return; + // Answering a blocking picker appends a user message mid-run. The active + // assistant message now sits above it, so continuing to write into it would + // render later tool calls before the answer that triggered them. Close it + // and open a new one underneath instead. + completeActiveAssistantMessage(COMPLETE_STATUS, getCurrentResponseContent(), { + collapseCopilotContent: false, + }); + } if (resumeLastAssistantMessage()) return; createAssistantMessage(); }; @@ -339,6 +348,12 @@ export const useCustomAssistantChatRuntime = ({ return; } + if (result?.content !== undefined || result?.text !== undefined) { + // Answering a picker completes the active message, so there may be none to + // write into and the returned result would be dropped. The error path below + // calls this for the same reason. + ensureAssistantMessage(); + } completeActiveAssistantMessage( result?.status ?? COMPLETE_STATUS, result?.content ??