From a38b1daec392a335016e41f052ae42d90c9ceccd Mon Sep 17 00:00:00 2001 From: Admire <64821731+LittleChenLiya@users.noreply.github.com> Date: Thu, 23 Jul 2026 08:51:14 +0800 Subject: [PATCH] fix(streaming): keep large file generation responsive (#4354) * fix(streaming): keep large file generation responsive * fix(streaming): address follow-up review feedback * fix(streaming): address final review feedback --- backend/AGENTS.md | 1 + .../harness/deerflow/runtime/runs/worker.py | 197 +++++++-- backend/tests/test_run_worker_rollback.py | 388 +++++++++++++++++- frontend/AGENTS.md | 1 + .../workspace/messages/message-group.tsx | 69 ++-- frontend/src/core/artifacts/utils.ts | 20 + .../tests/e2e/artifact-batched-stream.spec.ts | 234 +++++++++++ .../workspace/messages/message-group.test.ts | 46 ++- .../tests/unit/core/artifacts/utils.test.ts | 29 ++ 9 files changed, 929 insertions(+), 56 deletions(-) create mode 100644 frontend/tests/e2e/artifact-batched-stream.spec.ts diff --git a/backend/AGENTS.md b/backend/AGENTS.md index 147b79a29..aaad859a7 100644 --- a/backend/AGENTS.md +++ b/backend/AGENTS.md @@ -14,6 +14,7 @@ DeerFlow is a LangGraph-based AI super agent system with a full-stack architectu **Runtime**: - `make dev`, Docker dev, and production all run the agent runtime in Gateway via `RunManager` + `run_agent()` + `StreamBridge` (`packages/harness/deerflow/runtime/`). Nginx exposes that runtime at `/api/langgraph/*` and rewrites it to Gateway's native `/api/*` routers. +- Gateway streams `write_file` and `str_replace` argument deltas in bounded batches when clients also subscribe to `values`; messages-only consumers retain the original per-chunk contract, while `values` preserves the complete tool call. - Scheduled-task executions must reuse that same Gateway run lifecycle. The scheduler may decide *when* work runs, but it must dispatch through the existing run path rather than introducing a parallel execution stack. **Project Structure**: diff --git a/backend/packages/harness/deerflow/runtime/runs/worker.py b/backend/packages/harness/deerflow/runtime/runs/worker.py index 81b860fb5..bb8fb0342 100644 --- a/backend/packages/harness/deerflow/runtime/runs/worker.py +++ b/backend/packages/harness/deerflow/runtime/runs/worker.py @@ -20,6 +20,7 @@ import copy import inspect import logging import os +import sys import threading import weakref from collections.abc import AsyncIterator @@ -103,6 +104,126 @@ async def _checkpoint_thread_lock(thread_id: str) -> AsyncIterator[None]: # Valid stream_mode values for LangGraph's graph.astream() _VALID_LG_MODES = {"values", "updates", "checkpoints", "tasks", "debug", "messages", "custom"} +# Keep this streaming policy separate from middleware write-authorization sets. +_LARGE_FILE_TOOL_NAMES = frozenset({"str_replace", "write_file"}) +_LARGE_FILE_TOOL_BATCH_SIZE = 32 + + +@dataclass +class _LargeFileToolChunkBatcher: + """Batch file-body argument deltas to avoid quadratic browser parsing. + + Normal assistant text and non-file tool calls remain token-streamed. Large + file arguments still update progressively, but in bounded batches instead + of forcing the browser to reparse the growing JSON on every model token. + """ + + batch_size: int = _LARGE_FILE_TOOL_BATCH_SIZE + tool_names: dict[tuple[str, str, str], str] = field(default_factory=dict) + pending_identity: tuple[str, str, str] | None = None + pending_message: Any | None = None + pending_metadata: dict[str, Any] = field(default_factory=dict) + pending_count: int = 0 + + def push(self, chunk: Any) -> list[Any]: + if not isinstance(chunk, tuple) or len(chunk) != 2: + return [*self.flush(), chunk] + + message, metadata = chunk + message_id = getattr(message, "id", None) + tool_call_chunks = getattr(message, "tool_call_chunks", None) + if not isinstance(message_id, str) or not message_id or not isinstance(tool_call_chunks, list) or len(tool_call_chunks) != 1: + return [*self.flush(), chunk] + + tool_chunk = tool_call_chunks[0] + if not isinstance(tool_chunk, dict): + return [*self.flush(), chunk] + index = tool_chunk.get("index") + tool_call_id = tool_chunk.get("id") + if isinstance(index, int): + discriminator = f"index:{index}" + elif isinstance(tool_call_id, str) and tool_call_id: + discriminator = f"id:{tool_call_id}" + else: + discriminator = "single" + raw_namespace = None + if isinstance(metadata, dict): + raw_namespace = metadata.get("langgraph_checkpoint_ns") or metadata.get("checkpoint_ns") + namespace = raw_namespace if isinstance(raw_namespace, str) else "" + identity = (namespace, message_id, discriminator) + name_fragment = tool_chunk.get("name") + tool_name = self.tool_names.get(identity, "") + if tool_name not in _LARGE_FILE_TOOL_NAMES and isinstance(name_fragment, str) and name_fragment: + tool_name += name_fragment + if any(candidate.startswith(tool_name) for candidate in _LARGE_FILE_TOOL_NAMES): + self.tool_names[identity] = tool_name + else: + self.tool_names.pop(identity, None) + # Batching starts only after the accumulated name matches; split or + # incomplete name fragments stream per-chunk until then. + if tool_name not in _LARGE_FILE_TOOL_NAMES: + return [*self.flush(), chunk] + + model_copy = getattr(message, "model_copy", None) + if not callable(model_copy): + return [*self.flush(), chunk] + additional_kwargs = getattr(message, "additional_kwargs", None) + sanitized_additional_kwargs = additional_kwargs + if isinstance(additional_kwargs, dict) and ("function_call" in additional_kwargs or "tool_calls" in additional_kwargs): + sanitized_additional_kwargs = {key: value for key, value in additional_kwargs.items() if key not in {"function_call", "tool_calls"}} + has_non_tool_payload = bool(getattr(message, "content", None) or sanitized_additional_kwargs or getattr(message, "usage_metadata", None) or getattr(message, "response_metadata", None)) + outputs: list[Any] = [] + if self.pending_identity is not None and self.pending_identity != identity: + outputs.extend(self.flush()) + if has_non_tool_payload: + visible_message = model_copy( + update={ + "additional_kwargs": sanitized_additional_kwargs, + "invalid_tool_calls": [], + "tool_call_chunks": [], + "tool_calls": [], + } + ) + outputs.append((visible_message, metadata)) + + tool_only_message = model_copy( + update={ + "additional_kwargs": {}, + "content": "", + "invalid_tool_calls": [], + "response_metadata": {}, + "tool_calls": [], + "usage_metadata": None, + } + ) + self.pending_identity = identity + self.pending_message = tool_only_message if self.pending_message is None else self.pending_message + tool_only_message + if isinstance(metadata, dict): + self.pending_metadata.update(metadata) + self.pending_count += 1 + if self.pending_count >= self.batch_size: + outputs.extend(self.flush()) + return outputs + + def flush(self) -> list[Any]: + if self.pending_message is None: + return [] + chunk = (self.pending_message, self.pending_metadata) + self.pending_identity = None + self.pending_message = None + self.pending_metadata = {} + self.pending_count = 0 + return [chunk] + + def finish(self) -> list[Any]: + """Flush and release identities at a values or end-of-stream boundary. + + A regular batch-size or interleaved-mode flush must retain identities + because continuation chunks commonly omit the tool name. + """ + chunks = self.flush() + self.tool_names.clear() + return chunks def _build_runtime_context( @@ -539,40 +660,58 @@ async def run_agent( async def _stream_once(input_payload: Any, stream_config: RunnableConfig) -> None: nonlocal llm_error_fallback_message - async with _checkpoint_thread_lock(thread_id): - if len(lg_modes) == 1 and not stream_subgraphs: - # Single mode, no subgraphs: astream yields raw chunks - single_mode = lg_modes[0] - async for chunk in agent.astream(input_payload, config=stream_config, stream_mode=single_mode): + file_tool_chunk_batcher = _LargeFileToolChunkBatcher() if "values" in requested_modes else None + try: + async with _checkpoint_thread_lock(thread_id): + if len(lg_modes) == 1 and not stream_subgraphs: + # Single mode, no subgraphs: astream yields raw chunks + single_mode = lg_modes[0] + async for chunk in agent.astream(input_payload, config=stream_config, stream_mode=single_mode): + if record.abort_event.is_set(): + logger.info("Run %s abort requested — stopping", run_id) + break + llm_error_fallback_message = llm_error_fallback_message or _extract_llm_error_fallback_message(chunk, pre_existing_message_ids) + sse_event = _lg_mode_to_sse_event(single_mode) + await bridge.publish(run_id, sse_event, serialize(chunk, mode=single_mode)) + if single_mode == "custom": + await subagent_events.add(chunk) + return + # Multiple modes or subgraphs: astream yields tuples + async for item in agent.astream( + input_payload, + config=stream_config, + stream_mode=lg_modes, + subgraphs=stream_subgraphs, + ): if record.abort_event.is_set(): logger.info("Run %s abort requested — stopping", run_id) break + + mode, chunk = _unpack_stream_item(item, lg_modes, stream_subgraphs) + if mode is None: + continue + llm_error_fallback_message = llm_error_fallback_message or _extract_llm_error_fallback_message(chunk, pre_existing_message_ids) - sse_event = _lg_mode_to_sse_event(single_mode) - await bridge.publish(run_id, sse_event, serialize(chunk, mode=single_mode)) - if single_mode == "custom": + sse_event = _lg_mode_to_sse_event(mode) + if file_tool_chunk_batcher is not None and mode != "messages": + pending_chunks = file_tool_chunk_batcher.finish() if mode == "values" else file_tool_chunk_batcher.flush() + for publish_chunk in pending_chunks: + await bridge.publish(run_id, "messages", serialize(publish_chunk, mode="messages")) + chunks_to_publish = file_tool_chunk_batcher.push(chunk) if mode == "messages" and file_tool_chunk_batcher is not None else [chunk] + for publish_chunk in chunks_to_publish: + await bridge.publish(run_id, sse_event, serialize(publish_chunk, mode=mode)) + if mode == "custom": await subagent_events.add(chunk) - return - # Multiple modes or subgraphs: astream yields tuples - async for item in agent.astream( - input_payload, - config=stream_config, - stream_mode=lg_modes, - subgraphs=stream_subgraphs, - ): - if record.abort_event.is_set(): - logger.info("Run %s abort requested — stopping", run_id) - break - - mode, chunk = _unpack_stream_item(item, lg_modes, stream_subgraphs) - if mode is None: - continue - - llm_error_fallback_message = llm_error_fallback_message or _extract_llm_error_fallback_message(chunk, pre_existing_message_ids) - sse_event = _lg_mode_to_sse_event(mode) - await bridge.publish(run_id, sse_event, serialize(chunk, mode=mode)) - if mode == "custom": - await subagent_events.add(chunk) + finally: + stream_error = sys.exception() + if file_tool_chunk_batcher is not None: + try: + for publish_chunk in file_tool_chunk_batcher.finish(): + await bridge.publish(run_id, "messages", serialize(publish_chunk, mode="messages")) + except Exception: + if stream_error is None: + raise + logger.debug("Could not flush pending file-tool chunks for run %s", run_id, exc_info=True) # 7. Stream the requested turn, then optionally continue hidden goal turns. # Clear any stale stop_reason before the first (user-visible) turn only. diff --git a/backend/tests/test_run_worker_rollback.py b/backend/tests/test_run_worker_rollback.py index 9d498f84e..9082f2c20 100644 --- a/backend/tests/test_run_worker_rollback.py +++ b/backend/tests/test_run_worker_rollback.py @@ -6,7 +6,7 @@ from typing import Annotated, Any, NotRequired, TypedDict from unittest.mock import AsyncMock, call, patch import pytest -from langchain_core.messages import AIMessage, AnyMessage, HumanMessage +from langchain_core.messages import AIMessage, AIMessageChunk, AnyMessage, HumanMessage from langgraph.channels.delta import DeltaChannel from langgraph.checkpoint.base import empty_checkpoint from langgraph.checkpoint.memory import InMemorySaver @@ -31,6 +31,7 @@ from deerflow.runtime.runs.worker import ( _ensure_interrupted_title, _extract_llm_error_fallback_message, _install_runtime_context, + _LargeFileToolChunkBatcher, _rollback_to_pre_run_checkpoint, _try_extract_from_message, run_agent, @@ -93,6 +94,197 @@ def _build_message_append_graph(state_schema: type, checkpointer: Any): return builder.compile(checkpointer=checkpointer) +@pytest.mark.parametrize("tool_name", ["write_file", "str_replace"]) +def test_large_file_tool_chunk_batcher_streams_bounded_batches(tool_name: str): + batcher = _LargeFileToolChunkBatcher(batch_size=2) + first = AIMessageChunk( + content="", + id="ai-1", + tool_call_chunks=[ + { + "id": "call-1", + "index": 0, + "name": tool_name, + "args": '{"path":"/mnt/user-data/outputs/report.md","content":"Hel', + } + ], + ) + continuation = AIMessageChunk( + content="", + id="ai-1", + tool_call_chunks=[{"index": 0, "name": None, "args": 'lo"}'}], + ) + + assert batcher.push((first, {})) == [] + published = batcher.push((continuation, {})) + assert len(published) == 1 + message, metadata = published[0] + assert metadata == {} + assert message.tool_calls[0]["args"]["content"] == "Hello" + assert batcher.flush() == [] + + +def test_large_file_tool_chunk_batcher_preserves_visible_and_non_file_chunks(): + batcher = _LargeFileToolChunkBatcher() + visible_text = AIMessageChunk(content="Writing the report now.", id="ai-1") + search_tool = AIMessageChunk( + content="", + id="ai-2", + tool_call_chunks=[ + { + "id": "call-2", + "index": 0, + "name": "web_search", + "args": '{"query":"vector databases"}', + } + ], + ) + write_with_reasoning = AIMessageChunk( + content="", + id="ai-3", + additional_kwargs={"reasoning_content": "Choosing a filename."}, + tool_call_chunks=[ + { + "id": "call-3", + "index": 0, + "name": "write_file", + "args": '{"path":"/mnt/user-data/outputs/report.md"}', + } + ], + ) + + assert batcher.push((visible_text, {})) == [(visible_text, {})] + assert batcher.push((search_tool, {})) == [(search_tool, {})] + visible_reasoning = batcher.push((write_with_reasoning, {})) + assert len(visible_reasoning) == 1 + filtered_message, filtered_metadata = visible_reasoning[0] + assert filtered_metadata == {} + assert filtered_message.additional_kwargs == {"reasoning_content": "Choosing a filename."} + assert filtered_message.tool_call_chunks == [] + pending_file_chunks = batcher.flush() + assert len(pending_file_chunks) == 1 + assert pending_file_chunks[0][0].tool_call_chunks[0]["name"] == "write_file" + + +def test_large_file_tool_chunk_batcher_separates_subgraph_namespaces(): + batcher = _LargeFileToolChunkBatcher() + first = AIMessageChunk( + content="", + id="shared-ai-id", + tool_call_chunks=[{"id": "call-a", "index": 0, "name": "write_file", "args": '{"path":"a.md","content":"A'}], + ) + second = AIMessageChunk( + content="", + id="shared-ai-id", + tool_call_chunks=[{"id": "call-b", "index": 0, "name": "write_file", "args": '{"path":"b.md","content":"B'}], + ) + + assert batcher.push((first, {"langgraph_checkpoint_ns": "task-a"})) == [] + published = batcher.push((second, {"langgraph_checkpoint_ns": "task-b"})) + + assert len(published) == 1 + assert published[0][1]["langgraph_checkpoint_ns"] == "task-a" + assert batcher.flush()[0][1]["langgraph_checkpoint_ns"] == "task-b" + + +@pytest.mark.parametrize("metadata", [None, "not-a-dict"]) +def test_large_file_tool_chunk_batcher_accepts_non_dict_metadata(metadata: Any): + batcher = _LargeFileToolChunkBatcher(batch_size=1) + message = AIMessageChunk( + content="", + id="ai-file", + tool_call_chunks=[ + { + "id": "call-file", + "index": 0, + "name": "write_file", + "args": '{"path":"report.md","content":"draft"}', + } + ], + ) + + published = batcher.push((message, metadata)) + + assert len(published) == 1 + assert published[0][0].tool_call_chunks[0]["args"] == '{"path":"report.md","content":"draft"}' + assert published[0][1] == {} + + +def test_large_file_tool_chunk_batcher_does_not_retain_non_file_names(): + batcher = _LargeFileToolChunkBatcher() + + for index in range(100): + message = AIMessageChunk( + content="", + id=f"ai-{index}", + tool_call_chunks=[ + { + "id": f"call-{index}", + "index": 0, + "name": "web_search", + "args": '{"query":"deerflow"}', + } + ], + ) + + assert batcher.push((message, {})) == [(message, {})] + + assert batcher.tool_names == {} + + +def test_large_file_tool_chunk_batcher_starts_batching_after_split_name_matches(): + batcher = _LargeFileToolChunkBatcher(batch_size=1) + name_prefix = AIMessageChunk( + content="", + id="ai-file", + tool_call_chunks=[{"id": "call-file", "index": 0, "name": "write_", "args": ""}], + ) + name_suffix = AIMessageChunk( + content="", + id="ai-file", + tool_call_chunks=[ + { + "index": 0, + "name": "file", + "args": '{"path":"report.md","content":"draft"}', + } + ], + ) + + assert batcher.push((name_prefix, {})) == [(name_prefix, {})] + assert set(batcher.tool_names.values()) == {"write_"} + assert len(batcher.push((name_suffix, {}))) == 1 + assert set(batcher.tool_names.values()) == {"write_file"} + + +def test_large_file_tool_chunk_batcher_keeps_identity_across_batches_then_releases_it(): + batcher = _LargeFileToolChunkBatcher(batch_size=1) + first = AIMessageChunk( + content="", + id="ai-file", + tool_call_chunks=[ + { + "id": "call-file", + "index": 0, + "name": "write_file", + "args": '{"path":"report.md","content":"Hel', + } + ], + ) + continuation = AIMessageChunk( + content="", + id="ai-file", + tool_call_chunks=[{"index": 0, "name": None, "args": 'lo"}'}], + ) + + assert len(batcher.push((first, {}))) == 1 + assert len(batcher.push((continuation, {}))) == 1 + assert set(batcher.tool_names.values()) == {"write_file"} + + assert batcher.finish() == [] + assert batcher.tool_names == {} + + def test_build_runtime_context_includes_app_config_when_present(): app_config = object() @@ -136,6 +328,200 @@ def test_install_runtime_context_overrides_internal_pre_existing_message_ids(): assert config["context"][CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY] == frozenset({"old-ai"}) +@pytest.mark.anyio +async def test_run_agent_batches_incremental_file_args_and_keeps_complete_values(): + run_manager = RunManager() + record = await run_manager.create("thread-file-stream") + bridge = SimpleNamespace( + publish=AsyncMock(), + publish_end=AsyncMock(), + cleanup=AsyncMock(), + ) + complete_message = AIMessage( + content="", + id="ai-file", + tool_calls=[ + { + "id": "call-file", + "name": "write_file", + "args": { + "path": "/mnt/user-data/outputs/report.md", + "content": "Hello world", + }, + "type": "tool_call", + } + ], + ) + + class DummyAgent: + async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False): + del graph_input, config, stream_mode, subgraphs + yield ( + "messages", + ( + AIMessageChunk( + content="", + id="ai-file", + tool_call_chunks=[ + { + "id": "call-file", + "index": 0, + "name": "write_file", + "args": '{"path":"/mnt/user-data/outputs/report.md","content":"Hello', + } + ], + ), + {}, + ), + ) + yield ( + "messages", + ( + AIMessageChunk( + content="", + id="ai-file", + tool_call_chunks=[{"index": 0, "name": None, "args": ' world"}'}], + ), + {}, + ), + ) + yield ("values", {"messages": [complete_message]}) + + await run_agent( + bridge, + run_manager, + record, + ctx=RunContext(checkpointer=None), + agent_factory=lambda **_kwargs: DummyAgent(), + graph_input={}, + config={}, + stream_modes=["messages-tuple", "values"], + ) + + message_events = [call.args for call in bridge.publish.await_args_list if call.args[1] == "messages"] + assert len(message_events) == 1 + assert message_events[0][2][0]["tool_calls"][0]["args"]["content"] == "Hello world" + values_events = [call.args[2] for call in bridge.publish.await_args_list if call.args[1] == "values"] + assert any(event["messages"][0]["tool_calls"][0]["args"]["content"] == "Hello world" for event in values_events) + + +@pytest.mark.parametrize( + ("stream_error", "flush_publish_error", "expected_error"), + [ + (True, False, "stream failed"), + (True, True, "stream failed"), + (False, True, "flush publish failed"), + ], +) +@pytest.mark.anyio +async def test_run_agent_handles_pending_file_args_when_stream_or_flush_raises(stream_error: bool, flush_publish_error: bool, expected_error: str): + run_manager = RunManager() + record = await run_manager.create("thread-file-stream-error") + + async def publish(_run_id: str, event: str, _data: Any): + if flush_publish_error and event == "messages": + raise RuntimeError("flush publish failed") + + bridge = SimpleNamespace( + publish=AsyncMock(side_effect=publish), + publish_end=AsyncMock(), + cleanup=AsyncMock(), + ) + + class DummyAgent: + async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False): + del graph_input, config, stream_mode, subgraphs + yield ( + "messages", + ( + AIMessageChunk( + content="", + id="ai-file", + tool_call_chunks=[ + { + "id": "call-file", + "index": 0, + "name": "write_file", + "args": '{"path":"report.md","content":"partial', + } + ], + ), + {}, + ), + ) + if stream_error: + raise RuntimeError("stream failed") + + await run_agent( + bridge, + run_manager, + record, + ctx=RunContext(checkpointer=None), + agent_factory=lambda **_kwargs: DummyAgent(), + graph_input={}, + config={}, + stream_modes=["messages-tuple", "values"], + ) + + message_events = [call.args for call in bridge.publish.await_args_list if call.args[1] == "messages"] + assert len(message_events) == 1 + assert message_events[0][2][0]["tool_call_chunks"][0]["args"].endswith('"content":"partial') + error_events = [call.args for call in bridge.publish.await_args_list if call.args[1] == "error"] + assert error_events[0][2]["message"] == expected_error + + +@pytest.mark.anyio +async def test_run_agent_keeps_file_chunks_unbatched_without_values_mode(): + run_manager = RunManager() + record = await run_manager.create("thread-file-messages-only") + bridge = SimpleNamespace( + publish=AsyncMock(), + publish_end=AsyncMock(), + cleanup=AsyncMock(), + ) + chunks = [ + AIMessageChunk( + content="", + id="ai-file", + tool_call_chunks=[ + { + "id": "call-file", + "index": 0, + "name": "write_file", + "args": '{"path":"report.md","content":"Hel', + } + ], + ), + AIMessageChunk( + content="", + id="ai-file", + tool_call_chunks=[{"index": 0, "name": None, "args": 'lo"}'}], + ), + ] + + class DummyAgent: + async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False): + del graph_input, config, stream_mode, subgraphs + for chunk in chunks: + yield (chunk, {}) + + await run_agent( + bridge, + run_manager, + record, + ctx=RunContext(checkpointer=None), + agent_factory=lambda **_kwargs: DummyAgent(), + graph_input={}, + config={}, + stream_modes=["messages-tuple"], + ) + + message_events = [call.args[2] for call in bridge.publish.await_args_list if call.args[1] == "messages"] + assert len(message_events) == 2 + assert message_events[0][0]["tool_call_chunks"][0]["args"].endswith('"content":"Hel') + assert message_events[1][0]["tool_call_chunks"][0]["args"] == 'lo"}' + + @pytest.mark.anyio async def test_run_agent_threads_explicit_app_config_into_config_only_factory(): run_manager = RunManager() diff --git a/frontend/AGENTS.md b/frontend/AGENTS.md index e0864427d..dd960b549 100644 --- a/frontend/AGENTS.md +++ b/frontend/AGENTS.md @@ -65,6 +65,7 @@ The frontend is a stateful chat application. Users create **threads** (conversat 1. Optional composer helpers such as `core/input-polish` can rewrite the local draft before submission, and `core/voice-input` can transcribe browser microphone input into that same local draft; confirmed user input then flows to thread hooks (`core/threads/hooks.ts`) → LangGraph SDK streaming 2. Stream events update thread state (messages, artifacts, todos, goal) + File-tool artifact auto-open work must run in an effect with timer cleanup; never schedule timers while rendering streamed `write_file` or `str_replace` updates. 3. `useThreadHistory` loads persisted conversation pages from `GET /api/threads/{id}/messages/page`, preserving the backend's thread-global event `seq`; rendering overlays checkpoint/live copies at their matching canonical identities (a summarized checkpoint may contain a protected early input plus a recent tail), suppresses checkpoint/transient prefixes whose canonical position is still behind an unloaded cursor page instead of collapsing that unknown gap before a recent anchor, then adds optimistic messages without timestamp re-sorting. History invalidation preserves already-loaded pages so their established ordering positions are not discarded. 4. Stop actions call the LangGraph SDK stream stop path; `core/threads/hooks.ts` invalidates current-thread, thread-history, token-usage, and sidebar/search caches immediately and schedules one follow-up refetch because SDK stop may finish via abort + fire-and-forget cancel before backend title finalization commits 5. TanStack Query manages server state; localStorage stores user settings diff --git a/frontend/src/components/workspace/messages/message-group.tsx b/frontend/src/components/workspace/messages/message-group.tsx index e75106b60..a12096c00 100644 --- a/frontend/src/components/workspace/messages/message-group.tsx +++ b/frontend/src/components/workspace/messages/message-group.tsx @@ -15,7 +15,7 @@ import { SquareTerminalIcon, WrenchIcon, } from "lucide-react"; -import { memo, useMemo, useState } from "react"; +import { memo, useEffect, useMemo, useState } from "react"; import { ChainOfThought, @@ -26,7 +26,10 @@ import { } from "@/components/ai-elements/chain-of-thought"; import { CodeBlock } from "@/components/ai-elements/code-block"; import { Button } from "@/components/ui/button"; -import { resolveArtifactURL } from "@/core/artifacts/utils"; +import { + buildWriteFileArtifactURL, + resolveArtifactURL, +} from "@/core/artifacts/utils"; import { useI18n } from "@/core/i18n/hooks"; import { formatTokenCount } from "@/core/messages/usage"; import type { TokenDebugStep } from "@/core/messages/usage-model"; @@ -554,6 +557,40 @@ function ToolCall({ ) : ( fallback ); + const writeFilePath = + (name === "write_file" || name === "str_replace") && + typeof args.path === "string" + ? args.path + : undefined; + const writeFileArtifactUrl = writeFilePath + ? buildWriteFileArtifactURL({ + filepath: writeFilePath, + messageId, + toolCallId: id, + }) + : null; + const autoOpenArtifactUrl = + isLoading && + isLast && + autoOpen && + autoSelect && + writeFileArtifactUrl && + !result + ? writeFileArtifactUrl + : null; + + useEffect(() => { + if (!autoOpenArtifactUrl || selectedArtifact === autoOpenArtifactUrl) { + return; + } + + const timeout = window.setTimeout(() => { + select(autoOpenArtifactUrl, true); + setOpen(true); + }, 100); + + return () => window.clearTimeout(timeout); + }, [autoOpenArtifactUrl, select, selectedArtifact, setOpen]); if (name.startsWith("browser_")) { const shot = browserView?.screenshot; @@ -749,38 +786,24 @@ function ToolCall({ if (!description) { description = t.toolCalls.writeFile; } - const path: string | undefined = (args as { path: string })?.path; - if (isLoading && isLast && autoOpen && autoSelect && path && !result) { - setTimeout(() => { - const url = new URL( - `write-file:${path}?message_id=${messageId}&tool_call_id=${id}`, - ).toString(); - if (selectedArtifact === url) { - return; - } - select(url, true); - setOpen(true); - }, 100); - } return ( { - select( - new URL( - `write-file:${path}?message_id=${messageId}&tool_call_id=${id}`, - ).toString(), - ); + if (!writeFileArtifactUrl) { + return; + } + select(writeFileArtifactUrl); setOpen(true); }} > - {path && ( + {writeFilePath && ( - {path} + {writeFilePath} )} diff --git a/frontend/src/core/artifacts/utils.ts b/frontend/src/core/artifacts/utils.ts index 324d51925..eb31d3256 100644 --- a/frontend/src/core/artifacts/utils.ts +++ b/frontend/src/core/artifacts/utils.ts @@ -27,6 +27,26 @@ function encodeArtifactPath(filepath: string) { .join("/"); } +export function buildWriteFileArtifactURL({ + filepath, + messageId, + toolCallId, +}: { + filepath: string; + messageId?: string; + toolCallId?: string; +}) { + const url = new URL("write-file:/"); + url.pathname = filepath.replaceAll("%", "%25"); + if (messageId) { + url.searchParams.set("message_id", messageId); + } + if (toolCallId) { + url.searchParams.set("tool_call_id", toolCallId); + } + return url.toString(); +} + function decodeRelativeArtifactPath(filepath: string) { return filepath.split("/").map(decodePathSegment).join("/"); } diff --git a/frontend/tests/e2e/artifact-batched-stream.spec.ts b/frontend/tests/e2e/artifact-batched-stream.spec.ts new file mode 100644 index 000000000..e705a934a --- /dev/null +++ b/frontend/tests/e2e/artifact-batched-stream.spec.ts @@ -0,0 +1,234 @@ +import { createServer } from "node:http"; +import type { AddressInfo } from "node:net"; + +import { expect, test } from "@playwright/test"; + +import { mockLangGraphAPI } from "./utils/mock-api"; + +const THREAD_ID = "00000000-0000-0000-0000-000000004354"; +const RUN_ID = "00000000-0000-0000-0000-000000004355"; +const MISSING_PATH_THREAD_ID = "00000000-0000-0000-0000-000000004356"; +const ARTIFACT_PATH = "/artifact-fixtures/batched-report.md"; + +const INITIAL_MESSAGES = [ + { + type: "human", + id: "msg-human-batched-artifact", + content: [{ type: "text", text: "Create a batched markdown report" }], + }, +]; + +function batchedWriteFileStreamFrames() { + const chunks = [ + { + content: "", + additional_kwargs: {}, + response_metadata: {}, + type: "AIMessageChunk", + name: null, + id: "msg-ai-batched-artifact", + tool_calls: [ + { + name: "write_file", + args: { path: ARTIFACT_PATH, content: "Hello " }, + id: "call-batched-artifact", + type: "tool_call", + }, + ], + invalid_tool_calls: [], + usage_metadata: null, + tool_call_chunks: [ + { + name: "write_file", + args: `{"path":"${ARTIFACT_PATH}","content":"Hello `, + id: "call-batched-artifact", + index: 0, + type: "tool_call_chunk", + }, + ], + chunk_position: null, + }, + { + content: "", + additional_kwargs: {}, + response_metadata: {}, + type: "AIMessageChunk", + name: null, + id: "msg-ai-batched-artifact", + tool_calls: [], + invalid_tool_calls: [ + { + name: null, + args: 'world"}', + id: null, + error: null, + type: "invalid_tool_call", + }, + ], + usage_metadata: null, + tool_call_chunks: [ + { + name: null, + args: 'world"}', + id: null, + index: 0, + type: "tool_call_chunk", + }, + ], + chunk_position: null, + }, + ]; + const events = [ + { + event: "metadata", + data: { run_id: RUN_ID, thread_id: THREAD_ID }, + }, + { + event: "values", + data: { + messages: [ + ...INITIAL_MESSAGES, + { + type: "human", + id: "msg-human-batched-artifact-follow-up", + content: [{ type: "text", text: "Continue the report" }], + }, + ], + }, + }, + ...chunks.map((chunk) => ({ event: "messages", data: [chunk, {}] })), + ]; + + return events.map( + (event) => `event: ${event.event}\ndata: ${JSON.stringify(event.data)}\n\n`, + ); +} + +async function startBatchedWriteFileStreamServer() { + const frames = batchedWriteFileStreamFrames(); + const server = createServer((_request, response) => { + response.writeHead(200, { + "Access-Control-Allow-Origin": "*", + "Cache-Control": "no-cache", + "Content-Type": "text/event-stream", + }); + response.write(frames.slice(0, 3).join("")); + + const nextBatch = setTimeout(() => { + response.write(frames[3]); + }, 300); + const finishStream = setTimeout(() => { + response.end(); + }, 2_000); + response.once("close", () => { + clearTimeout(nextBatch); + clearTimeout(finishStream); + }); + }); + + await new Promise((resolve, reject) => { + const handleError = (error: Error) => reject(error); + server.once("error", handleError); + server.listen(0, "127.0.0.1", () => { + server.off("error", handleError); + resolve(); + }); + }); + + const { port } = server.address() as AddressInfo; + return { + url: `http://127.0.0.1:${port}/runs/stream`, + async close() { + server.closeAllConnections(); + await new Promise((resolve, reject) => { + server.close((error) => (error ? reject(error) : resolve())); + }); + }, + }; +} + +test("assembles streamed write-file argument deltas in the artifact preview", async ({ + page, +}) => { + let streamStarted = false; + let releasePostStreamHistory!: () => void; + const postStreamHistoryReleased = new Promise((resolve) => { + releasePostStreamHistory = resolve; + }); + const streamServer = await startBatchedWriteFileStreamServer(); + mockLangGraphAPI(page, { + threads: [ + { + thread_id: THREAD_ID, + title: "Batched artifact streaming", + messages: INITIAL_MESSAGES, + }, + ], + }); + await page.route("**/api/langgraph/threads/*/history", async (route) => { + if (streamStarted) { + await postStreamHistoryReleased; + } + return route.fallback(); + }); + await page.route("**/api/langgraph/threads/*/runs/stream", (route) => { + streamStarted = true; + return route.continue({ url: streamServer.url }); + }); + + try { + await page.goto(`/workspace/chats/${THREAD_ID}`); + + const textarea = page.getByPlaceholder(/how can i assist you/i); + await expect(textarea).toBeVisible({ timeout: 15_000 }); + await textarea.fill("Continue the report"); + await textarea.press("Enter"); + + await expect(page.getByText(ARTIFACT_PATH)).toBeVisible({ + timeout: 10_000, + }); + + const artifactsPanel = page.locator("#artifacts"); + await expect(artifactsPanel).toBeVisible(); + await expect(artifactsPanel.getByText("batched-report.md")).toBeVisible(); + await expect(artifactsPanel.getByText("Hello world")).toBeVisible(); + } finally { + releasePostStreamHistory(); + await streamServer.close(); + } +}); + +test("does not open an artifact for a file tool call without a path", async ({ + page, +}) => { + mockLangGraphAPI(page, { + threads: [ + { + thread_id: MISSING_PATH_THREAD_ID, + title: "File tool without a path", + messages: [ + ...INITIAL_MESSAGES, + { + type: "ai", + id: "msg-ai-missing-path", + content: "", + tool_calls: [ + { + id: "call-missing-path", + name: "write_file", + args: { description: "Write file" }, + }, + ], + }, + ], + }, + ], + }); + + await page.goto(`/workspace/chats/${MISSING_PATH_THREAD_ID}`); + + const writeFileStep = page.getByText("Write file", { exact: true }); + await expect(writeFileStep).toBeVisible({ timeout: 15_000 }); + await writeFileStep.click(); + await expect(page.locator("#artifacts")).toBeHidden(); +}); diff --git a/frontend/tests/unit/components/workspace/messages/message-group.test.ts b/frontend/tests/unit/components/workspace/messages/message-group.test.ts index 4b21c97b1..ad2eec510 100644 --- a/frontend/tests/unit/components/workspace/messages/message-group.test.ts +++ b/frontend/tests/unit/components/workspace/messages/message-group.test.ts @@ -1,25 +1,36 @@ import type { Message } from "@langchain/langgraph-sdk"; -import { describe, expect, it, rs } from "@rstest/core"; +import { afterEach, describe, expect, it, rs } from "@rstest/core"; import { createElement, type ComponentProps } from "react"; import { renderToStaticMarkup } from "react-dom/server"; import { MessageGroup } from "@/components/workspace/messages/message-group"; import { I18nContext } from "@/core/i18n/context"; +const artifactsMockState = rs.hoisted(() => ({ + autoOpen: false, + autoSelect: false, +})); + rs.mock("@/components/workspace/artifacts", () => ({ useArtifacts: () => ({ artifacts: [], setArtifacts: () => undefined, selectedArtifact: null, - autoSelect: false, + autoSelect: artifactsMockState.autoSelect, select: () => undefined, deselect: () => undefined, open: false, - autoOpen: false, + autoOpen: artifactsMockState.autoOpen, setOpen: () => undefined, }), })); +afterEach(() => { + artifactsMockState.autoOpen = false; + artifactsMockState.autoSelect = false; + rs.restoreAllMocks(); +}); + describe("MessageGroup", () => { it("renders assistant text attached to a tool-calling processing message", () => { const html = renderGroup([ @@ -92,6 +103,35 @@ describe("MessageGroup", () => { expect(html).toContain("1 more step"); }); + it("does not schedule artifact auto-open during render", () => { + artifactsMockState.autoOpen = true; + artifactsMockState.autoSelect = true; + const timeoutSpy = rs.spyOn(globalThis, "setTimeout"); + const html = renderGroup( + [ + { + id: "ai-write", + type: "ai", + content: "", + tool_calls: [ + { + id: "call-write", + name: "write_file", + args: { + path: "/mnt/user-data/outputs/report.md", + content: "# Report", + }, + }, + ], + } as Message, + ], + { isLoading: true }, + ); + + expect(html).toContain("/mnt/user-data/outputs/report.md"); + expect(timeoutSpy).not.toHaveBeenCalled(); + }); + it("keeps tool-calling assistant text visible when reasoning is also present", () => { const html = renderGroup([ { diff --git a/frontend/tests/unit/core/artifacts/utils.test.ts b/frontend/tests/unit/core/artifacts/utils.test.ts index 8c2aeced3..edb8b35dd 100644 --- a/frontend/tests/unit/core/artifacts/utils.test.ts +++ b/frontend/tests/unit/core/artifacts/utils.test.ts @@ -223,4 +223,33 @@ describe("artifact URL helpers", () => { ]), ).toBe("https://example.com/image.png"); }); + + test("builds encoded write-file URLs without undefined query parameters", async () => { + const { buildWriteFileArtifactURL } = await loadFreshArtifactUtils(); + const filepath = "/mnt/user-data/outputs/a b#c?%20.md"; + expect( + buildWriteFileArtifactURL({ + filepath: "/mnt/user-data/outputs/report.md", + messageId: "ai-1", + toolCallId: "call-1", + }), + ).toBe( + "write-file:/mnt/user-data/outputs/report.md?message_id=ai-1&tool_call_id=call-1", + ); + + const withIds = new URL( + buildWriteFileArtifactURL({ + filepath, + messageId: "message #1", + toolCallId: "call ?1", + }), + ); + expect(decodeURIComponent(withIds.pathname)).toBe(filepath); + expect(withIds.searchParams.get("message_id")).toBe("message #1"); + expect(withIds.searchParams.get("tool_call_id")).toBe("call ?1"); + + const withoutIds = buildWriteFileArtifactURL({ filepath }); + expect(withoutIds).not.toContain("undefined"); + expect(new URL(withoutIds).search).toBe(""); + }); });