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
This commit is contained in:
Admire 2026-07-23 08:51:14 +08:00 committed by GitHub
parent 7b330101d2
commit a38b1daec3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
9 changed files with 929 additions and 56 deletions

View File

@ -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**:

View File

@ -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.

View File

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

View File

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

View File

@ -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 (
<ChainOfThoughtStep
key={id}
className="cursor-pointer"
className={writeFileArtifactUrl ? "cursor-pointer" : undefined}
label={resolveLabel(description)}
icon={NotebookPenIcon}
onClick={() => {
select(
new URL(
`write-file:${path}?message_id=${messageId}&tool_call_id=${id}`,
).toString(),
);
if (!writeFileArtifactUrl) {
return;
}
select(writeFileArtifactUrl);
setOpen(true);
}}
>
{path && (
{writeFilePath && (
<ChainOfThoughtSearchResult className="cursor-pointer">
{path}
{writeFilePath}
</ChainOfThoughtSearchResult>
)}
</ChainOfThoughtStep>

View File

@ -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("/");
}

View File

@ -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<void>((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<void>((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<void>((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();
});

View File

@ -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([
{

View File

@ -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("");
});
});