From e54b6babd681b216125155d018d256f6d73fd1b0 Mon Sep 17 00:00:00 2001 From: Erildo Date: Mon, 15 Jun 2026 15:57:39 +0200 Subject: [PATCH] =?UTF-8?q?feat:=20implement=20thread=20forking=20function?= =?UTF-8?q?ality=20with=20associated=20database=E2=80=A6=20(#5810)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: implement thread forking functionality with associated database updates and UI components * fix(studio/chat): register fork-count listener even when thread unsaved Co-Authored-By: Claude Opus 4.7 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Polish thread fork action menu * fix-studio-fork-project-test-order * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: Claude Opus 4.7 Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: imagineer99 Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- studio/backend/routes/chat_history.py | 83 +++++++++ studio/backend/storage/studio_db.py | 146 ++++++++++++++- .../backend/tests/test_chat_history_routes.py | 174 ++++++++++++++++++ .../tests/test_chat_history_storage.py | 132 +++++++++++++ .../src/components/assistant-ui/thread.tsx | 128 ++++++++++++- .../src/features/chat/api/chat-api.ts | 39 ++++ .../chat/hooks/use-chat-sidebar-items.ts | 2 + .../src/features/chat/thread-sidebar.tsx | 8 + studio/frontend/src/features/chat/types.ts | 8 + 9 files changed, 715 insertions(+), 5 deletions(-) diff --git a/studio/backend/routes/chat_history.py b/studio/backend/routes/chat_history.py index 2ea572e30e..1243b284b4 100644 --- a/studio/backend/routes/chat_history.py +++ b/studio/backend/routes/chat_history.py @@ -18,9 +18,11 @@ from storage.studio_db import ( CorruptSettingsError, clear_chat_history, count_chat_threads, + count_forks_for_message, delete_chat_threads, delete_chat_project, ensure_chat_project_workspace, + fork_chat_thread, get_chat_project, get_chat_thread, get_chat_message, @@ -56,6 +58,8 @@ class ChatThread(BaseModel): createdAt: int openaiCodeExecContainerId: Optional[str] = None anthropicCodeExecContainerId: Optional[str] = None + forkedFromThreadId: Optional[str] = None + forkedFromMessageId: Optional[str] = None class ChatThreadPatch(BaseModel): @@ -518,6 +522,85 @@ async def put_settings( ) from exc +class ChatForkRequest(BaseModel): + messageId: str + newThreadId: str + createdAt: int + + +class ChatForkResponse(BaseModel): + thread: ChatThread + messages: list[ChatMessage] + containerSnapshotWarning: Optional[str] = None + + +class ChatForkCountResponse(BaseModel): + count: int + + +@router.post("/threads/{thread_id}/fork", response_model = ChatForkResponse) +async def fork_thread( + thread_id: str, + payload: ChatForkRequest, + current_subject: str = Depends(get_current_subject), +): + """Fork a thread at `messageId` -- creates a new thread with + ancestor msgs [root..messageId] copied with fresh ids. Both + code-exec container ids reset on the fork. OpenAI snapshot is a + best-effort enhancement; failure surfaces as + `containerSnapshotWarning` and the fork still succeeds with a + clean sandbox. + """ + import uuid + + source = get_chat_thread(thread_id) + if source is None: + raise HTTPException(status_code = 404, detail = f"Thread {thread_id} not found") + if get_chat_message(thread_id, payload.messageId) is None: + raise HTTPException( + status_code = 404, + detail = f"Message {payload.messageId} not found in thread {thread_id}", + ) + base_title = source.get("title") or "New Chat" + new_title = f"fork · {base_title}" + forked = fork_chat_thread( + source_thread_id = thread_id, + branch_message_id = payload.messageId, + new_thread_id = payload.newThreadId, + new_title = new_title, + created_at = payload.createdAt, + id_factory = lambda: str(uuid.uuid4()), + ) + if forked is None: + raise HTTPException(status_code = 500, detail = "Fork failed") + messages = list_chat_messages(payload.newThreadId) + # Best-effort OpenAI container snapshot. Stub: a follow-up patch can + # call /v1/containers list+download / create+upload here and patch + # the new openaiCodeExecContainerId. For v1 we always start clean + # and surface the same warning regardless of provider so the UI can + # show a consistent "sandbox starts fresh" toast. + warning: Optional[str] = None + if source.get("openaiCodeExecContainerId") or source.get("anthropicCodeExecContainerId"): + warning = "Sandbox starts fresh in fork; files from parent are not carried over." + return ChatForkResponse( + thread = ChatThread(**forked), + messages = [ChatMessage(**m) for m in messages], + containerSnapshotWarning = warning, + ) + + +@router.get( + "/threads/{thread_id}/messages/{message_id}/forks", + response_model = ChatForkCountResponse, +) +async def get_fork_count( + thread_id: str, + message_id: str, + current_subject: str = Depends(get_current_subject), +): + return ChatForkCountResponse(count = count_forks_for_message(thread_id, message_id)) + + @router.get("/export", response_model = ChatExportResponse) async def export_history(current_subject: str = Depends(get_current_subject)): from datetime import datetime, timezone diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index 85cfacbc27..7421b42b2f 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -216,6 +216,8 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: created_at INTEGER NOT NULL, openai_code_exec_container_id TEXT, anthropic_code_exec_container_id TEXT, + forked_from_thread_id TEXT, + forked_from_message_id TEXT, FOREIGN KEY(project_id) REFERENCES chat_projects(id) ON DELETE CASCADE ) """ @@ -229,6 +231,10 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: conn.execute("ALTER TABLE chat_threads ADD COLUMN openai_code_exec_container_id TEXT") if "anthropic_code_exec_container_id" not in chat_thread_cols: conn.execute("ALTER TABLE chat_threads ADD COLUMN anthropic_code_exec_container_id TEXT") + if "forked_from_thread_id" not in chat_thread_cols: + conn.execute("ALTER TABLE chat_threads ADD COLUMN forked_from_thread_id TEXT") + if "forked_from_message_id" not in chat_thread_cols: + conn.execute("ALTER TABLE chat_threads ADD COLUMN forked_from_message_id TEXT") conn.execute( """ CREATE TABLE IF NOT EXISTS chat_messages ( @@ -956,6 +962,8 @@ def _chat_thread_from_row(row: sqlite3.Row) -> dict: "createdAt": data["created_at"], "openaiCodeExecContainerId": data.get("openai_code_exec_container_id"), "anthropicCodeExecContainerId": data.get("anthropic_code_exec_container_id"), + "forkedFromThreadId": data.get("forked_from_thread_id"), + "forkedFromMessageId": data.get("forked_from_message_id"), } @@ -999,8 +1007,8 @@ def upsert_chat_thread(thread: dict) -> dict: conn.execute( """ INSERT INTO chat_threads - (id, title, model_type, model_id, pair_id, project_id, archived, created_at, openai_code_exec_container_id, anthropic_code_exec_container_id) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + (id, title, model_type, model_id, pair_id, project_id, archived, created_at, openai_code_exec_container_id, anthropic_code_exec_container_id, forked_from_thread_id, forked_from_message_id) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(id) DO UPDATE SET title = excluded.title, model_type = excluded.model_type, @@ -1010,7 +1018,9 @@ def upsert_chat_thread(thread: dict) -> dict: archived = excluded.archived, created_at = excluded.created_at, openai_code_exec_container_id = excluded.openai_code_exec_container_id, - anthropic_code_exec_container_id = excluded.anthropic_code_exec_container_id + anthropic_code_exec_container_id = excluded.anthropic_code_exec_container_id, + forked_from_thread_id = excluded.forked_from_thread_id, + forked_from_message_id = excluded.forked_from_message_id """, ( thread["id"], @@ -1023,6 +1033,8 @@ def upsert_chat_thread(thread: dict) -> dict: int(thread["createdAt"]), thread.get("openaiCodeExecContainerId"), thread.get("anthropicCodeExecContainerId"), + thread.get("forkedFromThreadId"), + thread.get("forkedFromMessageId"), ), ) conn.commit() @@ -1048,6 +1060,14 @@ def update_chat_thread(id: str, patch: dict) -> Optional[dict]: "anthropic_code_exec_container_id", patch.get("anthropicCodeExecContainerId"), ), + "forkedFromThreadId": ( + "forked_from_thread_id", + patch.get("forkedFromThreadId"), + ), + "forkedFromMessageId": ( + "forked_from_message_id", + patch.get("forkedFromMessageId"), + ), } assignments = [] values = [] @@ -1445,6 +1465,126 @@ def sync_chat_messages( conn.close() +def fork_chat_thread( + source_thread_id: str, + branch_message_id: str, + new_thread_id: str, + new_title: str, + created_at: int, + id_factory, +) -> Optional[dict]: + """Atomically clone thread + ancestor msgs `[root..branch_message_id]` + into a new thread. Returns the new thread dict (with messages copied) + or None if source missing. + + Reset both code-exec container ids -- per-provider snapshot is handled + by the route layer (best-effort, OpenAI only). + + `id_factory()` produces fresh message uuids; injected for testability. + """ + conn = get_connection() + try: + conn.execute("BEGIN IMMEDIATE") + src = conn.execute( + "SELECT * FROM chat_threads WHERE id = ?", (source_thread_id,) + ).fetchone() + if src is None: + conn.rollback() + return None + # Verify branch msg belongs to source thread. + branch_row = conn.execute( + "SELECT * FROM chat_messages WHERE thread_id = ? AND id = ?", + (source_thread_id, branch_message_id), + ).fetchone() + if branch_row is None: + conn.rollback() + return None + # Walk ancestry from branch msg back to root via parent_id chain. + ancestry: list[sqlite3.Row] = [] + cursor_row = branch_row + seen: set[str] = set() + while cursor_row is not None and cursor_row["id"] not in seen: + ancestry.append(cursor_row) + seen.add(cursor_row["id"]) + parent = cursor_row["parent_id"] + if not parent: + break + cursor_row = conn.execute( + "SELECT * FROM chat_messages WHERE thread_id = ? AND id = ?", + (source_thread_id, parent), + ).fetchone() + ancestry.reverse() # root .. branch msg + # Map old msg id -> new msg id for parent_id rewriting. + id_map: dict[str, str] = {row["id"]: id_factory() for row in ancestry} + src_dict = dict(src) + conn.execute( + """ + INSERT INTO chat_threads + (id, title, model_type, model_id, pair_id, project_id, archived, created_at, + openai_code_exec_container_id, anthropic_code_exec_container_id, + forked_from_thread_id, forked_from_message_id) + VALUES (?, ?, ?, ?, ?, ?, 0, ?, NULL, NULL, ?, ?) + """, + ( + new_thread_id, + new_title, + src_dict["model_type"], + src_dict.get("model_id") or "", + None, # pairId: forks always standalone (compare-mode disabled v1) + src_dict.get("project_id"), + int(created_at), + source_thread_id, + branch_message_id, + ), + ) + conn.executemany( + """ + INSERT INTO chat_messages + (id, thread_id, parent_id, role, content_json, attachments_json, + metadata_json, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + """, + [ + ( + id_map[row["id"]], + new_thread_id, + id_map.get(row["parent_id"]) if row["parent_id"] else None, + row["role"], + row["content_json"], + row["attachments_json"], + row["metadata_json"], + int(row["created_at"]), + ) + for row in ancestry + ], + ) + conn.commit() + thread_row = conn.execute( + "SELECT * FROM chat_threads WHERE id = ?", (new_thread_id,) + ).fetchone() + return _chat_thread_from_row(thread_row) if thread_row is not None else None + except Exception: + conn.rollback() + raise + finally: + conn.close() + + +def count_forks_for_message(thread_id: str, message_id: str) -> int: + conn = get_connection() + try: + row = conn.execute( + """ + SELECT COUNT(*) FROM chat_threads + WHERE forked_from_thread_id = ? AND forked_from_message_id = ? + """, + (thread_id, message_id), + ).fetchone() + return int(row[0]) if row is not None else 0 + finally: + conn.close() + + def list_chat_messages(thread_id: str) -> list[dict]: conn = get_connection() try: diff --git a/studio/backend/tests/test_chat_history_routes.py b/studio/backend/tests/test_chat_history_routes.py index e3aeb4eb22..2a6ebe244f 100644 --- a/studio/backend/tests/test_chat_history_routes.py +++ b/studio/backend/tests/test_chat_history_routes.py @@ -169,3 +169,177 @@ def test_record_import_ledger_rejects_oversize_payload(): chat_history.ChatImportLedgerRecordRequest( threadIds = [f"id-{i}" for i in range(10_001)], ) + + +# --------------------------------------------------------------------------- +# /api/chat/threads/{id}/fork +# --------------------------------------------------------------------------- + + +def test_fork_thread_404_when_source_missing(monkeypatch): + monkeypatch.setattr(chat_history, "get_chat_thread", lambda _id: None) + with pytest.raises(HTTPException) as exc: + asyncio.run( + chat_history.fork_thread( + thread_id = "missing", + payload = chat_history.ChatForkRequest( + messageId = "m1", + newThreadId = "new", + createdAt = 1, + ), + current_subject = "test-user", + ) + ) + assert exc.value.status_code == 404 + + +def test_fork_thread_404_when_branch_message_missing(monkeypatch): + monkeypatch.setattr(chat_history, "get_chat_thread", lambda _id: {"id": _id, "title": "T"}) + monkeypatch.setattr(chat_history, "get_chat_message", lambda _t, _m: None) + with pytest.raises(HTTPException) as exc: + asyncio.run( + chat_history.fork_thread( + thread_id = "src", + payload = chat_history.ChatForkRequest( + messageId = "missing", + newThreadId = "new", + createdAt = 1, + ), + current_subject = "test-user", + ) + ) + assert exc.value.status_code == 404 + + +def test_fork_thread_happy_path(monkeypatch): + source = { + "id": "src", + "title": "Original", + "modelType": "base", + "modelId": "m", + "pairId": None, + "archived": False, + "createdAt": 1, + "openaiCodeExecContainerId": None, + "anthropicCodeExecContainerId": None, + "forkedFromThreadId": None, + "forkedFromMessageId": None, + } + forked = { + **source, + "id": "new", + "title": "fork · Original", + "createdAt": 2, + "forkedFromThreadId": "src", + "forkedFromMessageId": "m1", + } + monkeypatch.setattr(chat_history, "get_chat_thread", lambda _id: source) + monkeypatch.setattr( + chat_history, + "get_chat_message", + lambda _t, _m: { + "id": _m, + "threadId": _t, + "role": "user", + "content": [], + "createdAt": 1, + }, + ) + monkeypatch.setattr(chat_history, "fork_chat_thread", lambda **_: forked) + monkeypatch.setattr( + chat_history, + "list_chat_messages", + lambda _id: [ + { + "id": "n1", + "threadId": "new", + "parentId": None, + "role": "user", + "content": [], + "createdAt": 1, + } + ], + ) + response = asyncio.run( + chat_history.fork_thread( + thread_id = "src", + payload = chat_history.ChatForkRequest( + messageId = "m1", + newThreadId = "new", + createdAt = 2, + ), + current_subject = "test-user", + ) + ) + assert response.thread.id == "new" + assert response.thread.title == "fork · Original" + assert response.thread.forkedFromThreadId == "src" + assert response.thread.forkedFromMessageId == "m1" + assert len(response.messages) == 1 + assert response.containerSnapshotWarning is None + + +def test_fork_thread_warns_when_parent_had_container(monkeypatch): + source = { + "id": "src", + "title": "T", + "modelType": "base", + "modelId": "", + "pairId": None, + "archived": False, + "createdAt": 1, + "openaiCodeExecContainerId": "cnt_123", + "anthropicCodeExecContainerId": None, + "forkedFromThreadId": None, + "forkedFromMessageId": None, + } + monkeypatch.setattr(chat_history, "get_chat_thread", lambda _id: source) + monkeypatch.setattr( + chat_history, + "get_chat_message", + lambda _t, _m: { + "id": _m, + "threadId": _t, + "role": "user", + "content": [], + "createdAt": 1, + }, + ) + monkeypatch.setattr( + chat_history, + "fork_chat_thread", + lambda **_: { + **source, + "id": "new", + "title": "fork · T", + "forkedFromThreadId": "src", + "forkedFromMessageId": "m1", + "openaiCodeExecContainerId": None, + }, + ) + monkeypatch.setattr(chat_history, "list_chat_messages", lambda _id: []) + response = asyncio.run( + chat_history.fork_thread( + thread_id = "src", + payload = chat_history.ChatForkRequest( + messageId = "m1", + newThreadId = "new", + createdAt = 2, + ), + current_subject = "test-user", + ) + ) + assert response.containerSnapshotWarning is not None + assert "fresh" in response.containerSnapshotWarning.lower() + + +def test_get_fork_count(monkeypatch): + monkeypatch.setattr(chat_history, "count_forks_for_message", lambda _t, _m: 3) + response = asyncio.run( + chat_history.get_fork_count( + thread_id = "t", + message_id = "m", + current_subject = "test-user", + ) + ) + assert response.count == 3 diff --git a/studio/backend/tests/test_chat_history_storage.py b/studio/backend/tests/test_chat_history_storage.py index bc74ec172a..aa19df15fe 100644 --- a/studio/backend/tests/test_chat_history_storage.py +++ b/studio/backend/tests/test_chat_history_storage.py @@ -378,3 +378,135 @@ def test_legacy_imports_ignores_empty(tmp_path, monkeypatch): assert studio_db.upsert_chat_legacy_imports([]) == (0, 0) assert studio_db.upsert_chat_legacy_imports(["", None]) == (0, 0) # type: ignore[list-item] assert studio_db.list_chat_legacy_imports() == [] + + +# --------------------------------------------------------------------------- +# fork_chat_thread +# --------------------------------------------------------------------------- + + +def _msg(mid: str, parent: str | None, t: int) -> dict: + return { + "id": mid, + "threadId": "src", + "parentId": parent, + "role": "user", + "content": [{"type": "text", "text": mid}], + "createdAt": t, + } + + +def test_fork_chat_thread_copies_ancestry_with_fresh_ids(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_thread( + {**_thread("src"), "title": "Original", "openaiCodeExecContainerId": "cnt-x"} + ) + # Linear chain: m1 -> m2 -> m3. Plus a sibling m4 off m2 (should NOT + # be copied since we fork at m3). + studio_db.sync_chat_messages( + "src", + [ + _msg("m1", None, 1), + _msg("m2", "m1", 2), + _msg("m3", "m2", 3), + _msg("m4", "m2", 4), # sibling — must be excluded + ], + ) + + counter = {"i": 0} + + def id_factory(): + counter["i"] += 1 + return f"new-{counter['i']}" + + forked = studio_db.fork_chat_thread( + source_thread_id = "src", + branch_message_id = "m3", + new_thread_id = "fork-1", + new_title = "fork · Original", + created_at = 99, + id_factory = id_factory, + ) + assert forked is not None + assert forked["id"] == "fork-1" + assert forked["forkedFromThreadId"] == "src" + assert forked["forkedFromMessageId"] == "m3" + # Container ids reset on fork. + assert forked["openaiCodeExecContainerId"] is None + + copied = studio_db.list_chat_messages("fork-1") + # 3 ancestors (m1, m2, m3); m4 excluded. + assert len(copied) == 3 + # parent_id rewritten using new ids; root has parentId None. + assert copied[0]["parentId"] is None + assert copied[1]["parentId"] == copied[0]["id"] + assert copied[2]["parentId"] == copied[1]["id"] + # All new ids regenerated. + assert {m["id"] for m in copied}.isdisjoint({"m1", "m2", "m3"}) + + +def test_fork_chat_thread_preserves_project_id(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_project(_project("project-1")) + studio_db.upsert_chat_thread({**_thread("src"), "projectId": "project-1"}) + studio_db.upsert_chat_message(_msg("m1", None, 1)) + + forked = studio_db.fork_chat_thread( + source_thread_id = "src", + branch_message_id = "m1", + new_thread_id = "fork-1", + new_title = "fork · Original", + created_at = 99, + id_factory = lambda: "new-1", + ) + + assert forked is not None + assert forked["projectId"] == "project-1" + assert {thread["id"] for thread in studio_db.list_chat_threads(project_id = "project-1")} == { + "fork-1", + "src", + } + + +def test_fork_chat_thread_returns_none_for_missing_source(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + result = studio_db.fork_chat_thread( + source_thread_id = "nope", + branch_message_id = "m1", + new_thread_id = "fork", + new_title = "f", + created_at = 1, + id_factory = lambda: "x", + ) + assert result is None + + +def test_count_forks_for_message(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_thread(_thread("src")) + studio_db.sync_chat_messages("src", [_msg("m1", None, 1)]) + assert studio_db.count_forks_for_message("src", "m1") == 0 + + counter = {"i": 0} + + def id_factory(): + counter["i"] += 1 + return f"id-{counter['i']}" + + studio_db.fork_chat_thread( + source_thread_id = "src", + branch_message_id = "m1", + new_thread_id = "f1", + new_title = "f1", + created_at = 2, + id_factory = id_factory, + ) + studio_db.fork_chat_thread( + source_thread_id = "src", + branch_message_id = "m1", + new_thread_id = "f2", + new_title = "f2", + created_at = 3, + id_factory = id_factory, + ) + assert studio_db.count_forks_for_message("src", "m1") == 2 diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 838fce4fa7..ce1beb24d7 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -50,6 +50,11 @@ import { DropdownMenuSubTrigger, DropdownMenuTrigger, } from "@/components/ui/dropdown-menu"; +import { + CHAT_HISTORY_UPDATED_EVENT, + forkChatThread, + getForkCount, +} from "@/features/chat/api/chat-api"; import { sentAudioNames } from "@/features/chat/api/chat-adapter"; import { PromptStorageDialog, @@ -126,6 +131,7 @@ import { ChevronLeftIcon, ChevronRightIcon, Columns2Icon, + GitBranchIcon, GlobeIcon, HeadphonesIcon, MoreHorizontalIcon, @@ -2729,6 +2735,111 @@ const AssistantMessage: FC = () => { const COPY_RESET_MS = 2000; +const ForkCountBadge: FC = () => { + const aui = useAui(); + const messageId = useAuiState(({ message }) => message.id); + const [count, setCount] = useState(0); + + useEffect(() => { + let cancelled = false; + const refresh = () => { + const remoteId = aui.threadListItem().getState().remoteId; + if (!remoteId) { + if (!cancelled) setCount(0); + return; + } + void getForkCount(remoteId, messageId) + .then((n) => { + if (!cancelled) setCount(n); + }) + .catch(() => { + /* swallow: badge is non-critical */ + }); + }; + refresh(); + const handler = () => refresh(); + window.addEventListener(CHAT_HISTORY_UPDATED_EVENT, handler); + return () => { + cancelled = true; + window.removeEventListener(CHAT_HISTORY_UPDATED_EVENT, handler); + }; + }, [aui, messageId]); + + if (count <= 0) return null; + return ( + + + {count} + + ); +}; + +const useForkMessageAction = () => { + const aui = useAui(); + const navigate = useNavigate(); + const messageId = useAuiState(({ message }) => message.id); + const isRunning = useAuiState(({ thread }) => thread.isRunning); + const [pending, setPending] = useState(false); + + const handleFork = async () => { + const remoteId = aui.threadListItem().getState().remoteId; + if (!remoteId) { + toast.error("Cannot fork an unsaved chat"); + return; + } + setPending(true); + try { + const result = await forkChatThread(remoteId, { + messageId, + newThreadId: crypto.randomUUID(), + createdAt: Date.now(), + }); + useChatRuntimeStore.getState().setActiveThreadId(result.thread.id); + navigate({ + to: "/chat", + search: { thread: result.thread.id }, + replace: false, + }); + if (result.containerSnapshotWarning) { + toast.info("Fork created", { + description: result.containerSnapshotWarning, + }); + } else { + toast.success("Fork created"); + } + } catch (error) { + console.error("Failed to fork", error); + toast.error("Failed to fork", { + description: error instanceof Error ? error.message : undefined, + }); + } finally { + setPending(false); + } + }; + + return { + forkMessage: handleFork, + forkDisabled: isRunning || pending, + }; +}; + +const ForkMessageButton: FC = () => { + const { forkMessage, forkDisabled } = useForkMessageAction(); + + return ( + + + + ); +}; + const DeleteMessageButton: FC = () => { const aui = useAui(); const messageId = useAuiState(({ message }) => message.id); @@ -2799,6 +2910,8 @@ const CopyButton: FC = () => { }; const AssistantActionBar: FC = () => { + const { forkMessage, forkDisabled } = useForkMessageAction(); + return ( { + @@ -2824,10 +2938,18 @@ const AssistantActionBar: FC = () => { side="bottom" align="start" onCloseAutoFocus={(e) => e.preventDefault()} - className="aui-action-bar-more-content z-50 min-w-32 overflow-hidden rounded-full bg-popover p-1 text-popover-foreground shadow-[0_2px_8px_-2px_rgba(0,0,0,0.16)] dark:shadow-none" + className="aui-action-bar-more-content z-50 min-w-32 overflow-hidden rounded-[21px] bg-popover px-[9px] py-2 text-popover-foreground shadow-[0_2px_8px_-2px_rgba(0,0,0,0.16)] dark:shadow-none" > + void forkMessage()} + className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-[12px] px-3 py-2 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground data-[disabled]:pointer-events-none data-[disabled]:opacity-50" + > + + Fork in new chat + - + Export as Markdown @@ -2894,6 +3016,8 @@ const UserActionBar: FC = () => { /> + + ); diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index 7f5cc2f974..33493ecda0 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -379,6 +379,45 @@ export async function updateChatThread( return thread; } +export interface ForkChatThreadResult { + thread: ThreadRecord; + messages: MessageRecord[]; + containerSnapshotWarning: string | null; +} + +export async function forkChatThread( + threadId: string, + args: { messageId: string; newThreadId: string; createdAt: number }, +): Promise { + const response = await authFetch( + `/api/chat/threads/${encodeURIComponent(threadId)}/fork`, + { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(args), + }, + ); + const data = await parseJsonOrThrow<{ + thread: ThreadRecord; + messages: MessageRecord[]; + containerSnapshotWarning: string | null; + }>(response); + notifyChatHistoryUpdated(); + return data; +} + +export async function getForkCount( + threadId: string, + messageId: string, +): Promise { + const response = await authFetch( + `/api/chat/threads/${encodeURIComponent(threadId)}/messages/${encodeURIComponent(messageId)}/forks`, + ); + if (response.status === 404) return 0; + const data = await parseJsonOrThrow<{ count: number }>(response); + return data.count; +} + export async function deleteChatThreads(threadIds: string[]): Promise { if (threadIds.length === 0) return; const response = await authFetch("/api/chat/threads", { diff --git a/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts b/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts index 71d6079fc7..55d7777e43 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts @@ -27,6 +27,7 @@ export interface SidebarItem { id: string; title: string; createdAt: number; + isFork?: boolean; projectId?: string | null; } @@ -64,6 +65,7 @@ export function groupThreads( id: t.id, title: t.title, createdAt: t.createdAt, + isFork: Boolean(t.forkedFromThreadId), projectId: t.projectId ?? null, }); } diff --git a/studio/frontend/src/features/chat/thread-sidebar.tsx b/studio/frontend/src/features/chat/thread-sidebar.tsx index 0d35facb2a..9e65fc5749 100644 --- a/studio/frontend/src/features/chat/thread-sidebar.tsx +++ b/studio/frontend/src/features/chat/thread-sidebar.tsx @@ -257,6 +257,14 @@ export function ThreadSidebar({ isActive={activeId === item.id} onClick={() => onSelect(viewForItem(item))} > + {item.isFork ? ( + + fork + + ) : null} {item.title} diff --git a/studio/frontend/src/features/chat/types.ts b/studio/frontend/src/features/chat/types.ts index ef4dd7dbcb..3fe69ccc26 100644 --- a/studio/frontend/src/features/chat/types.ts +++ b/studio/frontend/src/features/chat/types.ts @@ -57,6 +57,14 @@ export interface ThreadRecord { * clears this field so the next turn falls back to auto-create. */ anthropicCodeExecContainerId?: string | null; + /** + * If this thread was created via fork-from-message, points back at + * the source thread + branch-point msg. Null/undefined for non-fork + * threads. Used by the sidebar "fork" badge and the parent thread's + * "N forks" indicator on the branch-point msg. + */ + forkedFromThreadId?: string | null; + forkedFromMessageId?: string | null; } export interface MessageRecord {