From 186c6ea463ff85b467651185a258773ca1779b8d Mon Sep 17 00:00:00 2001 From: Ryker_Feng <90562015+18062706139fcz@users.noreply.github.com> Date: Mon, 6 Jul 2026 07:40:41 +0800 Subject: [PATCH] feat: add branching support for assistant turns (#3950) * feat: add assistant turn branching * fix(threads): skip workspace clone when branching from historical turn Workspace files are not checkpointed, so cloning them onto a branch rooted at an older assistant turn leaked files created in a later timeline. Restrict the best-effort workspace copy to branches taken from the latest turn; historical-turn branches now report workspace_clone_mode="skipped_historical_turn" and keep only the restored message history. * style(frontend): fix prettier formatting in e2e mock-api Collapse the branch-title normalization chain onto a single line to satisfy the frontend lint (prettier --check) CI gate. --------- Co-authored-by: Willem Jiang --- README.md | 2 + README_zh.md | 2 + backend/AGENTS.md | 2 +- backend/app/gateway/routers/threads.py | 266 ++++++++++++++++++ backend/tests/test_threads_router.py | 239 ++++++++++++++++ frontend/AGENTS.md | 1 + .../app/workspace/chats/[thread_id]/page.tsx | 38 +++ .../workspace/messages/message-list.tsx | 62 +++- frontend/src/core/i18n/locales/en-US.ts | 3 + frontend/src/core/i18n/locales/types.ts | 3 + frontend/src/core/i18n/locales/zh-CN.ts | 3 + frontend/src/core/threads/api.ts | 57 ++++ frontend/src/core/threads/hooks.ts | 31 +- frontend/tests/e2e/branch-thread.spec.ts | 66 +++++ frontend/tests/e2e/utils/mock-api.ts | 73 +++++ frontend/tests/unit/core/threads/api.test.ts | 59 ++++ 16 files changed, 896 insertions(+), 11 deletions(-) create mode 100644 frontend/tests/e2e/branch-thread.spec.ts diff --git a/README.md b/README.md index 132388721..c02f476bb 100644 --- a/README.md +++ b/README.md @@ -630,6 +630,8 @@ Gateway-generated follow-up suggestions now normalize both plain-string model ou Interrupted first-turn runs still persist a fallback conversation title, so stopping a streaming response does not leave the thread as "Untitled" after refresh. +In the Web UI, completed assistant turns can be branched into a new main conversation. The new thread starts from that turn's checkpoint. Because workspace files are not checkpointed, the branch only receives a best-effort copy of the current workspace when you branch from the latest turn; branching from an older turn keeps just the restored message history so the branch never inherits files that were created in a later part of the conversation. + ``` # Paths inside the sandbox container /mnt/skills/public diff --git a/README_zh.md b/README_zh.md index e3824eb23..4103136f6 100644 --- a/README_zh.md +++ b/README_zh.md @@ -484,6 +484,8 @@ Tools 也是同样的思路。DeerFlow 自带一组核心工具：网页搜索 Gateway 生成后续建议时，现在会先把普通字符串输出和 block/list 风格的富文本内容统一归一化，再去解析 JSON 数组响应，因此不同 provider 的内容包装方式不会再悄悄把建议吞掉。 +Web UI 支持从已完成的 assistant 回复分叉出一个新的主对话。新 thread 会从该回复对应的 checkpoint 开始，并尽力复制当前 thread 的工作区文件。 + ```text # sandbox 容器内的路径 /mnt/skills/public diff --git a/backend/AGENTS.md b/backend/AGENTS.md index 9069a2e07..483c5f5d7 100644 --- a/backend/AGENTS.md +++ b/backend/AGENTS.md @@ -306,7 +306,7 @@ CORS is same-origin by default when requests enter through nginx on port 2026. S | **Skills** (`/api/skills`) | `GET /` - list skills; `GET /{name}` - details; `PUT /{name}` - update enabled; `POST /install` - install from .skill archive (accepts standard optional frontmatter like `version`, `author`, `compatibility`) | | **Memory** (`/api/memory`) | `GET /` - memory data; `POST /reload` - force reload; `GET /config` - config; `GET /status` - config + data | | **Uploads** (`/api/threads/{id}/uploads`) | `POST /` - upload files (auto-converts PDF/PPT/Excel/Word); `GET /list` - list; `DELETE /{filename}` - delete | -| **Threads** (`/api/threads/{id}`) | `DELETE /` - remove DeerFlow-managed local thread data after LangGraph thread deletion; `GET /goal`, `PUT /goal`, `DELETE /goal` - read, set, and clear the active thread goal; unexpected failures are logged server-side and return a generic 500 detail | +| **Threads** (`/api/threads/{id}`) | `DELETE /` - remove DeerFlow-managed local thread data after LangGraph thread deletion; `POST /branches` - create a new main-thread branch from a completed assistant turn checkpoint. Workspace files are not checkpointed, so the branch only best-effort copies the current workspace when branching from the **latest** turn (`workspace_clone_mode="current_thread_best_effort"`); branching from an older/historical turn skips the copy (`workspace_clone_mode="skipped_historical_turn"`) so the branch never inherits files that only exist in a later timeline; `GET /goal`, `PUT /goal`, `DELETE /goal` - read, set, and clear the active thread goal; unexpected failures are logged server-side and return a generic 500 detail | | **Artifacts** (`/api/threads/{id}/artifacts`) | `GET /{path}` - serve artifacts; active content types (`text/html`, `application/xhtml+xml`, `image/svg+xml`) are always forced as download attachments to reduce XSS risk; `?download=true` still forces download for other file types | | **Suggestions** (`/api/suggestions`) | `GET /config` - returns global suggestions config boolean; `POST /threads/{id}/suggestions` - generate follow-up questions; rich list/block model content is normalized and inline reasoning (`...`, including unclosed/truncated blocks from reasoning models like MiniMax-M3) is stripped before JSON parsing | | **Thread Runs** (`/api/threads/{id}/runs`) | `POST /` - create background run; `POST /stream` - create + SSE stream; `POST /wait` - create + block; `POST /regenerate/prepare` - prepare clean input + checkpoint metadata for regenerating the latest assistant answer; `GET /` - list runs; `GET /{rid}` - run details; `POST /{rid}/cancel` - cancel; `GET /{rid}/join` - join SSE; `GET /{rid}/messages` - paginated messages `{data, has_more}`; `GET /{rid}/events` - full event stream; `GET /{rid}/workspace-changes` - workspace/output file change summary and optional diffs; `GET /../messages` - thread messages with feedback; `GET /../token-usage` - aggregate tokens | diff --git a/backend/app/gateway/routers/threads.py b/backend/app/gateway/routers/threads.py index b8ca4c0dd..7d0173c76 100644 --- a/backend/app/gateway/routers/threads.py +++ b/backend/app/gateway/routers/threads.py @@ -12,8 +12,11 @@ matching the LangGraph Platform wire format expected by the from __future__ import annotations +import copy import logging +import shutil import uuid +from pathlib import Path from typing import Any from fastapi import APIRouter, HTTPException, Request @@ -28,6 +31,7 @@ from deerflow.config.paths import Paths, get_paths from deerflow.runtime import serialize_channel_values_for_api from deerflow.runtime.goal import DEFAULT_MAX_GOAL_CONTINUATIONS, build_goal_state, ensure_thread_checkpoint, goal_thread_lock, read_thread_goal, write_thread_goal from deerflow.runtime.user_context import get_effective_user_id +from deerflow.utils.file_io import run_file_io from deerflow.utils.time import coerce_iso, now_iso logger = logging.getLogger(__name__) @@ -41,6 +45,9 @@ router = APIRouter(prefix="/api/threads", tags=["threads"]) # row-level invariant is still ``threads_meta.user_id`` populated from # the auth contextvar; this list closes the metadata-blob echo gap. _SERVER_RESERVED_METADATA_KEYS: frozenset[str] = frozenset({"owner_id", "user_id"}) +_SIDECAR_METADATA_KEY = "deerflow_sidecar" +_BRANCH_METADATA_KEY = "deerflow_branch" +_BRANCH_HISTORY_SCAN_LIMIT = 200 def _strip_reserved_metadata(metadata: dict[str, Any] | None) -> dict[str, Any]: @@ -50,6 +57,155 @@ def _strip_reserved_metadata(metadata: dict[str, Any] | None) -> dict[str, Any]: return {k: v for k, v in metadata.items() if k not in _SERVER_RESERVED_METADATA_KEYS} +def _message_id(message: Any) -> str | None: + if isinstance(message, dict): + raw = message.get("id") + else: + raw = getattr(message, "id", None) + return raw if isinstance(raw, str) and raw else None + + +def _message_type(message: Any) -> str | None: + if isinstance(message, dict): + raw = message.get("type") + else: + raw = getattr(message, "type", None) + return raw if isinstance(raw, str) and raw else None + + +def _message_additional_kwargs(message: Any) -> dict[str, Any]: + if isinstance(message, dict): + raw = message.get("additional_kwargs") + else: + raw = getattr(message, "additional_kwargs", None) + return raw if isinstance(raw, dict) else {} + + +def _is_branch_visible_message(message: Any) -> bool: + if _message_additional_kwargs(message).get("hide_from_ui") is True: + return False + return _message_type(message) in {"human", "ai"} + + +def _is_branch_assistant_message(message: Any) -> bool: + return _message_type(message) == "ai" + + +def _checkpoint_messages(checkpoint_tuple: Any) -> list[Any]: + checkpoint = getattr(checkpoint_tuple, "checkpoint", {}) or {} + channel_values = checkpoint.get("channel_values", {}) or {} + messages = channel_values.get("messages") or [] + return list(messages) if isinstance(messages, list) else [] + + +def _checkpoint_id(checkpoint_tuple: Any) -> str | None: + config = getattr(checkpoint_tuple, "config", {}) or {} + raw = config.get("configurable", {}).get("checkpoint_id") + return raw if isinstance(raw, str) and raw else None + + +def _matches_branch_target(messages: list[Any], target_message_ids: set[str]) -> bool: + if not target_message_ids: + return False + + index_by_id = {_message_id(message): index for index, message in enumerate(messages) if _message_id(message)} + if not target_message_ids.issubset(index_by_id.keys()): + return False + if any(not _is_branch_assistant_message(messages[index_by_id[message_id]]) for message_id in target_message_ids): + return False + + target_end_index = max(index_by_id[message_id] for message_id in target_message_ids) + return not any(_is_branch_visible_message(message) for message in messages[target_end_index + 1 :]) + + +async def _find_branch_checkpoint(checkpointer: Any, thread_id: str, target_message_ids: set[str]) -> Any: + config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}} + try: + async for checkpoint_tuple in checkpointer.alist(config, limit=_BRANCH_HISTORY_SCAN_LIMIT): + if _matches_branch_target(_checkpoint_messages(checkpoint_tuple), target_message_ids): + return checkpoint_tuple + except Exception: + logger.exception("Failed to scan branch checkpoint history for thread %s", sanitize_log_param(thread_id)) + raise HTTPException(status_code=500, detail="Failed to find branch checkpoint") + raise HTTPException(status_code=409, detail="This turn can no longer be branched from.") + + +async def _branch_targets_latest_turn(checkpointer: Any, thread_id: str, target_message_ids: set[str]) -> bool: + """Return True when the target turn is the final visible turn in the current state. + + ``alist`` yields newest-first; we take the newest checkpoint that actually holds + messages (thread creation writes an empty checkpoint that must be skipped) and + reuse ``_matches_branch_target`` to check the target turn is its tail. Used to + decide whether cloning the (uncheckpointed) workspace onto a branch is safe: only + a branch from the latest turn shares the current workspace timeline. On any lookup + failure we fail closed (treat as historical) so a branch from an older turn never + inherits a later timeline's workspace files. + """ + config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}} + try: + async for checkpoint_tuple in checkpointer.alist(config, limit=_BRANCH_HISTORY_SCAN_LIMIT): + messages = _checkpoint_messages(checkpoint_tuple) + if not messages: + continue + return _matches_branch_target(messages, target_message_ids) + except Exception: + logger.warning( + "Failed to resolve latest turn for thread %s; treating branch as historical", + sanitize_log_param(thread_id), + exc_info=True, + ) + return False + + +def _ignore_branch_user_data(directory: str, names: list[str]) -> set[str]: + ignored: set[str] = set() + base = Path(directory) + for name in names: + path = base / name + if name.startswith(".upload-") and name.endswith(".part"): + ignored.add(name) + elif path.is_symlink(): + ignored.add(name) + return ignored + + +def _copy_branch_user_data_sync(paths: Paths, source_thread_id: str, target_thread_id: str, *, user_id: str) -> str: + source = paths.sandbox_user_data_dir(source_thread_id, user_id=user_id) + target = paths.sandbox_user_data_dir(target_thread_id, user_id=user_id) + if not source.exists(): + return "not_found" + + shutil.copytree(source, target, ignore=_ignore_branch_user_data, dirs_exist_ok=True) + return "current_thread_best_effort" + + +async def _copy_branch_user_data(source_thread_id: str, target_thread_id: str) -> str: + paths = get_paths() + user_id = get_effective_user_id() + try: + return await run_file_io(_copy_branch_user_data_sync, paths, source_thread_id, target_thread_id, user_id=user_id) + except Exception: + logger.warning( + "Failed to copy user-data for branch %s -> %s", + sanitize_log_param(source_thread_id), + sanitize_log_param(target_thread_id), + exc_info=True, + ) + return "failed" + + +def _default_branch_display_name(source_title: Any, *, source_is_branch: bool = False) -> str | None: + if not isinstance(source_title, str): + return None + + display_name = source_title.strip() + if source_is_branch: + while display_name.lower().startswith("branch:"): + display_name = display_name[len("branch:") :].strip() + + return display_name or None + + # --------------------------------------------------------------------------- # Response / request models # --------------------------------------------------------------------------- @@ -181,6 +337,24 @@ class ThreadHistoryRequest(BaseModel): before: str | None = Field(default=None, description="Cursor for pagination") +class ThreadBranchRequest(BaseModel): + """Request body for creating a branch from a completed assistant turn.""" + + message_id: str = Field(..., min_length=1, description="Target assistant message ID to branch from") + message_ids: list[str] = Field(default_factory=list, description="All assistant message IDs in the target turn") + title: str | None = Field(default=None, max_length=256, description="Optional title for the branched thread") + + +class ThreadBranchResponse(BaseModel): + """Response model for a thread branch.""" + + thread_id: str + parent_thread_id: str + parent_checkpoint_id: str + branched_from_message_id: str + workspace_clone_mode: str + + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -367,6 +541,98 @@ async def create_thread(body: ThreadCreateRequest, request: Request) -> ThreadRe ) +@router.post("/{thread_id}/branches", response_model=ThreadBranchResponse) +@require_permission("threads", "write", owner_check=True, require_existing=True) +async def branch_thread(thread_id: str, body: ThreadBranchRequest, request: Request) -> ThreadBranchResponse: + """Create a new main-thread branch from a completed assistant turn.""" + from app.gateway.deps import get_thread_store + + checkpointer = get_checkpointer(request) + thread_store = get_thread_store(request) + + source_record = await thread_store.get(thread_id) + if source_record is None: + raise HTTPException(status_code=404, detail=f"Thread {thread_id} not found") + + source_metadata = source_record.get("metadata") or {} + if source_metadata.get(_SIDECAR_METADATA_KEY) is True: + raise HTTPException(status_code=409, detail="Branching is only available in the main conversation.") + + target_message_ids = {body.message_id, *body.message_ids} + checkpoint_tuple = await _find_branch_checkpoint(checkpointer, thread_id, target_message_ids) + parent_checkpoint_id = _checkpoint_id(checkpoint_tuple) + if not parent_checkpoint_id: + raise HTTPException(status_code=409, detail="This turn can no longer be branched from.") + + # Workspace files are not checkpointed, so they only reflect the *current* thread + # state. Cloning them onto a branch from an older turn would leak files created + # after that turn (message history rolls back, workspace would not). Restrict the + # best-effort clone to branches taken from the latest turn so history and workspace + # stay consistent. + branch_from_latest_turn = await _branch_targets_latest_turn(checkpointer, thread_id, target_message_ids) + + new_thread_id = str(uuid.uuid4()) + now = now_iso() + branch_metadata = { + _BRANCH_METADATA_KEY: True, + "branch_parent_thread_id": thread_id, + "branch_parent_checkpoint_id": parent_checkpoint_id, + "branch_parent_message_id": body.message_id, + "branch_created_at": now, + } + + display_name = body.title or _default_branch_display_name( + source_record.get("display_name"), + source_is_branch=source_metadata.get(_BRANCH_METADATA_KEY) is True, + ) + thread_owner_user_id = get_trusted_internal_owner_user_id(request) + thread_owner_kwargs = {"user_id": thread_owner_user_id} if thread_owner_user_id else {} + + checkpoint = copy.deepcopy(getattr(checkpoint_tuple, "checkpoint", {}) or {}) + metadata = copy.deepcopy(getattr(checkpoint_tuple, "metadata", {}) or {}) + checkpoint["id"] = str(uuid6()) + metadata.update( + { + "source": "branch", + "updated_at": now, + "created_at": now, + **branch_metadata, + } + ) + + write_config = {"configurable": {"thread_id": new_thread_id, "checkpoint_ns": ""}} + new_versions = dict(checkpoint.get("channel_versions", {}) or {}) + try: + await checkpointer.aput(write_config, checkpoint, metadata, new_versions) + except Exception: + logger.exception("Failed to write branch checkpoint for thread %s", sanitize_log_param(new_thread_id)) + raise HTTPException(status_code=500, detail="Failed to create branch") from None + + try: + await thread_store.create( + new_thread_id, + assistant_id=source_record.get("assistant_id"), + display_name=display_name, + metadata=branch_metadata, + **thread_owner_kwargs, + ) + except Exception: + logger.exception("Failed to write branch thread_meta for %s", sanitize_log_param(new_thread_id)) + raise HTTPException(status_code=500, detail="Failed to create branch") from None + + if branch_from_latest_turn: + workspace_clone_mode = await _copy_branch_user_data(thread_id, new_thread_id) + else: + workspace_clone_mode = "skipped_historical_turn" + return ThreadBranchResponse( + thread_id=new_thread_id, + parent_thread_id=thread_id, + parent_checkpoint_id=parent_checkpoint_id, + branched_from_message_id=body.message_id, + workspace_clone_mode=workspace_clone_mode, + ) + + @router.post("/search", response_model=list[ThreadResponse]) async def search_threads(body: ThreadSearchRequest, request: Request) -> list[ThreadResponse]: """Search and list threads. diff --git a/backend/tests/test_threads_router.py b/backend/tests/test_threads_router.py index b035c9bf7..d4e4d37a9 100644 --- a/backend/tests/test_threads_router.py +++ b/backend/tests/test_threads_router.py @@ -1,3 +1,4 @@ +import asyncio import re from types import SimpleNamespace from unittest.mock import patch @@ -6,6 +7,8 @@ import pytest from _router_auth_helpers import make_authed_test_app from fastapi import FastAPI, HTTPException from fastapi.testclient import TestClient +from langchain_core.messages import AIMessage, HumanMessage +from langgraph.checkpoint.base import empty_checkpoint from langgraph.checkpoint.memory import InMemorySaver from langgraph.store.memory import InMemoryStore @@ -65,6 +68,32 @@ def _build_thread_app() -> tuple[FastAPI, InMemoryStore, InMemorySaver]: return app, store, checkpointer +async def _write_checkpoint( + checkpointer: InMemorySaver, + thread_id: str, + checkpoint_id: str, + messages: list[object], + *, + step: int, +) -> dict: + checkpoint = empty_checkpoint() + checkpoint["id"] = checkpoint_id + checkpoint["channel_values"] = {"messages": messages} + checkpoint["channel_versions"] = {"messages": step} + return await checkpointer.aput( + {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}, + checkpoint, + { + "step": step, + "source": "loop", + "writes": {"test": {"messages": messages}}, + "parents": {}, + "created_at": f"2026-07-05T00:00:0{step}+00:00", + }, + {"messages": step}, + ) + + def test_delete_thread_data_removes_thread_directory(tmp_path): paths = Paths(tmp_path) thread_dir = paths.thread_dir("thread-cleanup") @@ -539,6 +568,216 @@ def test_get_thread_history_returns_iso_for_legacy_checkpoint_metadata() -> None assert _ISO_TIMESTAMP_RE.match(entry["created_at"]), entry +# ── branch threads from completed assistant turns ───────────────────────────── + + +def test_branch_thread_from_older_assistant_turn_creates_truncated_thread() -> None: + app, store, checkpointer = _build_thread_app() + source_thread_id = "source-thread" + + human_1 = HumanMessage(id="human-1", content="First question") + ai_1 = AIMessage(id="ai-1", content="First answer") + human_2 = HumanMessage(id="human-2", content="Second question") + ai_2 = AIMessage(id="ai-2", content="Second answer") + human_3 = HumanMessage(id="human-3", content="Third question") + ai_3 = AIMessage(id="ai-3", content="Third answer") + + async def _seed() -> None: + await _write_checkpoint(checkpointer, source_thread_id, "0001", [human_1, ai_1], step=1) + await _write_checkpoint(checkpointer, source_thread_id, "0002", [human_1, ai_1, human_2, ai_2], step=2) + await _write_checkpoint(checkpointer, source_thread_id, "0003", [human_1, ai_1, human_2, ai_2, human_3, ai_3], step=3) + + asyncio.run(_seed()) + + with TestClient(app) as client: + created = client.post("/api/threads", json={"thread_id": source_thread_id, "metadata": {}, "assistant_id": "agent"}) + assert created.status_code == 200, created.text + asyncio.run( + store.aput( + THREADS_NS, + source_thread_id, + { + "thread_id": source_thread_id, + "assistant_id": "agent", + "user_id": None, + "status": "idle", + "created_at": "2026-07-05T00:00:00Z", + "updated_at": "2026-07-05T00:00:00Z", + "display_name": "Original chat", + "metadata": {}, + }, + ) + ) + + response = client.post( + f"/api/threads/{source_thread_id}/branches", + json={"message_id": "ai-2", "message_ids": ["ai-2"]}, + ) + assert response.status_code == 200, response.text + body = response.json() + new_thread_id = body["thread_id"] + state_response = client.get(f"/api/threads/{new_thread_id}/state") + search_response = client.post("/api/threads/search", json={"limit": 10}) + + assert body["parent_thread_id"] == source_thread_id + assert body["parent_checkpoint_id"] == "0002" + assert body["branched_from_message_id"] == "ai-2" + assert body["workspace_clone_mode"] == "skipped_historical_turn" + + assert state_response.status_code == 200, state_response.text + messages = state_response.json()["values"]["messages"] + assert [message["id"] for message in messages] == ["human-1", "ai-1", "human-2", "ai-2"] + assert "Third answer" not in [message.get("content") for message in messages] + assert search_response.status_code == 200, search_response.text + branch_entry = next(item for item in search_response.json() if item["thread_id"] == new_thread_id) + assert branch_entry["values"]["title"] == "Original chat" + + +def test_branch_display_name_strips_legacy_branch_prefix_only_for_branch_sources() -> None: + assert threads._default_branch_display_name("Original chat") == "Original chat" + assert threads._default_branch_display_name("Branch: Original chat") == "Branch: Original chat" + assert threads._default_branch_display_name("Branch: Branch: Original chat", source_is_branch=True) == "Original chat" + + +def test_branch_thread_rejects_sidecar_threads() -> None: + app, _store, _checkpointer = _build_thread_app() + + with TestClient(app) as client: + created = client.post( + "/api/threads", + json={"thread_id": "sidecar-thread", "metadata": {"deerflow_sidecar": True}}, + ) + assert created.status_code == 200, created.text + + response = client.post( + "/api/threads/sidecar-thread/branches", + json={"message_id": "ai-1", "message_ids": ["ai-1"]}, + ) + + assert response.status_code == 409 + assert "main conversation" in response.json()["detail"] + + +def test_branch_thread_rejects_non_assistant_targets() -> None: + app, _store, checkpointer = _build_thread_app() + source_thread_id = "source-human-target" + human = HumanMessage(id="human-1", content="Question") + ai = AIMessage(id="ai-1", content="Answer") + + async def _seed() -> None: + await _write_checkpoint(checkpointer, source_thread_id, "0001", [human, ai], step=1) + + asyncio.run(_seed()) + + with TestClient(app) as client: + created = client.post("/api/threads", json={"thread_id": source_thread_id, "metadata": {}}) + assert created.status_code == 200, created.text + + response = client.post( + f"/api/threads/{source_thread_id}/branches", + json={"message_id": "human-1", "message_ids": ["human-1"]}, + ) + + assert response.status_code == 409 + assert "can no longer be branched" in response.json()["detail"] + + +def test_branch_thread_best_effort_copies_current_workspace(tmp_path) -> None: + paths = Paths(tmp_path) + app, _store, checkpointer = _build_thread_app() + source_thread_id = "source-with-files" + user_id = "branch-user" + + source_user_data = paths.sandbox_user_data_dir(source_thread_id, user_id=user_id) + source_outputs = paths.sandbox_outputs_dir(source_thread_id, user_id=user_id) + source_uploads = paths.sandbox_uploads_dir(source_thread_id, user_id=user_id) + source_outputs.mkdir(parents=True, exist_ok=True) + source_uploads.mkdir(parents=True, exist_ok=True) + (source_outputs / "result.txt").write_text("answer", encoding="utf-8") + (source_uploads / ".upload-stale.part").write_text("partial", encoding="utf-8") + + human = HumanMessage(id="human-file", content="Make a file") + ai = AIMessage(id="ai-file", content="Done") + + async def _seed() -> None: + await _write_checkpoint(checkpointer, source_thread_id, "0001", [human, ai], step=1) + + asyncio.run(_seed()) + + with ( + patch("app.gateway.routers.threads.get_paths", return_value=paths), + patch("app.gateway.routers.threads.get_effective_user_id", return_value=user_id), + TestClient(app) as client, + ): + created = client.post("/api/threads", json={"thread_id": source_thread_id, "metadata": {}}) + assert created.status_code == 200, created.text + + response = client.post( + f"/api/threads/{source_thread_id}/branches", + json={"message_id": "ai-file", "message_ids": ["ai-file"]}, + ) + + assert response.status_code == 200, response.text + body = response.json() + assert body["workspace_clone_mode"] == "current_thread_best_effort" + + target_user_data = paths.sandbox_user_data_dir(body["thread_id"], user_id=user_id) + assert target_user_data.exists() + assert (target_user_data / "outputs" / "result.txt").read_text(encoding="utf-8") == "answer" + assert not (target_user_data / "uploads" / ".upload-stale.part").exists() + assert source_user_data.exists() + + +def test_branch_thread_from_historical_turn_skips_workspace_clone(tmp_path) -> None: + """Branching from a non-latest turn must not clone the current workspace. + + Workspace files are not checkpointed, so cloning them onto a branch rooted at + an older turn would leak files created after that turn (regression for the + historical-turn workspace-leak review on PR #3950). + """ + paths = Paths(tmp_path) + app, _store, checkpointer = _build_thread_app() + source_thread_id = "source-historical" + user_id = "branch-user" + + source_outputs = paths.sandbox_outputs_dir(source_thread_id, user_id=user_id) + source_outputs.mkdir(parents=True, exist_ok=True) + # ``future.txt`` only exists in the current (latest) workspace timeline. + (source_outputs / "future.txt").write_text("future", encoding="utf-8") + + human_1 = HumanMessage(id="human-1", content="First question") + ai_1 = AIMessage(id="ai-1", content="First answer") + human_2 = HumanMessage(id="human-2", content="Second question") + ai_2 = AIMessage(id="ai-2", content="Second answer") + + async def _seed() -> None: + await _write_checkpoint(checkpointer, source_thread_id, "0001", [human_1, ai_1], step=1) + await _write_checkpoint(checkpointer, source_thread_id, "0002", [human_1, ai_1, human_2, ai_2], step=2) + + asyncio.run(_seed()) + + with ( + patch("app.gateway.routers.threads.get_paths", return_value=paths), + patch("app.gateway.routers.threads.get_effective_user_id", return_value=user_id), + TestClient(app) as client, + ): + created = client.post("/api/threads", json={"thread_id": source_thread_id, "metadata": {}}) + assert created.status_code == 200, created.text + + response = client.post( + f"/api/threads/{source_thread_id}/branches", + json={"message_id": "ai-1", "message_ids": ["ai-1"]}, + ) + + assert response.status_code == 200, response.text + body = response.json() + assert body["parent_checkpoint_id"] == "0001" + assert body["workspace_clone_mode"] == "skipped_historical_turn" + + target_user_data = paths.sandbox_user_data_dir(body["thread_id"], user_id=user_id) + assert not target_user_data.exists() + + # ── Metadata filter validation at API boundary ──────────────────────────────── diff --git a/frontend/AGENTS.md b/frontend/AGENTS.md index b46cb27eb..ad19d2b76 100644 --- a/frontend/AGENTS.md +++ b/frontend/AGENTS.md @@ -82,6 +82,7 @@ The frontend is a stateful chat application. Users create **threads** (conversat ### Interaction Ownership - `src/app/workspace/chats/[thread_id]/page.tsx` owns composer busy-state wiring. +- `src/app/workspace/chats/[thread_id]/page.tsx` owns branch-from-turn submission and navigation; sidecar `MessageList` instances do not receive the branch action. - `src/app/workspace/chats/[thread_id]/page.tsx` and `src/app/workspace/agents/[agent_name]/chats/[thread_id]/page.tsx` own active-goal display state for their composer overlays. - `src/core/threads/hooks.ts` owns pre-submit upload state and thread submission. diff --git a/frontend/src/app/workspace/chats/[thread_id]/page.tsx b/frontend/src/app/workspace/chats/[thread_id]/page.tsx index 909d9c57a..4049f99d9 100644 --- a/frontend/src/app/workspace/chats/[thread_id]/page.tsx +++ b/frontend/src/app/workspace/chats/[thread_id]/page.tsx @@ -2,6 +2,7 @@ import { useRouter } from "next/navigation"; import { useCallback, useEffect, useRef, useState } from "react"; +import { toast } from "sonner"; import { type PromptInputMessage } from "@/components/ai-elements/prompt-input"; import { SidebarTrigger } from "@/components/ui/sidebar"; @@ -37,6 +38,7 @@ import { useModels } from "@/core/models/hooks"; import { useNotification } from "@/core/notification/hooks"; import { useLocalSettings, useThreadSettings } from "@/core/settings"; import { + useBranchThread, useThreadMetadata, useThreadStream, useThreadTokenUsage, @@ -68,6 +70,7 @@ export default function ChatPage() { enabled: !isNewThread && !isMock, isMock, }); + const branchThread = useBranchThread(); const backendTokenUsage = threadTokenUsageToTokenUsage(threadTokenUsage.data); const mountedRef = useRef(false); useSpecificChatMode(); @@ -175,6 +178,32 @@ export default function ChatPage() { regenerateMessage(threadId, messageId, supersededMessageIds), [regenerateMessage, threadId], ); + const handleBranchTurn = useCallback( + async (messageId: string, messageIds: string[]) => { + if ( + isNewThread || + isMock || + env.NEXT_PUBLIC_STATIC_WEBSITE_ONLY === "true" + ) { + return; + } + + try { + const response = await branchThread.mutateAsync({ + threadId, + messageId, + messageIds, + }); + toast.success(t.conversation.branchCreated); + router.push(`/workspace/chats/${response.thread_id}`); + } catch (error) { + toast.error( + error instanceof Error ? error.message : t.conversation.branchFailed, + ); + } + }, + [branchThread, isMock, isNewThread, router, t, threadId], + ); const tokenUsageInlineMode = tokenUsageEnabled ? localSettings.tokenUsage.inlineMode @@ -246,6 +275,15 @@ export default function ChatPage() { !thread.isLoading } onRegenerateMessage={handleRegenerate} + canBranch={ + !isNewThread && + !isMock && + env.NEXT_PUBLIC_STATIC_WEBSITE_ONLY !== "true" && + !isUploading && + !thread.isLoading && + !branchThread.isPending + } + onBranchTurn={handleBranchTurn} />
void | Promise; + onBranchTurn?: ( + messageId: string, + messageIds: string[], + ) => void | Promise; canRegenerate?: boolean; + canBranch?: boolean; enableSidecarActions?: boolean; initialScroll?: ConversationProps["initial"]; resizeScroll?: ConversationProps["resize"]; @@ -248,6 +256,9 @@ export function MessageList({ const [regeneratingMessageId, setRegeneratingMessageId] = useState< string | null >(null); + const [branchingMessageId, setBranchingMessageId] = useState( + null, + ); const hasActiveAssistantText = useMemo(() => { let lastHumanIndex = -1; for (let i = groupedMessages.length - 1; i >= 0; i--) { @@ -427,22 +438,52 @@ export function MessageList({ enableRegenerateForTurn: boolean, ) => { const clipboardData = getAssistantTurnCopyData(messages, { isStreaming }); - const regenerateTarget = [...messages] + const actionTarget = [...messages] .reverse() .find((message) => message.type === "ai" && message.id); - const supersededMessageIds = messages + const assistantMessageIds = messages .filter((message) => message.type === "ai" && message.id) .map((message) => message.id) .filter((id): id is string => typeof id === "string"); - if (!clipboardData && !regenerateTarget) { + if (!clipboardData && !actionTarget) { return null; } return (
{clipboardData && } + {!isStreaming && actionTarget?.id && onBranchTurn && ( + + + + )} {enableRegenerateForTurn && - regenerateTarget?.id && + actionTarget?.id && onRegenerateMessage && (