feat: implement thread forking functionality with associated database… (#5810)
* 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 <noreply@anthropic.com> * [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 <noreply@anthropic.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: imagineer99 <samleejackson0@gmail.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
This commit is contained in:
Erildo
GitHub
parent
9413e72802
commit
e54b6babd6
9 changed files with 715 additions and 5 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
<span
|
||||
className="mx-1 inline-flex items-center gap-1 rounded-sm bg-primary/10 px-1.5 py-0.5 text-[10px] font-medium text-primary"
|
||||
title={`${count} fork${count === 1 ? "" : "s"} from this message`}
|
||||
>
|
||||
<GitBranchIcon strokeWidth={1.75} className="size-3" />
|
||||
{count}
|
||||
</span>
|
||||
);
|
||||
};
|
||||
|
||||
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 (
|
||||
<TooltipIconButton
|
||||
tooltip="Fork from here"
|
||||
disabled={forkDisabled}
|
||||
onClick={forkMessage}
|
||||
>
|
||||
<GitBranchIcon strokeWidth={1.75} className="size-icon" />
|
||||
</TooltipIconButton>
|
||||
);
|
||||
};
|
||||
|
||||
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 (
|
||||
<ActionBarPrimitive.Root
|
||||
hideWhenRunning={true}
|
||||
|
|
@ -2810,6 +2923,7 @@ const AssistantActionBar: FC = () => {
|
|||
<RefreshCwIcon strokeWidth={1.75} className="size-icon" />
|
||||
</TooltipIconButton>
|
||||
</ActionBarPrimitive.Reload>
|
||||
<ForkCountBadge />
|
||||
<DeleteMessageButton />
|
||||
<ActionBarMorePrimitive.Root>
|
||||
<ActionBarMorePrimitive.Trigger asChild={true}>
|
||||
|
|
@ -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"
|
||||
>
|
||||
<ActionBarMorePrimitive.Item
|
||||
disabled={forkDisabled}
|
||||
onSelect={() => 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"
|
||||
>
|
||||
<GitBranchIcon strokeWidth={1.75} className="size-icon" />
|
||||
Fork in new chat
|
||||
</ActionBarMorePrimitive.Item>
|
||||
<ActionBarPrimitive.ExportMarkdown asChild={true}>
|
||||
<ActionBarMorePrimitive.Item className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-full px-3 py-1.5 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground">
|
||||
<ActionBarMorePrimitive.Item 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">
|
||||
<HugeiconsIcon icon={Download01Icon} strokeWidth={1.75} className="size-icon" />
|
||||
Export as Markdown
|
||||
</ActionBarMorePrimitive.Item>
|
||||
|
|
@ -2894,6 +3016,8 @@ const UserActionBar: FC = () => {
|
|||
/>
|
||||
</TooltipIconButton>
|
||||
</ActionBarPrimitive.Edit>
|
||||
<ForkCountBadge />
|
||||
<ForkMessageButton />
|
||||
<DeleteMessageButton />
|
||||
</ActionBarPrimitive.Root>
|
||||
);
|
||||
|
|
|
|||
|
|
@ -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<ForkChatThreadResult> {
|
||||
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<number> {
|
||||
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<void> {
|
||||
if (threadIds.length === 0) return;
|
||||
const response = await authFetch("/api/chat/threads", {
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
});
|
||||
}
|
||||
|
|
|
|||
|
|
@ -257,6 +257,14 @@ export function ThreadSidebar({
|
|||
isActive={activeId === item.id}
|
||||
onClick={() => onSelect(viewForItem(item))}
|
||||
>
|
||||
{item.isFork ? (
|
||||
<span
|
||||
className="mr-1 rounded-sm bg-primary/10 px-1.5 py-0.5 text-[10px] font-semibold uppercase tracking-wide text-primary"
|
||||
title="Forked from another chat"
|
||||
>
|
||||
fork
|
||||
</span>
|
||||
) : null}
|
||||
<span>{item.title}</span>
|
||||
</SidebarMenuButton>
|
||||
<DropdownMenu>
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue