From 61ed4cac5143146c2cc80a016ea2874a1c5ef3b2 Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Fri, 22 May 2026 14:18:05 +0100 Subject: [PATCH] Studio: persist chat history in backend storage (#5272) * feat: Persist chat history in backend storage * [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 * Address chat tombstone batching review * fix: update desktop auth routes stub * chat db settings storage * chat db settings routes * chat db settings client * chat db settings store * chat db settings wiring * chat db history storage * chat db settings migration * chat db settings fallback * chat db container metadata * chat db legacy migration fixes * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * chat ci auth background reads * chat auth storage fixes * chat migration final fixes * chat export batch message lookup * chat history review fixes * chat prune sync fix * chat settings hydration retry * gate settings persistence * Scope chat-history rows by subject; fix hijack, clear-confirm, hydrate race Backend storage and routes: - chat_threads / chat_messages / chat_settings carry a NOT NULL subject column with composite PRIMARY KEY (id, subject). Two authenticated identities can no longer see or wipe each other's data. - Pre-existing rows on an existing studio.db migrate under sentinel subject __legacy_unscoped__ via rename + rebuild + copy; single-user installs see no behavior change. - ON CONFLICT(id, subject) DO UPDATE ... WHERE chat_messages.thread_id = excluded.thread_id refuses cross-thread re-parenting via upsert. upsert_chat_message + sync_chat_messages now raise ChatMessageThreadMismatch which the routes map to HTTP 409. - replace_thread_messages rejects body messages whose threadId does not match the URL thread (HTTP 400) instead of silently rewriting them. - DELETE /api/chat requires ?confirm=true, returns row count, logs the subject and count. - upsert_chat_settings_merge does read + deep-merge + write inside a single BEGIN IMMEDIATE so concurrent writers no longer drop each other's updates. The route delegates to this helper. - New POST /api/chat/messages:batch returns {thread_id -> messages[]} for many threads in one HTTP call. Subject-scoped. Unknown ids return empty lists instead of 404 so the sidebar/search caller can rebuild atomically. Frontend: - chat-runtime-store: hydrate-failure catch sets settingsHydrated:true so a transient backend blip no longer permanently disables persistence. setParams bumps inferenceParamMutationVersions unconditionally so a slow hydration response cannot clobber a pre-hydrate user edit. saveSettingsPatch replaces the serial chain with a debounced pendingPatch + deep merge; flush on beforeunload. - chat-history-storage: clearStoredChats returns ClearStoredChatsResult distinguishing backend / legacy / both outcomes. listStoredChatThreadsWithMessages uses the batched fetch (one HTTP call) instead of Promise.all per-thread; legacy Dexie fallback only fires when the batch result is empty. - chat-api: batchListChatMessages with graceful 404 / 405 fallback to per-thread listChatMessages for older servers. - chat-thread-tombstones: store {id, deletedAt} tuples with 90-day GC and a 5000-entry cap so localStorage stays bounded. Back-compat reads pre-fix plain strings. Adds removeChatThreadTombstones (rollback) and clearAllChatThreadTombstones (post-legacy-purge clean-up). - use-chat-sidebar-items: deleteChatItem tombstones synchronously BEFORE the backend round-trip and rolls back on failure (restores pre-PR optimistic UX). 300 ms trailing debounce on CHAT_HISTORY_UPDATED_EVENT plus requestSeq guard so stream-time event bursts produce at most one fetch per quiet window. Tests: - studio/backend/tests/pr5272_sim/ adds 64 regression tests covering schema migration from pre-fix shape, subject scoping, cross-thread hijack, bulk-replace mismatch, clear-confirm, concurrent settings, unicode + 2MB content + SQL-injection-safe binding, chunking boundary at 900 and 901 ids, batched endpoint (multi-subject + 1200 ids + per-thread order), and grep contracts for the frontend patches. test_chat_history_storage.py updated to pass subject. Verified locally on Linux + macOS + Windows GitHub Actions runners (staging fork): 64 pass + 2 from the PR's own backend test on all three OSes. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Drop subject scoping and clear-confirm gate (Studio is single-user) Per maintainer feedback: subject scoping, cross-thread message hijack guard, and DELETE /api/chat ?confirm=true gate are unnecessary because Studio is intentionally single-user (the client already shows a confirm dialog before clear-all). This commit reverts those backend changes and keeps only the non-multi-user pieces from the earlier fix commit: - studio_db.py: restored to pre-fix shape; adds upsert_chat_settings_merge which does atomic read + deep-merge + write under BEGIN IMMEDIATE so two concurrent slider drags cannot drop one another's updates. - routes/chat_history.py: restored; put_settings now calls the atomic merge instead of doing the read-merge-write across three separate connections. Adds POST /api/chat/messages:batch to collapse the sidebar/search rebuild from N round-trips to 1. - frontend/api/chat-api.ts: align batchListChatMessages request and response keys with the backend (threadIds / messagesByThreadId). - tests/test_chat_history_storage.py: add atomic-merge concurrency test, deep-merge nested-key test, and 901-id chunking-boundary test. - Drop the pr5272_sim test directory (those tests covered the reverted subject-scoping/hijack/confirm behavior). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix sidebar delete crash, keepalive on settings beforeunload flush, search rebuild race Two correctness bugs and one perf race surfaced by a fresh code review of the prior fix commit: - chat-api.ts: notifyChatHistoryUpdated was declared as a non-exported function, but use-chat-sidebar-items.ts imports it. The import would fail tsc with TS2305 and at runtime the optimistic-delete and delete-failure rollback paths would both throw. - chat-runtime-store.ts + chat-settings-api.ts + chat-settings-storage.ts: the beforeunload settings flush is now actually keepalive. Without it the browser cancels the in-flight PUT on tab close, so the last slider drag is silently dropped (which is exactly the case the debounce+beforeunload combination was meant to protect against). - use-chat-search-index.ts: rebuilds now coalesce with a 300ms trailing debounce and discard out-of-order responses via a requestSeq guard. Matches the sibling pattern in use-chat-sidebar-items.ts so two rapid CHAT_HISTORY_UPDATED_EVENTs (run-start + run-end save during a turn) cannot land with stale data winning. - chat-thread-tombstones.ts: drop dead clearAllChatThreadTombstones with no call sites; Dexie is never wiped so the function has no use. * fix(studio): protect chat persistence writes * fix(studio): align chat history clear semantics * fix(studio): show partial chat clear feedback * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(studio): preserve chat persistence fallbacks * fix(studio): harden chat thread persistence checks * Preserve chat message timestamps * Gate chat stream on history save * Make chat thread backfill best effort * Avoid chat message 404 probe * Tighten chat legacy fallbacks * chat: server-side ledger so legacy Dexie import is recoverable The boolean localStorage sentinel (unsloth_chat_legacy_imported_to_studio_db) made importLegacyChatsIfNeeded non-recoverable: deleting studio.db while the browser keeps the flag silently hides every legacy Dexie thread from the sidebar (verified by the 3-GPU validation probe; matches the third review comment on PR #5272). Same trap fires for browser-profile sync to a fresh machine and any other path that wipes studio.db while keeping IndexedDB. Source of truth moves into studio.db itself via a new chat_legacy_import_log table keyed by legacy thread id. The ledger disappears together with studio.db, so the next launch re-runs the import from whatever Dexie still holds. localStorage stays as a per-session perf hint only. Performance, all bounded by the three new fast-paths before any backend work: A) localStorage hint says "imported earlier in this session" -- 0 network, ~0 ms. Covers the warm sidebar mount. B) indexedDB.databases() reports no "unsloth-chat" DB -- 0 network, ~1 ms. Covers every new user who never had the old browser-only Studio (the common case after launch). C) db.threads.count() + db.messages.count() are both 0 -- 0 network, ~5 ms. Covers returning users who migrated long ago and Dexie was never repopulated. Only when all three miss does the code talk to the backend (GET /api/chat/import-ledger -> diff vs Dexie -> existing import path -> POST /api/chat/import-ledger to record what was just imported). Per-thread tracking is enough because Dexie is read-only after this PR; a thread's message set does not grow. Backend deployments that predate the import-ledger routes are handled transparently: the client treats 404/405 as an empty ledger and re-runs the (idempotent via UPSERT) import on next launch. Changes: - storage/studio_db.py: new chat_legacy_import_log table (WITHOUT ROWID, PK on legacy_thread_id) + list_chat_legacy_import_log() + record_chat_legacy_import_log() (idempotent batch UPSERT). - routes/chat_history.py: GET + POST /api/chat/import-ledger with the obvious request/response models. - frontend api/chat-api.ts: listChatImportLedger() (returns a Set for O(1) diff) + recordChatImportLedger(), both with 404/405 fallback. - frontend utils/chat-history-storage.ts: importLegacyChatsIfNeeded gains three fast-paths, ledger fetch on the slow path, and writes the ledger after a successful import. The localStorage helper is unchanged on the surface; it just stops being authoritative. - tests: 5 new test_legacy_import_log_* cases (empty default, record + list round-trip, idempotency, input dedup, empty/null ignore). All 9 pre-existing tests still pass. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Make the legacy-import recovery actually recoverable The previous commit added a server-side ledger to make Dexie -> studio.db import recoverable after a studio.db wipe, but the localStorage perf hint still short-circuited the import gate before the ledger was ever consulted. After a wipe, the hint stayed "true" and the bulk re-import never ran -- the ledger sat empty and only the per-thread lazy materialize-on-continue path restored data. Changes: - Remove the localStorage short-circuit from importLegacyChatsIfNeeded so the ledger is checked on every fresh tab. legacyChatImportPromise keeps the per-session cache; the hint now only matters for the listing paths. - Batch the slow path: one db.messages.where().anyOf().toArray() and one batchListChatMessages() instead of 2N round-trips. At 1k threads this drops a multi-second blocking import to a single request pair. - recordChatImportLedger returns {accepted, inserted, supported}. The localStorage hint is only flipped when supported is true, so old backends (404 / 405 / 501) no longer permanently poison recovery. - Ledger backfill: threads already present in chat_threads but missing from the ledger now get added too, so old-FE-then-new-FE deployments don't redo the diff every launch. - Backend response field renamed recorded -> {accepted, inserted}. accepted is the deduped non-empty input count; inserted is the rows actually new (via INSERT ... RETURNING). Bounded by Field(max_length= 10_000) on the request payload. - Storage helpers renamed: chat_legacy_import_log -> chat_legacy_imports, record_* -> upsert_* to match the existing noun/verb conventions. - DEXIE_DB_NAME exported from db.ts; duplicate constant in chat-history-storage.ts removed. - 3 new route-level tests for /api/chat/import-ledger covering the round-trip, the (accepted, inserted) split, and the 10k payload cap. All 18 chat-history tests pass. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: shine1i Co-authored-by: danielhanchen Co-authored-by: Daniel Han --- studio/backend/main.py | 3 + studio/backend/routes/__init__.py | 2 + studio/backend/routes/chat_history.py | 397 +++++++++ studio/backend/storage/studio_db.py | 664 ++++++++++++++- .../backend/tests/test_chat_history_routes.py | 109 +++ .../tests/test_chat_history_storage.py | 308 +++++++ studio/backend/tests/test_desktop_auth.py | 1 + studio/frontend/src/features/auth/api.ts | 40 +- .../src/features/chat/api/chat-adapter.ts | 390 +++++---- .../src/features/chat/api/chat-api.ts | 277 ++++++- .../features/chat/api/chat-settings-api.ts | 90 ++ .../src/features/chat/api/providers-api.ts | 452 +++++----- .../frontend/src/features/chat/chat-page.tsx | 115 +-- .../src/features/chat/chat-settings-sheet.tsx | 266 ++---- .../components/openai-code-exec-section.tsx | 64 +- studio/frontend/src/features/chat/db.ts | 6 +- .../chat/hooks/use-chat-search-index.ts | 88 +- .../chat/hooks/use-chat-sidebar-items.ts | 110 ++- studio/frontend/src/features/chat/index.ts | 2 + .../src/features/chat/runtime-provider.tsx | 428 ++++++---- .../chat/stores/chat-runtime-store.ts | 556 +++++++++---- .../chat/utils/chat-history-storage.ts | 773 ++++++++++++++++++ .../chat/utils/chat-settings-storage.ts | 403 +++++++++ .../chat/utils/chat-thread-tombstones.ts | 115 ++- .../features/chat/utils/clear-all-chats.ts | 13 +- .../chat/utils/delete-thread-message.ts | 60 +- .../chat/utils/export-chat-history.ts | 24 +- .../src/features/settings/tabs/chat-tab.tsx | 51 +- .../features/settings/tabs/general-tab.tsx | 6 +- 29 files changed, 4677 insertions(+), 1136 deletions(-) create mode 100644 studio/backend/routes/chat_history.py create mode 100644 studio/backend/tests/test_chat_history_routes.py create mode 100644 studio/backend/tests/test_chat_history_storage.py create mode 100644 studio/frontend/src/features/chat/api/chat-settings-api.ts create mode 100644 studio/frontend/src/features/chat/utils/chat-history-storage.ts create mode 100644 studio/frontend/src/features/chat/utils/chat-settings-storage.ts diff --git a/studio/backend/main.py b/studio/backend/main.py index d4593c2ab4..004ae404cd 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -114,6 +114,7 @@ from datetime import datetime # Import routers from routes import ( auth_router, + chat_history_router, data_recipe_router, datasets_router, export_router, @@ -367,6 +368,7 @@ _BODY_PROTECTED_PREFIXES = ( "/api/inference", "/api/data-recipe", "/api/datasets", + "/api/chat", "/api/train", "/api/export", ) @@ -509,6 +511,7 @@ app.add_middleware( app.include_router(auth_router, prefix = "/api/auth", tags = ["auth"]) app.include_router(training_router, prefix = "/api/train", tags = ["training"]) app.include_router(models_router, prefix = "/api/models", tags = ["models"]) +app.include_router(chat_history_router, prefix = "/api/chat", tags = ["chat"]) app.include_router(inference_router, prefix = "/api/inference", tags = ["inference"]) # Studio-only inference endpoints (cancel, etc.) are intentionally NOT # exposed on the /v1 OpenAI-compat prefix below. diff --git a/studio/backend/routes/__init__.py b/studio/backend/routes/__init__.py index 62320b9084..6bb5d15e8e 100644 --- a/studio/backend/routes/__init__.py +++ b/studio/backend/routes/__init__.py @@ -14,6 +14,7 @@ from routes.auth import router as auth_router from routes.data_recipe import router as data_recipe_router from routes.export import router as export_router from routes.training_history import router as training_history_router +from routes.chat_history import router as chat_history_router from routes.providers import router as providers_router __all__ = [ @@ -26,5 +27,6 @@ __all__ = [ "data_recipe_router", "export_router", "training_history_router", + "chat_history_router", "providers_router", ] diff --git a/studio/backend/routes/chat_history.py b/studio/backend/routes/chat_history.py new file mode 100644 index 0000000000..ed808040d2 --- /dev/null +++ b/studio/backend/routes/chat_history.py @@ -0,0 +1,397 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Chat history API routes backed by studio.db. +""" + +from typing import Any, Literal, Optional + +from fastapi import APIRouter, Depends, HTTPException, Query +from pydantic import BaseModel, ConfigDict, Field, ValidationError + +from auth.authentication import get_current_subject +from storage.studio_db import ( + ChatMessageConflictError, + CorruptSettingsError, + clear_chat_history, + count_chat_threads, + delete_chat_threads, + get_chat_thread, + get_chat_message, + list_chat_legacy_imports, + list_chat_settings, + list_chat_messages, + list_chat_messages_for_threads, + list_chat_threads, + sync_chat_messages, + update_chat_thread, + upsert_chat_legacy_imports, + upsert_chat_message, + upsert_chat_settings_merge, + upsert_chat_thread, +) + +router = APIRouter() + + +class ChatThread(BaseModel): + id: str + title: str = "New Chat" + modelType: Literal["base", "lora", "model1", "model2"] + modelId: str = "" + pairId: Optional[str] = None + archived: bool = False + createdAt: int + openaiCodeExecContainerId: Optional[str] = None + anthropicCodeExecContainerId: Optional[str] = None + + +class ChatThreadPatch(BaseModel): + title: Optional[str] = None + modelType: Optional[Literal["base", "lora", "model1", "model2"]] = None + modelId: Optional[str] = None + pairId: Optional[str] = None + archived: Optional[bool] = None + createdAt: Optional[int] = None + openaiCodeExecContainerId: Optional[str] = None + anthropicCodeExecContainerId: Optional[str] = None + + +class ChatMessage(BaseModel): + id: str + threadId: str + parentId: Optional[str] = None + role: str + content: Any = Field(default_factory = list) + attachments: Optional[Any] = None + metadata: Optional[dict[str, Any]] = None + createdAt: int + + +class ChatThreadListResponse(BaseModel): + threads: list[ChatThread] + + +class ChatMessageListResponse(BaseModel): + messages: list[ChatMessage] + + +class ChatMessageSyncRequest(BaseModel): + messages: list[ChatMessage] + pruneMissing: bool = False + + +class ChatDeleteRequest(BaseModel): + ids: list[str] + + +class ChatCountResponse(BaseModel): + count: int + + +class ChatExportResponse(BaseModel): + exportedAt: str + version: int + threadCount: int + threads: list[ChatThread] + messages: list[ChatMessage] + + +class ChatInferenceSettings(BaseModel): + model_config = ConfigDict(extra = "forbid") + + temperature: Optional[float] = None + topP: Optional[float] = None + topK: Optional[float] = None + minP: Optional[float] = None + repetitionPenalty: Optional[float] = None + presencePenalty: Optional[float] = None + maxSeqLength: Optional[float] = None + maxTokens: Optional[float] = None + systemPrompt: Optional[str] = None + trustRemoteCode: Optional[bool] = None + + +class ChatPreset(BaseModel): + model_config = ConfigDict(extra = "forbid") + + name: str + params: ChatInferenceSettings + + +class ChatSettingsPayload(BaseModel): + model_config = ConfigDict(extra = "forbid") + + inferenceParams: Optional[ChatInferenceSettings] = None + customPresets: Optional[list[ChatPreset]] = None + activePreset: Optional[str] = None + activePresetSource: Optional[Literal["builtin-default", "custom", "modified"]] = ( + None + ) + autoTitle: Optional[bool] = None + reasoningEffort: Optional[ + Literal["none", "minimal", "low", "medium", "high", "max", "xhigh"] + ] = None + preserveThinking: Optional[bool] = None + autoHealToolCalls: Optional[bool] = None + maxToolCallsPerMessage: Optional[int] = Field(default = None, ge = 1) + toolCallTimeout: Optional[int] = Field(default = None, ge = 1) + + +class ChatSettingsResponse(BaseModel): + settings: dict[str, Any] + + +class ChatMessagesBatchRequest(BaseModel): + threadIds: list[str] + + +class ChatMessagesBatchResponse(BaseModel): + messagesByThreadId: dict[str, list[ChatMessage]] + + +class ChatImportLedgerResponse(BaseModel): + # Plain list of legacy thread ids. Keeping the payload key-less keeps + # the client diff to a single Set construction. + threadIds: list[str] + + +class ChatImportLedgerRecordRequest(BaseModel): + # 10k cap keeps the request body bounded; real users have << 1k threads. + threadIds: list[str] = Field(default_factory = list, max_length = 10_000) + + +class ChatImportLedgerRecordResponse(BaseModel): + # accepted: deduped non-empty input count. inserted: rows actually new + # (ON CONFLICT DO NOTHING skips already-recorded ids). The client uses + # `accepted >= 0` as the "endpoint exists" signal and ignores the split + # otherwise. + accepted: int + inserted: int + + +@router.get("/threads", response_model = ChatThreadListResponse) +async def list_threads( + model_type: Optional[str] = Query(None), + pair_id: Optional[str] = Query(None), + include_archived: bool = Query(True), + current_subject: str = Depends(get_current_subject), +): + threads = list_chat_threads( + model_type = model_type, + pair_id = pair_id, + include_archived = include_archived, + ) + return ChatThreadListResponse(threads = [ChatThread(**t) for t in threads]) + + +@router.post("/threads", response_model = ChatThread) +async def save_thread( + payload: ChatThread, + current_subject: str = Depends(get_current_subject), +): + return ChatThread(**upsert_chat_thread(payload.model_dump())) + + +@router.get("/threads/{thread_id}", response_model = ChatThread) +async def get_thread( + thread_id: str, + current_subject: str = Depends(get_current_subject), +): + thread = get_chat_thread(thread_id) + if thread is None: + raise HTTPException(status_code = 404, detail = f"Thread {thread_id} not found") + return ChatThread(**thread) + + +@router.patch("/threads/{thread_id}", response_model = ChatThread) +async def patch_thread( + thread_id: str, + payload: ChatThreadPatch, + current_subject: str = Depends(get_current_subject), +): + patch = payload.model_dump(exclude_unset = True) + for field in ("title", "modelType", "modelId", "archived", "createdAt"): + if field in patch and patch[field] is None: + raise HTTPException(status_code = 400, detail = f"{field} cannot be null") + thread = update_chat_thread( + thread_id, + patch, + ) + if thread is None: + raise HTTPException(status_code = 404, detail = f"Thread {thread_id} not found") + return ChatThread(**thread) + + +@router.delete("/threads") +async def delete_threads( + payload: ChatDeleteRequest, + current_subject: str = Depends(get_current_subject), +): + delete_chat_threads(payload.ids) + return {"status": "deleted"} + + +@router.get("/threads/{thread_id}/messages", response_model = ChatMessageListResponse) +async def get_thread_messages( + thread_id: str, + current_subject: str = Depends(get_current_subject), +): + if get_chat_thread(thread_id) is None: + raise HTTPException(status_code = 404, detail = f"Thread {thread_id} not found") + return ChatMessageListResponse( + messages = [ChatMessage(**m) for m in list_chat_messages(thread_id)] + ) + + +@router.post("/messages:batch", response_model = ChatMessagesBatchResponse) +async def batch_thread_messages( + payload: ChatMessagesBatchRequest, + current_subject: str = Depends(get_current_subject), +): + """One round-trip per sidebar/search rebuild instead of N. Unknown thread + ids are returned as empty lists so callers don't need a pre-flight.""" + by_thread: dict[str, list[ChatMessage]] = {tid: [] for tid in payload.threadIds} + for m in list_chat_messages_for_threads(payload.threadIds): + tid = m["threadId"] + if tid in by_thread: + by_thread[tid].append(ChatMessage(**m)) + return ChatMessagesBatchResponse(messagesByThreadId = by_thread) + + +@router.get("/threads/{thread_id}/messages/{message_id}", response_model = ChatMessage) +async def get_thread_message( + thread_id: str, + message_id: str, + current_subject: str = Depends(get_current_subject), +): + if get_chat_thread(thread_id) is None: + raise HTTPException(status_code = 404, detail = f"Thread {thread_id} not found") + message = get_chat_message(thread_id, message_id) + if message is None: + raise HTTPException(status_code = 404, detail = f"Message {message_id} not found") + return ChatMessage(**message) + + +@router.put("/threads/{thread_id}/messages/{message_id}", response_model = ChatMessage) +async def save_thread_message( + thread_id: str, + message_id: str, + payload: ChatMessage, + current_subject: str = Depends(get_current_subject), +): + if thread_id != payload.threadId or message_id != payload.id: + raise HTTPException(status_code = 400, detail = "Message id mismatch") + if get_chat_thread(thread_id) is None: + raise HTTPException(status_code = 404, detail = f"Thread {thread_id} not found") + try: + return ChatMessage(**upsert_chat_message(payload.model_dump())) + except ChatMessageConflictError as exc: + raise HTTPException(status_code = 409, detail = str(exc)) from exc + + +@router.put("/threads/{thread_id}/messages", response_model = ChatMessageListResponse) +async def replace_thread_messages( + thread_id: str, + payload: ChatMessageSyncRequest, + current_subject: str = Depends(get_current_subject), +): + mismatched_ids = [ + message.id for message in payload.messages if message.threadId != thread_id + ] + if mismatched_ids: + preview = ", ".join(mismatched_ids[:5]) + suffix = ( + "" if len(mismatched_ids) <= 5 else f" (+{len(mismatched_ids) - 5} more)" + ) + raise HTTPException( + status_code = 400, + detail = f"Message threadId mismatch: {preview}{suffix}", + ) + if get_chat_thread(thread_id) is None: + raise HTTPException(status_code = 404, detail = f"Thread {thread_id} not found") + messages = [message.model_dump() for message in payload.messages] + try: + return ChatMessageListResponse( + messages = [ + ChatMessage(**m) + for m in sync_chat_messages( + thread_id, + messages, + prune_missing = payload.pruneMissing, + ) + ] + ) + except ChatMessageConflictError as exc: + raise HTTPException(status_code = 409, detail = str(exc)) from exc + + +@router.get("/count", response_model = ChatCountResponse) +async def count_threads(current_subject: str = Depends(get_current_subject)): + return ChatCountResponse(count = count_chat_threads()) + + +@router.get("/import-ledger", response_model = ChatImportLedgerResponse) +async def get_import_ledger(current_subject: str = Depends(get_current_subject)): + """Legacy-Dexie import ledger. Returns the set of legacy thread ids + already copied into chat_threads / chat_messages. The frontend + uses this on every fresh tab open to decide whether to re-run the + Dexie -> studio.db import. Source of truth lives inside studio.db + so a studio.db wipe makes the import recoverable.""" + return ChatImportLedgerResponse(threadIds = list_chat_legacy_imports()) + + +@router.post("/import-ledger", response_model = ChatImportLedgerRecordResponse) +async def record_import_ledger( + payload: ChatImportLedgerRecordRequest, + current_subject: str = Depends(get_current_subject), +): + """Mark each legacy thread id as imported. Idempotent.""" + accepted, inserted = upsert_chat_legacy_imports(payload.threadIds) + return ChatImportLedgerRecordResponse(accepted = accepted, inserted = inserted) + + +@router.delete("") +async def clear_history(current_subject: str = Depends(get_current_subject)): + clear_chat_history() + return {"status": "deleted"} + + +@router.get("/settings", response_model = ChatSettingsResponse) +async def get_settings(current_subject: str = Depends(get_current_subject)): + return ChatSettingsResponse(settings = list_chat_settings()) + + +@router.put("/settings", response_model = ChatSettingsResponse) +async def put_settings( + payload: dict[str, Any], + current_subject: str = Depends(get_current_subject), +): + try: + parsed = ChatSettingsPayload.model_validate(payload) + except ValidationError as exc: + raise HTTPException(status_code = 400, detail = exc.errors()) from exc + # Atomic read + deep-merge + write inside one BEGIN IMMEDIATE so two + # concurrent slider drags can't drop each other's updates. + try: + return ChatSettingsResponse( + settings = upsert_chat_settings_merge(parsed.model_dump(exclude_unset = True)) + ) + except CorruptSettingsError as exc: + raise HTTPException(status_code = 409, detail = str(exc)) from exc + + +@router.get("/export", response_model = ChatExportResponse) +async def export_history(current_subject: str = Depends(get_current_subject)): + from datetime import datetime, timezone + + threads = list_chat_threads(include_archived = True) + messages = list_chat_messages_for_threads([thread["id"] for thread in threads]) + return ChatExportResponse( + exportedAt = datetime.now(timezone.utc).isoformat(), + version = 1, + threadCount = len(threads), + threads = [ChatThread(**thread) for thread in threads], + messages = [ChatMessage(**message) for message in messages], + ) diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index 8dc29a9f24..de89b6cbd2 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -19,7 +19,7 @@ import threading from datetime import datetime, timezone logger = logging.getLogger(__name__) -from typing import Optional +from typing import Any, Iterable, Optional from utils.paths import studio_db_path, ensure_dir @@ -54,6 +54,7 @@ def _denied_path_prefixes() -> list[str]: _schema_lock = threading.Lock() _schema_ready = False +_SQLITE_IN_CHUNK_SIZE = 900 def _ensure_schema(conn: sqlite3.Connection) -> None: @@ -118,6 +119,92 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: ) """ ) + conn.execute( + """ + CREATE TABLE IF NOT EXISTS chat_threads ( + id TEXT NOT NULL PRIMARY KEY, + title TEXT NOT NULL, + model_type TEXT NOT NULL, + model_id TEXT, + pair_id TEXT, + archived INTEGER NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + openai_code_exec_container_id TEXT, + anthropic_code_exec_container_id TEXT + ) + """ + ) + chat_thread_cols = { + row[1] for row in conn.execute("PRAGMA table_info(chat_threads)").fetchall() + } + if "openai_code_exec_container_id" not in chat_thread_cols: + 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" + ) + conn.execute( + """ + CREATE TABLE IF NOT EXISTS chat_messages ( + id TEXT NOT NULL PRIMARY KEY, + thread_id TEXT NOT NULL REFERENCES chat_threads(id) ON DELETE CASCADE, + parent_id TEXT, + role TEXT NOT NULL, + content_json TEXT NOT NULL, + attachments_json TEXT, + metadata_json TEXT, + created_at INTEGER NOT NULL + ) + """ + ) + conn.execute( + "CREATE INDEX IF NOT EXISTS idx_chat_threads_model_type_created_at ON chat_threads(model_type, created_at)" + ) + conn.execute( + "CREATE INDEX IF NOT EXISTS idx_chat_threads_pair_id ON chat_threads(pair_id)" + ) + conn.execute( + "CREATE INDEX IF NOT EXISTS idx_chat_messages_thread_id_created_at ON chat_messages(thread_id, created_at)" + ) + conn.execute( + """ + CREATE TABLE IF NOT EXISTS chat_settings ( + key TEXT NOT NULL PRIMARY KEY, + value_json TEXT NOT NULL, + updated_at TEXT NOT NULL + ) + """ + ) + conn.execute( + """ + CREATE TABLE IF NOT EXISTS chat_settings_quarantine ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + key TEXT NOT NULL, + value_json TEXT NOT NULL, + reason TEXT NOT NULL, + quarantined_at TEXT NOT NULL + ) + """ + ) + # Server-side import ledger so a studio.db wipe correctly re-triggers + # the legacy Dexie import. The previous boolean localStorage sentinel + # (`unsloth_chat_legacy_imported_to_studio_db`) is non-recoverable: + # if studio.db is recreated while the browser keeps the flag, legacy + # Dexie threads are silently hidden from the sidebar. The ledger + # lives inside studio.db so it disappears together with the data it + # is supposed to track, which is the recovery the boolean lacked. + # Keyed by legacy thread id; per-thread is sufficient because Dexie + # is read-only after this PR (a thread's message set does not grow). + conn.execute( + """ + CREATE TABLE IF NOT EXISTS chat_legacy_imports ( + legacy_thread_id TEXT NOT NULL PRIMARY KEY, + imported_at INTEGER NOT NULL + ) WITHOUT ROWID + """ + ) def get_connection() -> sqlite3.Connection: @@ -575,3 +662,578 @@ def remove_scan_folder(id: int) -> None: conn.commit() finally: conn.close() + + +def _json_loads(value: str | None, fallback): + if value is None: + return fallback + try: + return json.loads(value) + except (json.JSONDecodeError, TypeError): + return fallback + + +def _chat_thread_from_row(row: sqlite3.Row) -> dict: + data = dict(row) + return { + "id": data["id"], + "title": data["title"], + "modelType": data["model_type"], + "modelId": data.get("model_id") or "", + "pairId": data.get("pair_id") or None, + "archived": bool(data["archived"]), + "createdAt": data["created_at"], + "openaiCodeExecContainerId": data.get("openai_code_exec_container_id"), + "anthropicCodeExecContainerId": data.get("anthropic_code_exec_container_id"), + } + + +def _chat_message_from_row(row: sqlite3.Row) -> dict: + data = dict(row) + message = { + "id": data["id"], + "threadId": data["thread_id"], + "parentId": data.get("parent_id"), + "role": data["role"], + "content": _json_loads(data.get("content_json"), []), + "createdAt": data["created_at"], + } + attachments = _json_loads(data.get("attachments_json"), None) + metadata = _json_loads(data.get("metadata_json"), None) + if attachments is not None: + message["attachments"] = attachments + if metadata is not None: + message["metadata"] = metadata + return message + + +def upsert_chat_thread(thread: dict) -> dict: + conn = get_connection() + try: + conn.execute( + """ + INSERT INTO chat_threads + (id, title, model_type, model_id, pair_id, archived, created_at, openai_code_exec_container_id, anthropic_code_exec_container_id) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(id) DO UPDATE SET + title = excluded.title, + model_type = excluded.model_type, + model_id = excluded.model_id, + pair_id = excluded.pair_id, + 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 + """, + ( + thread["id"], + thread.get("title") or "New Chat", + thread["modelType"], + thread.get("modelId") or "", + thread.get("pairId"), + 1 if thread.get("archived") else 0, + int(thread["createdAt"]), + thread.get("openaiCodeExecContainerId"), + thread.get("anthropicCodeExecContainerId"), + ), + ) + conn.commit() + return get_chat_thread(thread["id"]) or thread + finally: + conn.close() + + +def update_chat_thread(id: str, patch: dict) -> Optional[dict]: + allowed = { + "title": ("title", patch.get("title")), + "modelType": ("model_type", patch.get("modelType")), + "modelId": ("model_id", patch.get("modelId")), + "pairId": ("pair_id", patch.get("pairId")), + "archived": ("archived", 1 if patch.get("archived") else 0), + "createdAt": ("created_at", patch.get("createdAt")), + "openaiCodeExecContainerId": ( + "openai_code_exec_container_id", + patch.get("openaiCodeExecContainerId"), + ), + "anthropicCodeExecContainerId": ( + "anthropic_code_exec_container_id", + patch.get("anthropicCodeExecContainerId"), + ), + } + assignments = [] + values = [] + for key, (column, value) in allowed.items(): + if key in patch: + assignments.append(f"{column} = ?") + values.append(value) + if not assignments: + return get_chat_thread(id) + + conn = get_connection() + try: + conn.execute( + f"UPDATE chat_threads SET {', '.join(assignments)} WHERE id = ?", + (*values, id), + ) + conn.commit() + row = conn.execute("SELECT * FROM chat_threads WHERE id = ?", (id,)).fetchone() + return _chat_thread_from_row(row) if row is not None else None + finally: + conn.close() + + +def get_chat_thread(id: str) -> Optional[dict]: + conn = get_connection() + try: + row = conn.execute("SELECT * FROM chat_threads WHERE id = ?", (id,)).fetchone() + return _chat_thread_from_row(row) if row is not None else None + finally: + conn.close() + + +def list_chat_threads( + model_type: str | None = None, + pair_id: str | None = None, + include_archived: bool = True, +) -> list[dict]: + clauses = [] + values: list[object] = [] + if model_type is not None: + clauses.append("model_type = ?") + values.append(model_type) + if pair_id is not None: + clauses.append("pair_id = ?") + values.append(pair_id) + if not include_archived: + clauses.append("archived = 0") + where = f"WHERE {' AND '.join(clauses)}" if clauses else "" + conn = get_connection() + try: + rows = conn.execute( + f"SELECT * FROM chat_threads {where} ORDER BY created_at DESC", + values, + ).fetchall() + return [_chat_thread_from_row(row) for row in rows] + finally: + conn.close() + + +def delete_chat_threads(ids: list[str]) -> None: + if not ids: + return + conn = get_connection() + try: + conn.executemany("DELETE FROM chat_threads WHERE id = ?", [(id,) for id in ids]) + conn.commit() + finally: + conn.close() + + +def clear_chat_history() -> None: + conn = get_connection() + try: + conn.execute("DELETE FROM chat_threads") + conn.commit() + finally: + conn.close() + + +def count_chat_threads() -> int: + conn = get_connection() + try: + return int(conn.execute("SELECT COUNT(*) FROM chat_threads").fetchone()[0]) + finally: + conn.close() + + +class ChatMessageConflictError(RuntimeError): + """Raised when a chat message id already belongs to another thread.""" + + +class CorruptSettingsError(RuntimeError): + """Raised when a partial settings patch would overwrite corrupt settings.""" + + +def _parse_chat_setting_json(key: str, value_json: str) -> tuple[bool, Any]: + try: + return True, json.loads(value_json) + except (json.JSONDecodeError, TypeError) as exc: + logger.warning( + "Corrupt chat_settings JSON; quarantining key=%s error=%s", + key, + exc, + ) + return False, None + + +def _load_chat_settings_for_merge( + conn: sqlite3.Connection, +) -> tuple[dict[str, Any], set[str]]: + rows = conn.execute("SELECT key, value_json FROM chat_settings").fetchall() + current: dict[str, Any] = {} + corrupt: set[str] = set() + now = datetime.now(timezone.utc).isoformat() + for row in rows: + ok, value = _parse_chat_setting_json(row["key"], row["value_json"]) + if ok: + current[row["key"]] = value + continue + corrupt.add(row["key"]) + conn.execute( + """ + INSERT INTO chat_settings_quarantine + (key, value_json, reason, quarantined_at) + VALUES (?, ?, ?, ?) + """, + (row["key"], row["value_json"], "json_decode_error", now), + ) + conn.execute( + "DELETE FROM chat_settings WHERE key = ? AND value_json = ?", + (row["key"], row["value_json"]), + ) + return current, corrupt + + +def _raise_if_chat_message_thread_conflicts( + conn: sqlite3.Connection, + thread_id: str, + message_ids: list[str], +) -> None: + unique_ids = list(dict.fromkeys(message_ids)) + if not unique_ids: + return + conflicts: list[str] = [] + for start in range(0, len(unique_ids), _SQLITE_IN_CHUNK_SIZE): + chunk = unique_ids[start : start + _SQLITE_IN_CHUNK_SIZE] + placeholders = ",".join("?" for _ in chunk) + rows = conn.execute( + f""" + SELECT id FROM chat_messages + WHERE id IN ({placeholders}) AND thread_id != ? + ORDER BY id + """, + (*chunk, thread_id), + ).fetchall() + conflicts.extend(row["id"] for row in rows) + if conflicts: + preview = ", ".join(conflicts[:5]) + suffix = "" if len(conflicts) <= 5 else f" (+{len(conflicts) - 5} more)" + raise ChatMessageConflictError( + f"Message id already belongs to another thread: {preview}{suffix}" + ) + + +def upsert_chat_message(message: dict) -> dict: + conn = get_connection() + try: + conn.execute("BEGIN IMMEDIATE") + _raise_if_chat_message_thread_conflicts( + conn, + message["threadId"], + [message["id"]], + ) + conn.execute( + """ + INSERT INTO chat_messages + (id, thread_id, parent_id, role, content_json, attachments_json, metadata_json, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(id) DO UPDATE SET + parent_id = excluded.parent_id, + role = excluded.role, + content_json = excluded.content_json, + attachments_json = excluded.attachments_json, + metadata_json = excluded.metadata_json, + created_at = excluded.created_at + WHERE excluded.thread_id = chat_messages.thread_id + """, + ( + message["id"], + message["threadId"], + message.get("parentId"), + message["role"], + json.dumps(message.get("content", [])), + json.dumps(message.get("attachments")) + if message.get("attachments") is not None + else None, + json.dumps(message.get("metadata")) + if message.get("metadata") is not None + else None, + int(message["createdAt"]), + ), + ) + conn.commit() + return message + except Exception: + conn.rollback() + raise + finally: + conn.close() + + +def sync_chat_messages( + thread_id: str, + messages: list[dict], + prune_missing: bool = False, +) -> list[dict]: + conn = get_connection() + try: + conn.execute("BEGIN IMMEDIATE") + _raise_if_chat_message_thread_conflicts( + conn, + thread_id, + [m["id"] for m in messages], + ) + if prune_missing: + conn.execute("DELETE FROM chat_messages WHERE thread_id = ?", (thread_id,)) + conn.executemany( + """ + INSERT INTO chat_messages + (id, thread_id, parent_id, role, content_json, attachments_json, metadata_json, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(id) DO UPDATE SET + parent_id = excluded.parent_id, + role = excluded.role, + content_json = excluded.content_json, + attachments_json = excluded.attachments_json, + metadata_json = excluded.metadata_json, + created_at = excluded.created_at + WHERE excluded.thread_id = chat_messages.thread_id + """, + [ + ( + m["id"], + thread_id, + m.get("parentId"), + m["role"], + json.dumps(m.get("content", [])), + json.dumps(m.get("attachments")) + if m.get("attachments") is not None + else None, + json.dumps(m.get("metadata")) + if m.get("metadata") is not None + else None, + int(m["createdAt"]), + ) + for m in messages + ], + ) + conn.commit() + return list_chat_messages(thread_id) + except ChatMessageConflictError: + conn.rollback() + raise + except sqlite3.Error: + logger.exception("Failed to sync chat messages for thread %s", thread_id) + conn.rollback() + raise + finally: + conn.close() + + +def list_chat_messages(thread_id: str) -> list[dict]: + conn = get_connection() + try: + rows = conn.execute( + """ + SELECT * FROM chat_messages + WHERE thread_id = ? + ORDER BY created_at ASC, id ASC + """, + (thread_id,), + ).fetchall() + return [_chat_message_from_row(row) for row in rows] + finally: + conn.close() + + +def get_chat_message(thread_id: str, message_id: str) -> Optional[dict]: + conn = get_connection() + try: + row = conn.execute( + """ + SELECT * FROM chat_messages + WHERE thread_id = ? AND id = ? + """, + (thread_id, message_id), + ).fetchone() + return _chat_message_from_row(row) if row is not None else None + finally: + conn.close() + + +def list_chat_messages_for_threads(thread_ids: list[str]) -> list[dict]: + if not thread_ids: + return [] + unique_thread_ids = list(dict.fromkeys(thread_ids)) + messages: list[dict] = [] + conn = get_connection() + try: + for start in range(0, len(unique_thread_ids), _SQLITE_IN_CHUNK_SIZE): + chunk = unique_thread_ids[start : start + _SQLITE_IN_CHUNK_SIZE] + placeholders = ",".join("?" for _ in chunk) + rows = conn.execute( + f""" + SELECT * FROM chat_messages + WHERE thread_id IN ({placeholders}) + ORDER BY created_at ASC, id ASC + """, + chunk, + ).fetchall() + messages.extend(_chat_message_from_row(row) for row in rows) + return sorted( + messages, + key = lambda message: (message["createdAt"], message["id"]), + ) + finally: + conn.close() + + +def list_chat_settings() -> dict[str, Any]: + conn = get_connection() + try: + rows = conn.execute( + "SELECT key, value_json FROM chat_settings ORDER BY key" + ).fetchall() + settings: dict[str, Any] = {} + for row in rows: + settings[row["key"]] = _json_loads(row["value_json"], None) + return settings + finally: + conn.close() + + +def upsert_chat_settings(settings: dict[str, Any]) -> dict[str, Any]: + if not settings: + return list_chat_settings() + conn = get_connection() + try: + now = datetime.now(timezone.utc).isoformat() + conn.executemany( + """ + INSERT INTO chat_settings (key, value_json, updated_at) + VALUES (?, ?, ?) + ON CONFLICT(key) DO UPDATE SET + value_json = excluded.value_json, + updated_at = excluded.updated_at + """, + [(key, json.dumps(value), now) for key, value in settings.items()], + ) + conn.commit() + return list_chat_settings() + finally: + conn.close() + + +def _deep_merge_settings( + current: dict[str, Any], updates: dict[str, Any] +) -> dict[str, Any]: + merged = dict(current) + for key, value in updates.items(): + current_value = merged.get(key) + if isinstance(current_value, dict) and isinstance(value, dict): + merged[key] = _deep_merge_settings(current_value, value) + else: + merged[key] = value + return merged + + +def upsert_chat_settings_merge(updates: dict[str, Any]) -> dict[str, Any]: + """Atomic read-merge-write under BEGIN IMMEDIATE so two concurrent writers + cannot drop one another's updates.""" + if not updates: + return list_chat_settings() + conn = get_connection() + try: + conn.execute("BEGIN IMMEDIATE") + current, corrupt = _load_chat_settings_for_merge(conn) + unsafe_partial_keys = [ + key + for key, value in updates.items() + if key in corrupt and isinstance(value, dict) + ] + if unsafe_partial_keys: + conn.commit() + keys = ", ".join(sorted(unsafe_partial_keys)) + raise CorruptSettingsError( + f"Cannot apply partial settings patch to corrupt key(s): {keys}" + ) + merged = _deep_merge_settings(current, updates) + now = datetime.now(timezone.utc).isoformat() + conn.executemany( + """ + INSERT INTO chat_settings (key, value_json, updated_at) + VALUES (?, ?, ?) + ON CONFLICT(key) DO UPDATE SET + value_json = excluded.value_json, + updated_at = excluded.updated_at + """, + [(key, json.dumps(value), now) for key, value in merged.items()], + ) + conn.commit() + return merged + except CorruptSettingsError: + raise + except Exception: + conn.rollback() + raise + finally: + conn.close() + + +# --------------------------------------------------------------------------- +# Legacy Dexie import ledger +# --------------------------------------------------------------------------- +# See the schema comment in _ensure_schema() for the recovery rationale. + + +def list_chat_legacy_imports() -> list[str]: + """Return the legacy_thread_id of every thread already imported. + + Cheap: scans a single small PK-only table. The frontend stuffs the + result into a Set before walking Dexie, so the diff is O(|Dexie|). + """ + conn = get_connection() + try: + rows = conn.execute( + "SELECT legacy_thread_id FROM chat_legacy_imports" + ).fetchall() + return [row[0] for row in rows] + finally: + conn.close() + + +def upsert_chat_legacy_imports(legacy_thread_ids: list[str]) -> tuple[int, int]: + """Mark each given legacy thread id as imported. Idempotent. + + Returns (accepted, inserted): + - accepted: number of non-empty deduped input ids + - inserted: number of rows that were actually new (not already in ledger) + + ON CONFLICT DO NOTHING keeps the existing imported_at when an id is + recorded twice. INSERT...RETURNING reports only the rows that were + actually inserted, so callers can distinguish first-time imports + from idempotent re-runs without an extra SELECT. + """ + ids = list(dict.fromkeys(tid for tid in legacy_thread_ids if tid)) + if not ids: + return 0, 0 + ts = int(datetime.now(timezone.utc).timestamp() * 1000) + conn = get_connection() + try: + inserted = 0 + for tid in ids: + row = conn.execute( + """ + INSERT INTO chat_legacy_imports (legacy_thread_id, imported_at) + VALUES (?, ?) + ON CONFLICT(legacy_thread_id) DO NOTHING + RETURNING legacy_thread_id + """, + (tid, ts), + ).fetchone() + if row is not None: + inserted += 1 + conn.commit() + return len(ids), inserted + finally: + conn.close() diff --git a/studio/backend/tests/test_chat_history_routes.py b/studio/backend/tests/test_chat_history_routes.py new file mode 100644 index 0000000000..9337544638 --- /dev/null +++ b/studio/backend/tests/test_chat_history_routes.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import asyncio +import os +import sys + +import pytest +from fastapi import HTTPException + +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +from routes import chat_history + + +def _message(message_id: str, thread_id: str) -> chat_history.ChatMessage: + return chat_history.ChatMessage( + id = message_id, + threadId = thread_id, + parentId = None, + role = "user", + content = [{"type": "text", "text": "hello"}], + createdAt = 1_700_000_000_000, + ) + + +def test_replace_thread_messages_rejects_body_thread_mismatch(monkeypatch): + called = False + + def fake_get_chat_thread(thread_id: str): + return {"id": thread_id} + + def fake_sync_chat_messages(*args, **kwargs): + nonlocal called + called = True + return [] + + monkeypatch.setattr(chat_history, "get_chat_thread", fake_get_chat_thread) + monkeypatch.setattr(chat_history, "sync_chat_messages", fake_sync_chat_messages) + + with pytest.raises(HTTPException) as exc_info: + asyncio.run( + chat_history.replace_thread_messages( + "thread-1", + chat_history.ChatMessageSyncRequest( + messages = [_message("msg-1", "thread-2")], + pruneMissing = True, + ), + current_subject = "test-user", + ) + ) + + assert exc_info.value.status_code == 400 + assert "Message threadId mismatch" in str(exc_info.value.detail) + assert called is False + + +# --------------------------------------------------------------------------- +# /api/chat/import-ledger +# --------------------------------------------------------------------------- + + +def test_get_import_ledger_round_trips_through_storage(monkeypatch): + seen: list[str] = [] + + def fake_list(): + return list(seen) + + monkeypatch.setattr(chat_history, "list_chat_legacy_imports", fake_list) + + response = asyncio.run(chat_history.get_import_ledger(current_subject = "test-user")) + assert response.threadIds == [] + + seen.extend(["legacy-a", "legacy-b"]) + response = asyncio.run(chat_history.get_import_ledger(current_subject = "test-user")) + assert response.threadIds == ["legacy-a", "legacy-b"] + + +def test_record_import_ledger_returns_accepted_and_inserted(monkeypatch): + captured: list[list[str]] = [] + + def fake_upsert(thread_ids): + captured.append(list(thread_ids)) + # Pretend two of the three were already in the ledger. + return (len(thread_ids), max(0, len(thread_ids) - 2)) + + monkeypatch.setattr(chat_history, "upsert_chat_legacy_imports", fake_upsert) + + response = asyncio.run( + chat_history.record_import_ledger( + payload = chat_history.ChatImportLedgerRecordRequest( + threadIds = ["a", "b", "c"], + ), + current_subject = "test-user", + ) + ) + assert response.accepted == 3 + assert response.inserted == 1 + assert captured == [["a", "b", "c"]] + + +def test_record_import_ledger_rejects_oversize_payload(): + from pydantic import ValidationError + + with pytest.raises(ValidationError): + chat_history.ChatImportLedgerRecordRequest( + threadIds = [f"id-{i}" for i in range(10_001)], + ) diff --git a/studio/backend/tests/test_chat_history_storage.py b/studio/backend/tests/test_chat_history_storage.py new file mode 100644 index 0000000000..123dbf1b96 --- /dev/null +++ b/studio/backend/tests/test_chat_history_storage.py @@ -0,0 +1,308 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import threading + +import pytest + +from storage import studio_db + + +def _reset_studio_db(tmp_path, monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + +def _thread(thread_id: str = "thread-1") -> dict: + return { + "id": thread_id, + "title": "Test Chat", + "modelType": "base", + "modelId": "test-model", + "pairId": None, + "archived": False, + "createdAt": 1_700_000_000_000, + } + + +def _message( + message_id: str, + created_at: int, + content: str, + thread_id: str = "thread-1", +) -> dict: + return { + "id": message_id, + "threadId": thread_id, + "parentId": None, + "role": "user", + "content": [{"type": "text", "text": content}], + "createdAt": created_at, + } + + +def test_sync_chat_messages_upserts_without_pruning(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_thread(_thread()) + studio_db.sync_chat_messages( + "thread-1", + [ + _message("msg-1", 1, "keep me"), + _message("msg-2", 2, "old text"), + ], + prune_missing = True, + ) + + messages = studio_db.sync_chat_messages( + "thread-1", + [_message("msg-2", 2, "updated text")], + ) + + by_id = {message["id"]: message for message in messages} + assert set(by_id) == {"msg-1", "msg-2"} + assert by_id["msg-2"]["content"] == [{"type": "text", "text": "updated text"}] + + +def test_sync_chat_messages_prunes_when_requested(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_thread(_thread()) + studio_db.sync_chat_messages( + "thread-1", + [ + _message("msg-1", 1, "delete me"), + _message("msg-2", 2, "keep me"), + ], + ) + + messages = studio_db.sync_chat_messages( + "thread-1", + [_message("msg-2", 2, "keep me")], + prune_missing = True, + ) + + assert [message["id"] for message in messages] == ["msg-2"] + + +def test_upsert_chat_message_rejects_cross_thread_id_conflict(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_thread(_thread("thread-1")) + studio_db.upsert_chat_thread(_thread("thread-2")) + studio_db.upsert_chat_message(_message("msg-1", 1, "original", "thread-1")) + + with pytest.raises(studio_db.ChatMessageConflictError): + studio_db.upsert_chat_message(_message("msg-1", 2, "moved", "thread-2")) + + assert [m["id"] for m in studio_db.list_chat_messages("thread-1")] == ["msg-1"] + assert studio_db.list_chat_messages("thread-2") == [] + + +def test_sync_chat_messages_detects_conflict_before_prune(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_thread(_thread("thread-1")) + studio_db.upsert_chat_thread(_thread("thread-2")) + studio_db.sync_chat_messages( + "thread-1", + [_message("keep-me", 1, "keep", "thread-1")], + ) + studio_db.upsert_chat_message(_message("conflict", 2, "other", "thread-2")) + + with pytest.raises(studio_db.ChatMessageConflictError): + studio_db.sync_chat_messages( + "thread-1", + [_message("conflict", 3, "bad", "thread-1")], + prune_missing = True, + ) + + assert [m["id"] for m in studio_db.list_chat_messages("thread-1")] == ["keep-me"] + assert [m["id"] for m in studio_db.list_chat_messages("thread-2")] == ["conflict"] + + +def test_settings_merge_atomic_under_concurrency(tmp_path, monkeypatch): + """Two threads writing distinct keys must not drop each other's update.""" + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_settings_merge({"inferenceParams": {}}) + + barrier = threading.Barrier(2) + + def writer(key: str, value: float) -> None: + barrier.wait() + studio_db.upsert_chat_settings_merge({"inferenceParams": {key: value}}) + + t1 = threading.Thread(target = writer, args = ("temperature", 0.7)) + t2 = threading.Thread(target = writer, args = ("topP", 0.9)) + t1.start() + t2.start() + t1.join() + t2.join() + + merged = studio_db.list_chat_settings()["inferenceParams"] + assert merged.get("temperature") == 0.7 + assert merged.get("topP") == 0.9 + + +def test_settings_merge_preserves_nested_keys(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_settings_merge( + {"inferenceParams": {"temperature": 0.5, "topP": 0.8}} + ) + studio_db.upsert_chat_settings_merge({"inferenceParams": {"temperature": 0.9}}) + + params = studio_db.list_chat_settings()["inferenceParams"] + assert params == {"temperature": 0.9, "topP": 0.8} + + +def test_settings_merge_quarantines_corrupt_json_and_rejects_partial_patch( + tmp_path, + monkeypatch, +): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_settings_merge( + {"inferenceParams": {"temperature": 0.5, "topP": 0.8}} + ) + conn = studio_db.get_connection() + try: + conn.execute( + "UPDATE chat_settings SET value_json = ? WHERE key = ?", + ('{"temperature": 0.5', "inferenceParams"), + ) + conn.commit() + finally: + conn.close() + + with pytest.raises(studio_db.CorruptSettingsError): + studio_db.upsert_chat_settings_merge({"inferenceParams": {"temperature": 0.9}}) + + conn = studio_db.get_connection() + try: + quarantined = conn.execute( + "SELECT key, value_json, reason FROM chat_settings_quarantine" + ).fetchall() + remaining = conn.execute( + "SELECT key FROM chat_settings WHERE key = ?", + ("inferenceParams",), + ).fetchall() + finally: + conn.close() + assert [row["key"] for row in quarantined] == ["inferenceParams"] + assert quarantined[0]["reason"] == "json_decode_error" + assert remaining == [] + + +def test_settings_merge_replaces_corrupt_scalar_after_quarantine(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_settings_merge({"autoTitle": False}) + conn = studio_db.get_connection() + try: + conn.execute( + "UPDATE chat_settings SET value_json = ? WHERE key = ?", + ("not-json", "autoTitle"), + ) + conn.commit() + finally: + conn.close() + + settings = studio_db.upsert_chat_settings_merge({"autoTitle": True}) + + assert settings["autoTitle"] is True + conn = studio_db.get_connection() + try: + quarantined = conn.execute( + "SELECT key, reason FROM chat_settings_quarantine" + ).fetchall() + finally: + conn.close() + assert [(row["key"], row["reason"]) for row in quarantined] == [ + ("autoTitle", "json_decode_error") + ] + + +def test_list_chat_messages_for_threads_chunks_over_900_ids(tmp_path, monkeypatch): + """SQLite host-parameter limit is 999 on older builds; chunk at 900.""" + _reset_studio_db(tmp_path, monkeypatch) + n = 901 + for i in range(n): + studio_db.upsert_chat_thread( + { + "id": f"t-{i}", + "title": "T", + "modelType": "base", + "modelId": "m", + "pairId": None, + "archived": False, + "createdAt": 1_700_000_000_000 + i, + } + ) + studio_db.upsert_chat_message( + { + "id": f"m-{i}", + "threadId": f"t-{i}", + "parentId": None, + "role": "user", + "content": [{"type": "text", "text": "hi"}], + "createdAt": 1_700_000_000_000 + i, + } + ) + out = studio_db.list_chat_messages_for_threads([f"t-{i}" for i in range(n)]) + assert len(out) == n + assert {m["threadId"] for m in out} == {f"t-{i}" for i in range(n)} + + +# --------------------------------------------------------------------------- +# Legacy Dexie import ledger +# --------------------------------------------------------------------------- + + +def test_legacy_imports_empty_by_default(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + assert studio_db.list_chat_legacy_imports() == [] + + +def test_legacy_imports_records_and_lists(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + accepted, inserted = studio_db.upsert_chat_legacy_imports( + ["legacy-a", "legacy-b", "legacy-c"], + ) + assert accepted == 3 + assert inserted == 3 + assert set(studio_db.list_chat_legacy_imports()) == { + "legacy-a", + "legacy-b", + "legacy-c", + } + + +def test_legacy_imports_is_idempotent(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + accepted1, inserted1 = studio_db.upsert_chat_legacy_imports( + ["legacy-a", "legacy-b"], + ) + accepted2, inserted2 = studio_db.upsert_chat_legacy_imports( + ["legacy-b", "legacy-c"], + ) + assert (accepted1, inserted1) == (2, 2) + # legacy-b is already in the ledger, only legacy-c is genuinely new. + assert (accepted2, inserted2) == (2, 1) + assert set(studio_db.list_chat_legacy_imports()) == { + "legacy-a", + "legacy-b", + "legacy-c", + } + + +def test_legacy_imports_dedups_input(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + accepted, inserted = studio_db.upsert_chat_legacy_imports( + ["x", "x", "y", "x"], + ) + # accepted is the deduped non-empty input size; inserted is the rows + # actually new in the ledger after ON CONFLICT DO NOTHING. + assert accepted == 2 + assert inserted == 2 + assert set(studio_db.list_chat_legacy_imports()) == {"x", "y"} + + +def test_legacy_imports_ignores_empty(tmp_path, monkeypatch): + _reset_studio_db(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() == [] diff --git a/studio/backend/tests/test_desktop_auth.py b/studio/backend/tests/test_desktop_auth.py index 913c3cc355..ab1a03eeda 100644 --- a/studio/backend/tests/test_desktop_auth.py +++ b/studio/backend/tests/test_desktop_auth.py @@ -431,6 +431,7 @@ def test_desktop_capabilities_json_reports_rollout_safe_flags(): def test_health_response_reports_desktop_capability_fields(monkeypatch): router_stub = SimpleNamespace( auth_router = APIRouter(), + chat_history_router = APIRouter(), data_recipe_router = APIRouter(), datasets_router = APIRouter(), export_router = APIRouter(), diff --git a/studio/frontend/src/features/auth/api.ts b/studio/frontend/src/features/auth/api.ts index 1c0909454d..95296d3ab9 100644 --- a/studio/frontend/src/features/auth/api.ts +++ b/studio/frontend/src/features/auth/api.ts @@ -18,6 +18,9 @@ type RefreshResponse = { }; let isRedirecting = false; +let refreshInflight: Promise | null = null; +let refreshInflightToken: string | null = null; +let logoutGeneration = 0; const TAURI_FETCH_RETRY_DELAYS_MS = [250, 750, 1500] as const; @@ -25,6 +28,10 @@ function wait(ms: number): Promise { return new Promise((resolve) => setTimeout(resolve, ms)); } +function clearAuthTokensIfCurrent(refreshToken: string | null): void { + if (!refreshToken || getRefreshToken() === refreshToken) clearAuthTokens(); +} + async function fetchWithTauriNetworkRetry( input: RequestInfo | URL, init?: RequestInit, @@ -98,19 +105,15 @@ async function retryWithTauriAutoAuth( return null; } -// Singleflight: the backend consumes the refresh token atomically, so -// concurrent callers must share one in-flight promise (loser would 401). -let refreshInflight: Promise | null = null; -// Bumped by logout(); a refresh that resolves after logout drops its -// new tokens instead of silently re-auth-ing the SPA. -let logoutGeneration = 0; - export async function refreshSession(): Promise { - if (refreshInflight) return refreshInflight; + const refreshToken = getRefreshToken(); + if (!refreshToken) return false; + if (refreshInflight && refreshInflightToken === refreshToken) { + return refreshInflight; + } + const startGeneration = logoutGeneration; - refreshInflight = (async () => { - const refreshToken = getRefreshToken(); - if (!refreshToken) return false; + const promise = (async () => { try { const response = await fetchWithTauriNetworkRetry( apiUrl("/api/auth/refresh"), @@ -121,11 +124,12 @@ export async function refreshSession(): Promise { }, ); if (!response.ok) { - clearAuthTokens(); + clearAuthTokensIfCurrent(refreshToken); return false; } const payload = (await response.json()) as RefreshResponse; if (startGeneration !== logoutGeneration) return false; + if (getRefreshToken() !== refreshToken) return false; storeAuthTokens(payload.access_token, payload.refresh_token); setMustChangePassword(payload.must_change_password ?? false); return true; @@ -133,10 +137,15 @@ export async function refreshSession(): Promise { return false; } })(); + refreshInflight = promise; + refreshInflightToken = refreshToken; try { - return await refreshInflight; + return await promise; } finally { - refreshInflight = null; + if (refreshInflight === promise) { + refreshInflight = null; + refreshInflightToken = null; + } } } @@ -181,12 +190,13 @@ export async function authFetch( } if (response.status !== 401) return response; + const refreshToken = getRefreshToken(); const refreshed = await refreshSession(); if (!refreshed) { if (isTauri) { return (await retryWithTauriAutoAuth(resolvedInput, init)) ?? response; } - clearAuthTokens(); + clearAuthTokensIfCurrent(refreshToken); void redirectToAuth(); return response; } diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 7814f1892d..a5a77b0863 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -1,34 +1,11 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import type { ChatModelAdapter } from "@assistant-ui/react"; -import type { MessageTiming, ToolCallMessagePart } from "@assistant-ui/core"; -import { toast } from "@/lib/toast"; import { getAuthToken } from "@/features/auth/session"; import { apiUrl } from "@/lib/api-base"; -import { - generateAudio, - listCachedGguf, - listCachedModels, - listGgufVariants, - loadModel, - streamChatCompletions, - validateModel, -} from "./chat-api"; -import { pickFriendlyContainerName } from "../lib/friendly-names"; -import { - createOpenAIContainer, - listOpenAIContainers, -} from "./openai-containers"; -import { - encryptProviderApiKey, - isProviderKeyRotationError, -} from "./providers-api"; -import { db } from "../db"; -import type { - OpenAIChatCompletionsRequest, - OpenAIMessageContent, -} from "../types/api"; +import { toast } from "@/lib/toast"; +import type { MessageTiming, ToolCallMessagePart } from "@assistant-ui/core"; +import type { ChatModelAdapter } from "@assistant-ui/react"; import { getExternalProviderApiKey, isCustomProviderType, @@ -38,6 +15,7 @@ import { supportsProviderPromptCaching, toExternalBackendProviderType, } from "../external-providers"; +import { pickFriendlyContainerName } from "../lib/friendly-names"; import { EXTERNAL_MAX_OUTPUT_TOKENS, clampReasoningEffortToLevels, @@ -51,12 +29,38 @@ import { import { useChatRuntimeStore } from "../stores/chat-runtime-store"; import { useExternalProvidersStore } from "../stores/external-providers-store"; import { isMultimodalResponse } from "../types/api"; +import type { + OpenAIChatCompletionsRequest, + OpenAIMessageContent, +} from "../types/api"; import type { ChatModelSummary } from "../types/runtime"; import { getImageInputUnavailableReason } from "../utils/image-input-support"; +import { + getStoredChatThread, + listStoredChatThreads, + updateStoredChatThread, +} from "../utils/chat-history-storage"; import { hasClosedThinkTag, parseAssistantContent, } from "../utils/parse-assistant-content"; +import { + generateAudio, + listCachedGguf, + listCachedModels, + listGgufVariants, + loadModel, + streamChatCompletions, + validateModel, +} from "./chat-api"; +import { + createOpenAIContainer, + listOpenAIContainers, +} from "./openai-containers"; +import { + encryptProviderApiKey, + isProviderKeyRotationError, +} from "./providers-api"; /** Server-side usage data from llama-server (via stream_options.include_usage). */ interface ServerUsage { @@ -119,11 +123,42 @@ export function isContextLimitError(message: string): boolean { ); } +function wait(ms: number): Promise { + return new Promise((resolve) => setTimeout(resolve, ms)); +} + +async function updateStoredChatThreadEventually( + threadId: string, + patch: Parameters[1], +): Promise { + for (let attempt = 0; attempt < 10; attempt++) { + const updated = await updateStoredChatThread(threadId, patch).catch( + () => undefined, + ); + if (updated) return; + await wait(50); + } +} + /** Parse "Title: ...\nURL: ...\nSnippet: ..." blocks into source content parts. */ -function parseSourcesFromResult(raw: string): { type: "source"; sourceType: "url"; id: string; url: string; title: string; metadata?: { description: string } }[] { +function parseSourcesFromResult(raw: string): { + type: "source"; + sourceType: "url"; + id: string; + url: string; + title: string; + metadata?: { description: string }; +}[] { if (!raw) return []; const blocks = raw.split(/\n---\n/).filter(Boolean); - const sources: { type: "source"; sourceType: "url"; id: string; url: string; title: string; metadata?: { description: string } }[] = []; + const sources: { + type: "source"; + sourceType: "url"; + id: string; + url: string; + title: string; + metadata?: { description: string }; + }[] = []; for (const block of blocks) { const titleMatch = block.match(/Title:\s*(.+)/); const urlMatch = block.match(/URL:\s*(.+)/); @@ -264,7 +299,7 @@ function collectImageParts( message: RunMessage, ): Array<{ type: "image_url"; image_url: { url: string } }> { const parts: Array<{ type: "image_url"; image_url: { url: string } }> = []; - + for (const part of message.content ?? []) { if (part.type === "image" && "image" in part) { const src = (part as { image: string }).image; @@ -278,7 +313,7 @@ function collectImageParts( } } } - + if ("attachments" in message && (message.attachments?.length ?? 0) > 0) { for (const attachment of message.attachments ?? []) { for (const part of attachment.content ?? []) { @@ -298,7 +333,7 @@ function collectImageParts( } } } - + return parts; } @@ -388,7 +423,12 @@ function findLatestUserAudioBase64(messages: RunMessages): string | undefined { for (const part of message.content ?? []) { if (part.type === "audio" && "audio" in part) { - const audioPart = (part as unknown as { type: "audio"; audio: string | { data: string; format: string } }).audio; + const audioPart = ( + part as unknown as { + type: "audio"; + audio: string | { data: string; format: string }; + } + ).audio; const raw = typeof audioPart === "string" ? audioPart : audioPart?.data; if (raw) return raw.startsWith("data:") ? raw.split(",")[1] : raw; } @@ -407,7 +447,7 @@ async function resolveUseAdapter( return undefined; } try { - const thread = await db.threads.get(threadId); + const thread = await getStoredChatThread(threadId); if (!thread?.pairId) { return undefined; } @@ -426,8 +466,14 @@ async function resolveUseAdapter( function waitForModelReady(abortSignal?: AbortSignal): Promise { return new Promise((resolve, reject) => { const check = () => { - if (abortSignal?.aborted) { reject(new Error("Aborted")); return; } - if (!useChatRuntimeStore.getState().modelLoading) { resolve(); return; } + if (abortSignal?.aborted) { + reject(new Error("Aborted")); + return; + } + if (!useChatRuntimeStore.getState().modelLoading) { + resolve(); + return; + } setTimeout(check, 500); }; check(); @@ -515,12 +561,17 @@ async function autoLoadSmallestModel(): Promise<{ gguf_variant: variant.quant, trust_remote_code: trustRemoteCode, }); - useChatRuntimeStore.getState().setCheckpoint(repo.repo_id, variant.quant); + useChatRuntimeStore + .getState() + .setCheckpoint(repo.repo_id, variant.quant); const store = useChatRuntimeStore.getState(); store.setModelRequiresTrustRemoteCode( loadResp.requires_trust_remote_code ?? false, ); - store.setParams({ ...store.params, maxTokens: loadResp.context_length ?? 131072 }); + store.setParams({ + ...store.params, + maxTokens: loadResp.context_length ?? 131072, + }); // Add model to store so the selector shows the name const autoModel: ChatModelSummary = { id: repo.repo_id, @@ -538,12 +589,16 @@ async function autoLoadSmallestModel(): Promise<{ } useChatRuntimeStore.setState({ ggufContextLength: loadResp.context_length ?? 131072, - ggufMaxContextLength: loadResp.max_context_length ?? loadResp.context_length ?? 131072, + ggufMaxContextLength: + loadResp.max_context_length ?? + loadResp.context_length ?? + 131072, supportsReasoning: loadResp.supports_reasoning ?? false, reasoningAlwaysOn: loadResp.reasoning_always_on ?? false, reasoningEnabled: loadResp.supports_reasoning ?? false, reasoningStyle: loadResp.reasoning_style ?? "enable_thinking", - supportsPreserveThinking: loadResp.supports_preserve_thinking ?? false, + supportsPreserveThinking: + loadResp.supports_preserve_thinking ?? false, supportsTools: loadResp.supports_tools ?? false, toolsEnabled: loadResp.supports_tools ?? false, codeToolsEnabled: loadResp.supports_tools ?? false, @@ -554,7 +609,9 @@ async function autoLoadSmallestModel(): Promise<{ loadedChatTemplateOverride: null, loadedIsMultimodal: isMultimodalResponse(loadResp), }); - toast.success(`Loaded ${repo.repo_id} (${variant.quant})`, { id: toastId }); + toast.success(`Loaded ${repo.repo_id} (${variant.quant})`, { + id: toastId, + }); return { loaded: true, blockedByTrustRemoteCode: false }; } } catch { @@ -566,7 +623,9 @@ async function autoLoadSmallestModel(): Promise<{ // Fall back to safetensors models if (modelRepos.length > 0) { - const sorted = [...modelRepos].sort((a, b) => a.size_bytes - b.size_bytes); + const sorted = [...modelRepos].sort( + (a, b) => a.size_bytes - b.size_bytes, + ); for (const repo of sorted) { if (loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS) break; try { @@ -601,7 +660,8 @@ async function autoLoadSmallestModel(): Promise<{ reasoningAlwaysOn: sfLoadResp.reasoning_always_on ?? false, reasoningEnabled: sfLoadResp.supports_reasoning ?? false, reasoningStyle: sfLoadResp.reasoning_style ?? "enable_thinking", - supportsPreserveThinking: sfLoadResp.supports_preserve_thinking ?? false, + supportsPreserveThinking: + sfLoadResp.supports_preserve_thinking ?? false, supportsTools: sfLoadResp.supports_tools ?? false, // Parity with the GGUF branch above. toolsEnabled: sfLoadResp.supports_tools ?? false, @@ -646,7 +706,8 @@ async function autoLoadSmallestModel(): Promise<{ // No cached models found — try downloading a small default GGUF toast("Downloading a small model…", { id: toastId, - description: "No downloaded models found. Fetching Gemma-4-E2B-it (UD-Q4_K_XL).", + description: + "No downloaded models found. Fetching Gemma-4-E2B-it (UD-Q4_K_XL).", duration: 30000, }); try { @@ -671,12 +732,17 @@ async function autoLoadSmallestModel(): Promise<{ gguf_variant: "UD-Q4_K_XL", trust_remote_code: trustRemoteCode, }); - useChatRuntimeStore.getState().setCheckpoint("unsloth/gemma-4-E2B-it-GGUF", "UD-Q4_K_XL"); + useChatRuntimeStore + .getState() + .setCheckpoint("unsloth/gemma-4-E2B-it-GGUF", "UD-Q4_K_XL"); const store = useChatRuntimeStore.getState(); store.setModelRequiresTrustRemoteCode( loadResp.requires_trust_remote_code ?? false, ); - store.setParams({ ...store.params, maxTokens: loadResp.context_length ?? 131072 }); + store.setParams({ + ...store.params, + maxTokens: loadResp.context_length ?? 131072, + }); const defaultModel: ChatModelSummary = { id: "unsloth/gemma-4-E2B-it-GGUF", name: loadResp.display_name ?? "gemma-4-E2B-it-GGUF", @@ -689,7 +755,8 @@ async function autoLoadSmallestModel(): Promise<{ } useChatRuntimeStore.setState({ ggufContextLength: loadResp.context_length ?? 131072, - ggufMaxContextLength: loadResp.max_context_length ?? loadResp.context_length ?? 131072, + ggufMaxContextLength: + loadResp.max_context_length ?? loadResp.context_length ?? 131072, supportsReasoning: loadResp.supports_reasoning ?? false, reasoningAlwaysOn: loadResp.reasoning_always_on ?? false, reasoningEnabled: loadResp.supports_reasoning ?? false, @@ -720,8 +787,7 @@ async function autoLoadSmallestModel(): Promise<{ hadNonTrustFailure = true; return { loaded: false, - blockedByTrustRemoteCode: - blockedByTrustRemoteCode && !hadNonTrustFailure, + blockedByTrustRemoteCode: blockedByTrustRemoteCode && !hadNonTrustFailure, }; } } @@ -729,6 +795,7 @@ async function autoLoadSmallestModel(): Promise<{ export function createOpenAIStreamAdapter(): ChatModelAdapter { return { async *run({ messages, abortSignal, unstable_threadId }) { + await useChatRuntimeStore.getState().hydratePersistedSettings(); let runtime = useChatRuntimeStore.getState(); // Capture the thread ID once at the start so it stays stable even if // the user switches chats while waiting for model load / auto-load. @@ -763,11 +830,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { // Re-read store after potential auto-load / model ready wait runtime = useChatRuntimeStore.getState(); const { params } = runtime; - const { - supportsTools, - toolsEnabled, - codeToolsEnabled, - } = runtime; + const { supportsTools, toolsEnabled, codeToolsEnabled } = runtime; const externalSelection = parseExternalModelId(params.checkpoint); const isExternalRequest = externalSelection !== null; if ( @@ -857,7 +920,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { if (audioBase64) { const audioName = runtime.pendingAudioName; if (audioName) { - const lastUserMsg = [...messages].reverse().find((m) => m.role === "user"); + const lastUserMsg = [...messages] + .reverse() + .find((m) => m.role === "user"); if (lastUserMsg) sentAudioNames.set(lastUserMsg.id, audioName); } runtime.clearPendingAudio(); @@ -905,8 +970,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } catch (err) { if (!abortSignal.aborted) { toast.error("Audio generation failed", { - description: - err instanceof Error ? err.message : "Unknown error", + description: err instanceof Error ? err.message : "Unknown error", }); } throw err; @@ -967,7 +1031,10 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { // Tool call content parts — accumulated and yielded cumulatively. // result is set directly on the tool-call part when tool_end arrives. const toolCallParts: ToolCallMessagePart[] = []; - let serverMetadata: { usage?: ServerUsage; timings?: ServerTimings } | null = null; + let serverMetadata: { + usage?: ServerUsage; + timings?: ServerTimings; + } | null = null; // Per-run cancellation token so a delayed stop POST cannot match // the next run on the same thread. @@ -1042,16 +1109,17 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { NonNullable, "none" | "minimal" | "low" | "medium" | "high" | "max" | "xhigh" >; - const fallbackExternalEffort = - (externalReasoningCaps.reasoningEffortLevels[0] ?? - "low") as RequestReasoningEffort; + const fallbackExternalEffort = (externalReasoningCaps + .reasoningEffortLevels[0] ?? "low") as RequestReasoningEffort; const selectedExternalEffort: RequestReasoningEffort = clampReasoningEffortToLevels( reasoningEffort, externalReasoningCaps.reasoningEffortLevels, ) as RequestReasoningEffort; const localReasoningEffort = - reasoningEffort === "low" || reasoningEffort === "medium" || reasoningEffort === "high" + reasoningEffort === "low" || + reasoningEffort === "medium" || + reasoningEffort === "high" ? reasoningEffort : "low"; const externalReasoningEnabled = @@ -1078,7 +1146,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { ); if (codeExecEnabledForThisTurn && resolvedThreadId) { try { - const thread = await db.threads.get(resolvedThreadId); + const thread = await getStoredChatThread(resolvedThreadId); openaiCodeExecContainerId = thread?.openaiCodeExecContainerId ?? null; anthropicCodeExecContainerId = @@ -1114,11 +1182,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { openaiCodeExecContainerId && !activeContainerIds.has(openaiCodeExecContainerId) ) { - void db.threads - .update(resolvedThreadId, { - openaiCodeExecContainerId: null, - }) - .catch(() => {}); + void updateStoredChatThreadEventually(resolvedThreadId, { + openaiCodeExecContainerId: null, + }).catch(() => {}); openaiCodeExecContainerId = null; } } @@ -1136,32 +1202,28 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { externalProvider.providerType === "openai" ) { try { - const others = await db.threads - .orderBy("createdAt") - .reverse() - .toArray(); + const others = await listStoredChatThreads({ + includeArchived: true, + }); for (const t of others) { if (t.id === resolvedThreadId) continue; if (!t.openaiCodeExecContainerId) continue; - // Skip inherited ids that are not in the active - // container set — they would 400 on send. Also - // null them on the source thread so the next - // inheritance pass doesn't re-pick the same dead id. + // Skip ids not in active set; null on source thread so + // the next pass doesn't re-pick a dead id. if ( activeContainerIds && !activeContainerIds.has(t.openaiCodeExecContainerId) ) { - void db.threads - .update(t.id, { openaiCodeExecContainerId: null }) + void updateStoredChatThreadEventually(t.id, { + openaiCodeExecContainerId: null, + }) .catch(() => {}); continue; } openaiCodeExecContainerId = t.openaiCodeExecContainerId; - void db.threads - .update(resolvedThreadId, { - openaiCodeExecContainerId, - }) - .catch(() => {}); + void updateStoredChatThreadEventually(resolvedThreadId, { + openaiCodeExecContainerId, + }).catch(() => {}); break; } } catch { @@ -1180,8 +1242,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { externalProvider.providerType === "openai" ) { const ttl = externalProvider.openaiContainerTtlMinutes; - const ttlToUse = - typeof ttl === "number" && ttl >= 1 ? ttl : 20; + const ttlToUse = typeof ttl === "number" && ttl >= 1 ? ttl : 20; try { const created = await createOpenAIContainer( { @@ -1198,10 +1259,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { }, ); openaiCodeExecContainerId = created.id; - void db.threads - .update(resolvedThreadId, { - openaiCodeExecContainerId: created.id, - }) + void updateStoredChatThreadEventually(resolvedThreadId, { + openaiCodeExecContainerId: created.id, + }) .catch(() => {}); } catch { // Fall back to backend's container_auto path on @@ -1254,7 +1314,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { // schema — for Anthropic that's the entries appended to // body["tools"] inside _stream_anthropic. ...((toolsEnabled && - providerSupportsBuiltinWebSearch(externalProvider.providerType)) || + providerSupportsBuiltinWebSearch( + externalProvider.providerType, + )) || (codeToolsEnabled && providerSupportsBuiltinCodeExecution( externalProvider.providerType, @@ -1360,7 +1422,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { : {} : { enable_thinking: reasoningEnabled } : {}), - ...(supportsPreserveThinking ? { preserve_thinking: preserveThinking } : {}), + ...(supportsPreserveThinking + ? { preserve_thinking: preserveThinking } + : {}), ...(supportsTools && (toolsEnabled || codeToolsEnabled) ? { enable_tools: true, @@ -1368,8 +1432,10 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { ...(toolsEnabled ? ["web_search"] : []), ...(codeToolsEnabled ? ["python", "terminal"] : []), ], - auto_heal_tool_calls: useChatRuntimeStore.getState().autoHealToolCalls, - max_tool_calls_per_message: useChatRuntimeStore.getState().maxToolCallsPerMessage, + auto_heal_tool_calls: + useChatRuntimeStore.getState().autoHealToolCalls, + max_tool_calls_per_message: + useChatRuntimeStore.getState().maxToolCallsPerMessage, tool_call_timeout: (() => { const mins = useChatRuntimeStore.getState().toolCallTimeout; return mins >= 9999 ? 9999 : mins * 60; @@ -1389,16 +1455,20 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { for await (const chunk of stream) { // Handle tool status events - const toolStatusText = (chunk as unknown as { _toolStatus?: string })._toolStatus; + const toolStatusText = ( + chunk as unknown as { _toolStatus?: string } + )._toolStatus; if (toolStatusText !== undefined) { runtime.setToolStatus(toolStatusText || null); continue; } - + // Emit tool-call content parts for assistant-ui. // On tool_start: add a new tool-call part (renders in "running" state). // On tool_end: set result on the existing part (transitions to "complete"). - const toolEvent = (chunk as unknown as { _toolEvent?: Record })._toolEvent; + const toolEvent = ( + chunk as unknown as { _toolEvent?: Record } + )._toolEvent; if (toolEvent !== undefined) { // OpenAI shell-tool container persistence — see // ThreadRecord.openaiCodeExecContainerId. The backend @@ -1414,29 +1484,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { externalProvider?.providerType === "anthropic" ? "anthropicCodeExecContainerId" : "openaiCodeExecContainerId"; - // On the first turn of a brand-new thread the row - // may not be in Dexie yet when this SSE event - // fires — db.threads.update silently affects 0 - // rows, the next turn re-reads null, and Anthropic - // auto-creates a fresh container. Retry briefly so - // assistant-ui's own DexieAdapter.initialize lands - // the row first (with the correct modelType for - // base / lora / compare contexts) and our update - // sticks on a subsequent attempt. - try { - for (let attempt = 0; attempt < 10; attempt++) { - const affected = await db.threads.update( - resolvedThreadId, - { [field]: newContainerId }, - ); - if (affected > 0) break; - await new Promise((resolve) => - setTimeout(resolve, 50), - ); - } - } catch { - /* best-effort: container reuse is an optimization */ - } + void updateStoredChatThreadEventually(resolvedThreadId, { + [field]: newContainerId, + }).catch(() => {}); } continue; } @@ -1446,17 +1496,19 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { externalProvider?.providerType === "anthropic" ? "anthropicCodeExecContainerId" : "openaiCodeExecContainerId"; - void db.threads - .update(resolvedThreadId, { - [field]: null, - }) + void updateStoredChatThreadEventually(resolvedThreadId, { + [field]: null, + }) .catch(() => {}); } continue; } if (toolEvent.type === "tool_start") { - const id = (toolEvent.tool_call_id as string) || `${toolEvent.tool_name}_${Date.now()}`; - const toolArgs = (toolEvent.arguments ?? {}) as ToolCallMessagePart["args"]; + const id = + (toolEvent.tool_call_id as string) || + `${toolEvent.tool_name}_${Date.now()}`; + const toolArgs = (toolEvent.arguments ?? + {}) as ToolCallMessagePart["args"]; toolCallParts.push({ type: "tool-call" as const, toolCallId: id, @@ -1465,21 +1517,29 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { args: toolArgs, }); } else if (toolEvent.type === "tool_end") { - const id = (toolEvent.tool_call_id as string) || - toolCallParts[toolCallParts.length - 1]?.toolCallId || ""; - const idx = toolCallParts.findIndex((p) => p.toolCallId === id); + const id = + (toolEvent.tool_call_id as string) || + toolCallParts[toolCallParts.length - 1]?.toolCallId || + ""; + const idx = toolCallParts.findIndex( + (p) => p.toolCallId === id, + ); if (idx !== -1) { const rawResult = (toolEvent.result as string) ?? ""; const imgMarker = "\n__IMAGES__:"; const imgIdx = rawResult.lastIndexOf(imgMarker); - let parsedResult: string | { text: string; images: string[]; sessionId: string }; + let parsedResult: + | string + | { text: string; images: string[]; sessionId: string }; if (imgIdx !== -1) { const text = rawResult.slice(0, imgIdx); // Fall back to "_default" to match the backend sandbox directory // used when no session_id is provided (see tools.py _get_workdir). const sessionId = resolvedThreadId || "_default"; try { - const images = JSON.parse(rawResult.slice(imgIdx + imgMarker.length)) as string[]; + const images = JSON.parse( + rawResult.slice(imgIdx + imgMarker.length), + ) as string[]; parsedResult = { text, images, sessionId }; } catch { parsedResult = rawResult; @@ -1487,7 +1547,10 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } else { parsedResult = rawResult; } - toolCallParts[idx] = { ...toolCallParts[idx], result: parsedResult }; + toolCallParts[idx] = { + ...toolCallParts[idx], + result: parsedResult, + }; } } // Yield cumulative state so tool UI updates (tools first, text after) @@ -1495,7 +1558,11 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { yield { content: [...toolCallParts, ...textParts], metadata: { - timing: buildTiming(streamStartTime, totalChunks, firstTokenTime), + timing: buildTiming( + streamStartTime, + totalChunks, + firstTokenTime, + ), custom: { reasoningDuration }, }, }; @@ -1506,7 +1573,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { if (chunk.choices?.length === 0 && chunk.usage) { serverMetadata = { usage: chunk.usage, - timings: (chunk as Record).timings as ServerTimings | undefined, + timings: (chunk as Record).timings as + | ServerTimings + | undefined, }; continue; } @@ -1617,11 +1686,20 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } const parts = parseAssistantContent(cumulativeText); - if (parts.some((part) => part.type === "reasoning") && !reasoningStartAt) { + if ( + parts.some((part) => part.type === "reasoning") && + !reasoningStartAt + ) { reasoningStartAt = Date.now(); } - if (hasClosedThinkTag(cumulativeText) && reasoningStartAt && !reasoningDuration) { - reasoningDuration = Math.round((Date.now() - reasoningStartAt) / 1000); + if ( + hasClosedThinkTag(cumulativeText) && + reasoningStartAt && + !reasoningDuration + ) { + reasoningDuration = Math.round( + (Date.now() - reasoningStartAt) / 1000, + ); } if (parts.length > 0 || toolCallParts.length > 0) { @@ -1671,12 +1749,14 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { ) { return []; } - return parseSourcesFromResult(typeof tc.result === "string" ? tc.result : ""); + return parseSourcesFromResult( + typeof tc.result === "string" ? tc.result : "", + ); }); const meta = serverMetadata; - const finalTokenCount = meta?.usage?.completion_tokens - ?? estimateTokenCount(cumulativeText); + const finalTokenCount = + meta?.usage?.completion_tokens ?? estimateTokenCount(cumulativeText); const finalTokPerSec = meta?.timings?.predicted_per_second; const serverPromptEvalTime = meta?.timings?.prompt_ms; @@ -1716,19 +1796,23 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { custom: { reasoningDuration, serverTimings: meta?.timings ?? undefined, - contextUsage: meta?.usage ? { - promptTokens: meta.usage.prompt_tokens, - completionTokens: meta.usage.completion_tokens, - totalTokens: meta.usage.total_tokens, - cachedTokens: meta.timings?.cache_n ?? 0, - modelId: params.checkpoint, - } : undefined, + contextUsage: meta?.usage + ? { + promptTokens: meta.usage.prompt_tokens, + completionTokens: meta.usage.completion_tokens, + totalTokens: meta.usage.total_tokens, + cachedTokens: meta.timings?.cache_n ?? 0, + modelId: params.checkpoint, + } + : undefined, timing: finalTiming, }, }, }; } catch (err) { - settleFirstTokenErr(err instanceof Error ? err : new Error("Generation failed")); + settleFirstTokenErr( + err instanceof Error ? err : new Error("Generation failed"), + ); if (!abortSignal.aborted) { const msg = err instanceof Error ? err.message : String(err); if (isContextLimitError(msg)) { @@ -1739,7 +1823,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { toast.error("Context limit reached", { description: "The conversation has filled the model's context window. " + - "Increase \"Context Length\" in the chat Settings panel (⚙ in the top-right), " + + 'Increase "Context Length" in the chat Settings panel (⚙ in the top-right), ' + "or start a new chat.", duration: 8000, }); @@ -1756,14 +1840,12 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { runtime.setToolStatus(null); clearTimeout(warmupTimer); if (waitingFirstChunk) { - if (!firstTokenSettled) { - if (abortSignal.aborted) { - settleFirstTokenErr(new Error("Cancelled")); - } else { - settleFirstTokenErr(new Error("No tokens received")); - } - } else { + if (firstTokenSettled) { settleFirstTokenOk(); + } else if (abortSignal.aborted) { + settleFirstTokenErr(new Error("Cancelled")); + } else { + settleFirstTokenErr(new Error("No tokens received")); } } runtime.setThreadRunning(threadKey, false); diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index f842144723..81303d9311 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -3,6 +3,7 @@ import { authFetch } from "@/features/auth"; import { formatFastApiDetail } from "@/lib/format-fastapi-error"; +import type { MessageRecord, ModelType, ThreadRecord } from "../types"; import type { AudioGenerationResponse, GgufVariantsResponse, @@ -17,6 +18,14 @@ import type { ValidateModelResponse, } from "../types/api"; +export const CHAT_HISTORY_UPDATED_EVENT = "unsloth-chat-history-updated"; + +export function notifyChatHistoryUpdated(): void { + if (typeof window !== "undefined") { + window.dispatchEvent(new Event(CHAT_HISTORY_UPDATED_EVENT)); + } +} + function parseErrorText(status: number, body: unknown): string { if (body && typeof body === "object") { const detail = (body as { detail?: unknown }).detail; @@ -41,7 +50,9 @@ export async function listModels(): Promise { return parseJsonOrThrow(response); } -export async function listLoras(outputsDir?: string): Promise { +export async function listLoras( + outputsDir?: string, +): Promise { const query = outputsDir ? `?${new URLSearchParams({ outputs_dir: outputsDir }).toString()}` : ""; @@ -104,13 +115,19 @@ export async function getGgufDownloadProgress( repoId: string, variant: string, expectedBytes: number, -): Promise<{ downloaded_bytes: number; expected_bytes: number; progress: number }> { +): Promise<{ + downloaded_bytes: number; + expected_bytes: number; + progress: number; +}> { const params = new URLSearchParams({ repo_id: repoId, variant, expected_bytes: String(expectedBytes), }); - const response = await authFetch(`/api/models/gguf-download-progress?${params}`); + const response = await authFetch( + `/api/models/gguf-download-progress?${params}`, + ); return parseJsonOrThrow(response); } @@ -205,7 +222,10 @@ export async function listCachedModels(): Promise { return data.cached; } -export async function deleteCachedModel(repoId: string, variant?: string): Promise { +export async function deleteCachedModel( + repoId: string, + variant?: string, +): Promise { const payload: Record = { repo_id: repoId }; if (variant) payload.variant = variant; const response = await authFetch("/api/models/delete-cached", { @@ -263,6 +283,246 @@ export async function removeScanFolder(id: number): Promise { await parseJsonOrThrow(response); } +export async function listChatThreads( + args: { + modelType?: ModelType; + pairId?: string; + includeArchived?: boolean; + } = {}, +): Promise { + const params = new URLSearchParams(); + if (args.modelType) params.set("model_type", args.modelType); + if (args.pairId) params.set("pair_id", args.pairId); + if (args.includeArchived !== undefined) { + params.set("include_archived", String(args.includeArchived)); + } + const qs = params.toString(); + const response = await authFetch(`/api/chat/threads${qs ? `?${qs}` : ""}`); + const data = await parseJsonOrThrow<{ threads: ThreadRecord[] }>(response); + return data.threads; +} + +export async function getChatThread( + threadId: string, +): Promise { + const response = await authFetch( + `/api/chat/threads/${encodeURIComponent(threadId)}`, + ); + if (response.status === 404) return null; + return parseJsonOrThrow(response); +} + +export async function saveChatThread( + thread: ThreadRecord, +): Promise { + const response = await authFetch("/api/chat/threads", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(thread), + }); + const savedThread = await parseJsonOrThrow(response); + notifyChatHistoryUpdated(); + return savedThread; +} + +export async function updateChatThread( + threadId: string, + patch: Partial, +): Promise { + const response = await authFetch( + `/api/chat/threads/${encodeURIComponent(threadId)}`, + { + method: "PATCH", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(patch), + }, + ); + const thread = await parseJsonOrThrow(response); + notifyChatHistoryUpdated(); + return thread; +} + +export async function deleteChatThreads(threadIds: string[]): Promise { + if (threadIds.length === 0) return; + const response = await authFetch("/api/chat/threads", { + method: "DELETE", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ ids: threadIds }), + }); + await parseJsonOrThrow(response); + notifyChatHistoryUpdated(); +} + +export async function listChatMessages( + threadId: string, +): Promise { + const response = await authFetch( + `/api/chat/threads/${encodeURIComponent(threadId)}/messages`, + ); + if (response.status === 404) return []; + const data = await parseJsonOrThrow<{ messages: MessageRecord[] }>(response); + return data.messages; +} + +/** + * Fetch messages for many threads in one HTTP call. Falls back to + * per-thread listChatMessages on 404/405 (older servers without the + * batch route). + */ +export async function batchListChatMessages( + threadIds: string[], +): Promise> { + const out = new Map(); + if (threadIds.length === 0) return out; + const response = await authFetch("/api/chat/messages:batch", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ threadIds }), + }); + if (response.status === 404 || response.status === 405) { + // Older server: fall back to per-thread fetches. + const per = await Promise.all( + threadIds.map(async (id) => [id, await listChatMessages(id)] as const), + ); + for (const [id, msgs] of per) out.set(id, msgs); + return out; + } + const data = await parseJsonOrThrow<{ + messagesByThreadId: Record; + }>(response); + for (const id of threadIds) { + out.set(id, data.messagesByThreadId[id] ?? []); + } + return out; +} + +export async function getChatMessage( + threadId: string, + messageId: string, +): Promise { + const response = await authFetch( + `/api/chat/threads/${encodeURIComponent(threadId)}/messages/${encodeURIComponent(messageId)}`, + ); + if (response.status === 404) return null; + return parseJsonOrThrow(response); +} + +export async function saveChatMessage( + message: MessageRecord, +): Promise { + const response = await authFetch( + `/api/chat/threads/${encodeURIComponent(message.threadId)}/messages/${encodeURIComponent(message.id)}`, + { + method: "PUT", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(message), + }, + ); + const savedMessage = await parseJsonOrThrow(response); + notifyChatHistoryUpdated(); + return savedMessage; +} + +export async function syncChatMessages( + threadId: string, + messages: MessageRecord[], + options: { pruneMissing?: boolean } = {}, +): Promise { + const response = await authFetch( + `/api/chat/threads/${encodeURIComponent(threadId)}/messages`, + { + method: "PUT", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + messages, + pruneMissing: options.pruneMissing ?? false, + }), + }, + ); + const data = await parseJsonOrThrow<{ messages: MessageRecord[] }>(response); + notifyChatHistoryUpdated(); + return data.messages; +} + +export async function countBackendChats(): Promise { + const response = await authFetch("/api/chat/count"); + const data = await parseJsonOrThrow<{ count: number }>(response); + return data.count; +} + +export async function clearBackendChats( + options: { notify?: boolean } = {}, +): Promise { + const response = await authFetch("/api/chat", { method: "DELETE" }); + await parseJsonOrThrow(response); + if (options.notify !== false) { + notifyChatHistoryUpdated(); + } +} + +export async function buildBackendChatExport(): Promise<{ + exportedAt: string; + version: number; + threadCount: number; + threads: ThreadRecord[]; + messages: MessageRecord[]; +}> { + const response = await authFetch("/api/chat/export"); + return parseJsonOrThrow(response); +} + +// Legacy-Dexie import ledger. The server-side source of truth that +// replaces the boolean localStorage sentinel +// (`unsloth_chat_legacy_imported_to_studio_db`) so a studio.db wipe +// makes the import recoverable. +export async function listChatImportLedger(): Promise> { + const response = await authFetch("/api/chat/import-ledger"); + // Backend deployments that don't have this endpoint yet behave the + // same as an empty ledger -- caller treats every legacy thread as + // un-imported and tries to import. The UPSERT semantics in + // syncChatMessages prevent duplicates, so this fallback is safe. + if (response.status === 404 || response.status === 405) return new Set(); + const data = await parseJsonOrThrow<{ threadIds: string[] }>(response); + return new Set(data.threadIds); +} + +export interface RecordChatImportLedgerResult { + accepted: number; + inserted: number; + // false when the backend predates /api/chat/import-ledger (404/405/501) + // so the caller can avoid poisoning the localStorage perf hint -- the + // next launch will retry the (idempotent) import. + supported: boolean; +} + +export async function recordChatImportLedger( + threadIds: string[], +): Promise { + if (threadIds.length === 0) { + return { accepted: 0, inserted: 0, supported: true }; + } + const response = await authFetch("/api/chat/import-ledger", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ threadIds }), + }); + if ( + response.status === 404 || + response.status === 405 || + response.status === 501 + ) { + return { accepted: 0, inserted: 0, supported: false }; + } + const data = await parseJsonOrThrow<{ accepted: number; inserted: number }>( + response, + ); + return { + accepted: data.accepted, + inserted: data.inserted, + supported: true, + }; +} + export interface BrowseEntry { name: string; has_models: boolean; @@ -382,12 +642,17 @@ export async function* streamChatCompletions( } // Tool status events are custom SSE payloads, not OpenAI chunks if ("type" in parsed && parsed.type === "tool_status") { - yield { _toolStatus: parsed.content ?? "" } as unknown as OpenAIChatChunk; + yield { + _toolStatus: parsed.content ?? "", + } as unknown as OpenAIChatChunk; separatorIndex = buffer.search(/\r?\n\r?\n/); continue; } // Tool start/end events carry full input/output for the tool outputs panel - if ("type" in parsed && (parsed.type === "tool_start" || parsed.type === "tool_end")) { + if ( + "type" in parsed && + (parsed.type === "tool_start" || parsed.type === "tool_end") + ) { yield { _toolEvent: parsed } as unknown as OpenAIChatChunk; separatorIndex = buffer.search(/\r?\n\r?\n/); continue; diff --git a/studio/frontend/src/features/chat/api/chat-settings-api.ts b/studio/frontend/src/features/chat/api/chat-settings-api.ts new file mode 100644 index 0000000000..4208c90c76 --- /dev/null +++ b/studio/frontend/src/features/chat/api/chat-settings-api.ts @@ -0,0 +1,90 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { authFetch } from "@/features/auth"; +import type { ChatPresetSource } from "../presets/preset-policy"; +import type { ReasoningEffort } from "../stores/chat-runtime-store"; +import type { InferenceParams } from "../types/runtime"; + +export type PersistedInferenceParams = Partial< + Omit +>; + +export interface PersistedChatPreset { + name: string; + params: PersistedInferenceParams; +} + +export interface PersistedChatSettings { + inferenceParams?: PersistedInferenceParams; + customPresets?: PersistedChatPreset[]; + activePreset?: string; + activePresetSource?: ChatPresetSource; + autoTitle?: boolean; + reasoningEffort?: ReasoningEffort; + preserveThinking?: boolean; + autoHealToolCalls?: boolean; + maxToolCallsPerMessage?: number; + toolCallTimeout?: number; +} + +interface ChatSettingsResponse { + settings: PersistedChatSettings; +} + +function parseErrorText(status: number, body: unknown): string { + if ( + body && + typeof body === "object" && + "detail" in body && + typeof body.detail === "string" + ) { + return body.detail; + } + if ( + body && + typeof body === "object" && + "detail" in body && + body.detail != null + ) { + return `Request failed (${status}): ${JSON.stringify(body.detail)}`; + } + if ( + body && + typeof body === "object" && + "message" in body && + typeof body.message === "string" + ) { + return body.message; + } + return `Request failed (${status})`; +} + +async function parseJsonOrThrow(response: Response): Promise { + const body = await response.json().catch(() => null); + if (!response.ok) { + throw new Error(parseErrorText(response.status, body)); + } + return body as T; +} + +export async function getChatSettings(): Promise { + const response = await authFetch("/api/chat/settings"); + const data = await parseJsonOrThrow(response); + return data.settings; +} + +export async function saveChatSettingsPatch( + patch: PersistedChatSettings, + options: { keepalive?: boolean } = {}, +): Promise { + const response = await authFetch("/api/chat/settings", { + method: "PUT", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(patch), + // keepalive lets the PUT survive a tab close from the beforeunload flush. + keepalive: options.keepalive, + }); + const data = await parseJsonOrThrow(response); + return data.settings; +} diff --git a/studio/frontend/src/features/chat/api/providers-api.ts b/studio/frontend/src/features/chat/api/providers-api.ts index e0faac27b4..0b86c7ee82 100644 --- a/studio/frontend/src/features/chat/api/providers-api.ts +++ b/studio/frontend/src/features/chat/api/providers-api.ts @@ -1,226 +1,226 @@ -// SPDX-License-Identifier: AGPL-3.0-only -// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 - -import forge from "node-forge"; -import { authFetch } from "@/features/auth"; -import { formatFastApiDetail } from "@/lib/format-fastapi-error"; - -export interface ProviderRegistryEntry { - provider_type: string; - display_name: string; - base_url: string; - default_models: string[]; - supports_streaming: boolean; - supports_vision: boolean; - supports_tool_calling: boolean; - /** remote = fetch /models; curated = huge catalogs — UI uses defaults + manual IDs only */ - model_list_mode?: "remote" | "curated"; -} - -export interface ProviderConfig { - id: string; - provider_type: string; - display_name: string; - base_url: string; - is_enabled: boolean; - created_at: string; - updated_at: string; -} - -export interface ProviderModelInfo { - id: string; - display_name: string; - context_length?: number | null; - owned_by?: string | null; -} - -export interface ProviderTestResult { - success: boolean; - message: string; - models_count?: number | null; -} - -function parseErrorText(status: number, body: unknown): string { - if (body && typeof body === "object") { - const detail = (body as { detail?: unknown }).detail; - const formatted = formatFastApiDetail(detail); - if (formatted) return formatted; - const message = (body as { message?: unknown }).message; - if (typeof message === "string" && message) return message; - } - return `Request failed (${status})`; -} - -async function parseJsonOrThrow(response: Response): Promise { - const body = await response.json().catch(() => null); - if (!response.ok) { - throw new Error(parseErrorText(response.status, body)); - } - return body as T; -} - -export function isProviderKeyRotationError(error: unknown): boolean { - if (!(error instanceof Error)) return false; - const normalized = error.message.toLowerCase(); - return ( - normalized.includes("public key may have changed") || - normalized.includes("server key may have changed") - ); -} - -let cachedPublicKeyPem: string | null = null; -let cachedForgeKey: forge.pki.rsa.PublicKey | null = null; - -export function clearProviderPublicKeyCache(): void { - cachedPublicKeyPem = null; - cachedForgeKey = null; -} - -async function importProviderPublicKey( - forceRefresh = false, -): Promise { - if (!forceRefresh && cachedForgeKey) { - return cachedForgeKey; - } - const response = await authFetch("/api/providers/public-key"); - const body = await parseJsonOrThrow<{ public_key: string }>(response); - const publicKeyPem = body.public_key?.trim(); - if (!publicKeyPem) { - throw new Error("Provider public key is missing."); - } - if (!forceRefresh && cachedPublicKeyPem === publicKeyPem && cachedForgeKey) { - return cachedForgeKey; - } - const forgeKey = forge.pki.publicKeyFromPem(publicKeyPem); - cachedPublicKeyPem = publicKeyPem; - cachedForgeKey = forgeKey; - return forgeKey; -} - -export async function encryptProviderApiKey( - plaintextApiKey: string, - forceRefresh = false, -): Promise { - const key = await importProviderPublicKey(forceRefresh); - const encrypted = key.encrypt(plaintextApiKey, "RSA-OAEP", { - md: forge.md.sha256.create(), - mgf1: { md: forge.md.sha256.create() }, - }); - return forge.util.encode64(encrypted); -} - -export async function listProviderRegistry(): Promise { - const response = await authFetch("/api/providers/registry"); - return parseJsonOrThrow(response); -} - -export async function listProviderConfigs(): Promise { - const response = await authFetch("/api/providers/"); - return parseJsonOrThrow(response); -} - -export async function createProviderConfig(payload: { - providerType: string; - displayName: string; - baseUrl?: string | null; -}): Promise { - const response = await authFetch("/api/providers/", { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ - provider_type: payload.providerType, - display_name: payload.displayName, - base_url: payload.baseUrl ?? null, - }), - }); - return parseJsonOrThrow(response); -} - -export async function deleteProviderConfig(providerId: string): Promise { - const response = await authFetch(`/api/providers/${providerId}`, { - method: "DELETE", - }); - if (!response.ok) { - const body = await response.json().catch(() => null); - throw new Error(parseErrorText(response.status, body)); - } -} - -export async function updateProviderConfig( - providerId: string, - payload: { - displayName?: string; - baseUrl?: string | null; - isEnabled?: boolean; - }, -): Promise { - const response = await authFetch(`/api/providers/${providerId}`, { - method: "PUT", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ - ...(payload.displayName === undefined ? {} : { display_name: payload.displayName }), - ...(payload.baseUrl === undefined ? {} : { base_url: payload.baseUrl }), - ...(payload.isEnabled === undefined ? {} : { is_enabled: payload.isEnabled }), - }), - }); - return parseJsonOrThrow(response); -} - -async function withApiKeyEncryptionRetry( - plaintextApiKey: string, - call: (encryptedApiKey: string | null) => Promise, -): Promise { - // Empty key (local providers): skip RSA round-trip and let the backend omit auth. - if (!plaintextApiKey) { - return await call(null); - } - try { - const encrypted = await encryptProviderApiKey(plaintextApiKey, false); - return await call(encrypted); - } catch (error) { - if (!isProviderKeyRotationError(error)) { - throw error; - } - clearProviderPublicKeyCache(); - const encrypted = await encryptProviderApiKey(plaintextApiKey, true); - return await call(encrypted); - } -} - -export async function testProviderConnection(payload: { - providerType: string; - apiKey: string; - baseUrl?: string | null; -}): Promise { - return withApiKeyEncryptionRetry(payload.apiKey, async (encryptedApiKey) => { - const response = await authFetch("/api/providers/test", { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ - provider_type: payload.providerType, - encrypted_api_key: encryptedApiKey, - base_url: payload.baseUrl ?? null, - }), - }); - return parseJsonOrThrow(response); - }); -} - -export async function listProviderModels(payload: { - providerType: string; - apiKey: string; - baseUrl?: string | null; -}): Promise { - return withApiKeyEncryptionRetry(payload.apiKey, async (encryptedApiKey) => { - const response = await authFetch("/api/providers/models", { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ - provider_type: payload.providerType, - encrypted_api_key: encryptedApiKey, - base_url: payload.baseUrl ?? null, - }), - }); - return parseJsonOrThrow(response); - }); -} +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import forge from "node-forge"; +import { authFetch } from "@/features/auth"; +import { formatFastApiDetail } from "@/lib/format-fastapi-error"; + +export interface ProviderRegistryEntry { + provider_type: string; + display_name: string; + base_url: string; + default_models: string[]; + supports_streaming: boolean; + supports_vision: boolean; + supports_tool_calling: boolean; + /** remote = fetch /models; curated = huge catalogs — UI uses defaults + manual IDs only */ + model_list_mode?: "remote" | "curated"; +} + +export interface ProviderConfig { + id: string; + provider_type: string; + display_name: string; + base_url: string; + is_enabled: boolean; + created_at: string; + updated_at: string; +} + +export interface ProviderModelInfo { + id: string; + display_name: string; + context_length?: number | null; + owned_by?: string | null; +} + +export interface ProviderTestResult { + success: boolean; + message: string; + models_count?: number | null; +} + +function parseErrorText(status: number, body: unknown): string { + if (body && typeof body === "object") { + const detail = (body as { detail?: unknown }).detail; + const formatted = formatFastApiDetail(detail); + if (formatted) return formatted; + const message = (body as { message?: unknown }).message; + if (typeof message === "string" && message) return message; + } + return `Request failed (${status})`; +} + +async function parseJsonOrThrow(response: Response): Promise { + const body = await response.json().catch(() => null); + if (!response.ok) { + throw new Error(parseErrorText(response.status, body)); + } + return body as T; +} + +export function isProviderKeyRotationError(error: unknown): boolean { + if (!(error instanceof Error)) return false; + const normalized = error.message.toLowerCase(); + return ( + normalized.includes("public key may have changed") || + normalized.includes("server key may have changed") + ); +} + +let cachedPublicKeyPem: string | null = null; +let cachedForgeKey: forge.pki.rsa.PublicKey | null = null; + +export function clearProviderPublicKeyCache(): void { + cachedPublicKeyPem = null; + cachedForgeKey = null; +} + +async function importProviderPublicKey( + forceRefresh = false, +): Promise { + if (!forceRefresh && cachedForgeKey) { + return cachedForgeKey; + } + const response = await authFetch("/api/providers/public-key"); + const body = await parseJsonOrThrow<{ public_key: string }>(response); + const publicKeyPem = body.public_key?.trim(); + if (!publicKeyPem) { + throw new Error("Provider public key is missing."); + } + if (!forceRefresh && cachedPublicKeyPem === publicKeyPem && cachedForgeKey) { + return cachedForgeKey; + } + const forgeKey = forge.pki.publicKeyFromPem(publicKeyPem); + cachedPublicKeyPem = publicKeyPem; + cachedForgeKey = forgeKey; + return forgeKey; +} + +export async function encryptProviderApiKey( + plaintextApiKey: string, + forceRefresh = false, +): Promise { + const key = await importProviderPublicKey(forceRefresh); + const encrypted = key.encrypt(plaintextApiKey, "RSA-OAEP", { + md: forge.md.sha256.create(), + mgf1: { md: forge.md.sha256.create() }, + }); + return forge.util.encode64(encrypted); +} + +export async function listProviderRegistry(): Promise { + const response = await authFetch("/api/providers/registry"); + return parseJsonOrThrow(response); +} + +export async function listProviderConfigs(): Promise { + const response = await authFetch("/api/providers/"); + return parseJsonOrThrow(response); +} + +export async function createProviderConfig(payload: { + providerType: string; + displayName: string; + baseUrl?: string | null; +}): Promise { + const response = await authFetch("/api/providers/", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + provider_type: payload.providerType, + display_name: payload.displayName, + base_url: payload.baseUrl ?? null, + }), + }); + return parseJsonOrThrow(response); +} + +export async function deleteProviderConfig(providerId: string): Promise { + const response = await authFetch(`/api/providers/${providerId}`, { + method: "DELETE", + }); + if (!response.ok) { + const body = await response.json().catch(() => null); + throw new Error(parseErrorText(response.status, body)); + } +} + +export async function updateProviderConfig( + providerId: string, + payload: { + displayName?: string; + baseUrl?: string | null; + isEnabled?: boolean; + }, +): Promise { + const response = await authFetch(`/api/providers/${providerId}`, { + method: "PUT", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + ...(payload.displayName === undefined ? {} : { display_name: payload.displayName }), + ...(payload.baseUrl === undefined ? {} : { base_url: payload.baseUrl }), + ...(payload.isEnabled === undefined ? {} : { is_enabled: payload.isEnabled }), + }), + }); + return parseJsonOrThrow(response); +} + +async function withApiKeyEncryptionRetry( + plaintextApiKey: string, + call: (encryptedApiKey: string | null) => Promise, +): Promise { + // Empty key (local providers): skip RSA round-trip and let the backend omit auth. + if (!plaintextApiKey) { + return await call(null); + } + try { + const encrypted = await encryptProviderApiKey(plaintextApiKey, false); + return await call(encrypted); + } catch (error) { + if (!isProviderKeyRotationError(error)) { + throw error; + } + clearProviderPublicKeyCache(); + const encrypted = await encryptProviderApiKey(plaintextApiKey, true); + return await call(encrypted); + } +} + +export async function testProviderConnection(payload: { + providerType: string; + apiKey: string; + baseUrl?: string | null; +}): Promise { + return withApiKeyEncryptionRetry(payload.apiKey, async (encryptedApiKey) => { + const response = await authFetch("/api/providers/test", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + provider_type: payload.providerType, + encrypted_api_key: encryptedApiKey, + base_url: payload.baseUrl ?? null, + }), + }); + return parseJsonOrThrow(response); + }); +} + +export async function listProviderModels(payload: { + providerType: string; + apiKey: string; + baseUrl?: string | null; +}): Promise { + return withApiKeyEncryptionRetry(payload.apiKey, async (encryptedApiKey) => { + const response = await authFetch("/api/providers/models", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + provider_type: payload.providerType, + encrypted_api_key: encryptedApiKey, + base_url: payload.baseUrl ?? null, + }), + }); + return parseJsonOrThrow(response); + }); +} diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index a62270ed50..2f062568dd 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -1,6 +1,7 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +import type { ChatSearch } from "@/app/routes/chat"; import { type DeletedModelRef, type ExternalModelOption, @@ -9,22 +10,22 @@ import { ModelSelector, } from "@/components/assistant-ui/model-selector"; import { Thread } from "@/components/assistant-ui/thread"; +import { useSidebar } from "@/components/ui/sidebar"; +import { Tooltip, TooltipContent } from "@/components/ui/tooltip"; import { NativeModelChip } from "@/features/native-intents/components/native-model-chip"; import { NativeModelDropOverlay } from "@/features/native-intents/components/native-model-drop-overlay"; +import { useNativeIntentStore } from "@/features/native-intents/store"; +import type { NativeIntent } from "@/features/native-intents/types"; import { useChooseNativeModel } from "@/features/native-intents/use-native-dialogs"; import { useNativeModelDrop } from "@/features/native-intents/use-native-drop"; import { useNativePathLeasesSupported } from "@/features/native-intents/use-native-readiness"; -import { useNativeIntentStore } from "@/features/native-intents/store"; -import type { NativeIntent } from "@/features/native-intents/types"; +import { GuidedTour, useGuidedTourController } from "@/features/tour"; import { isTauri } from "@/lib/api-base"; import { cn } from "@/lib/utils"; -import { GuidedTour, useGuidedTourController } from "@/features/tour"; -import { useSidebar } from "@/components/ui/sidebar"; import { CustomizeIcon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; -import { Tooltip, TooltipContent } from "@/components/ui/tooltip"; -import { Tooltip as TooltipPrimitive } from "radix-ui"; import { useNavigate, useSearch } from "@tanstack/react-router"; +import { Tooltip as TooltipPrimitive } from "radix-ui"; import { type ReactElement, memo, @@ -35,18 +36,21 @@ import { useState, } from "react"; import { toast } from "@/lib/toast"; -import type { ChatSearch } from "@/app/routes/chat"; import { listLocalModels } from "./api/chat-api"; import { ChatSettingsPanel } from "./chat-settings-sheet"; import { CopyableErrorChip } from "@/components/ui/copyable-error-chip"; import { ContextUsageBar } from "./components/context-usage-bar"; import { ModelLoadInlineStatus } from "./components/model-load-status"; -import { db } from "./db"; import { buildExternalModelId, isExternalModelId, parseExternalModelId, } from "./external-providers"; +import { useChatModelRuntime } from "./hooks/use-chat-model-runtime"; +import { + clearTrainingCompareHandoff, + getTrainingCompareHandoff, +} from "./lib/training-compare-handoff"; import { clampReasoningEffortToLevels, getExternalReasoningCapabilities, @@ -54,11 +58,6 @@ import { providerSupportsBuiltinCodeExecution, providerSupportsBuiltinWebSearch, } from "./provider-capabilities"; -import { useChatModelRuntime } from "./hooks/use-chat-model-runtime"; -import { - clearTrainingCompareHandoff, - getTrainingCompareHandoff, -} from "./lib/training-compare-handoff"; import { ChatRuntimeProvider } from "./runtime-provider"; import { type CompareHandle, @@ -76,6 +75,12 @@ import { import { useExternalProvidersStore } from "./stores/external-providers-store"; import { buildChatTourSteps } from "./tour"; import type { ChatView, MessageRecord } from "./types"; +import { + getStoredChatThread, + isExpectedBackgroundChatStorageError, + listStoredChatMessages, + listStoredChatThreads, +} from "./utils/chat-history-storage"; type LoraCandidate = { id: string; @@ -322,15 +327,15 @@ const LoraCompareContent = memo(function LoraCompareContent({ useEffect(() => { let isActive = true; - db.threads - .where("pairId") - .equals(pairId) - .toArray() - .then((threads) => { - if (!isActive) return; - setBaseThreadId(threads.find((t) => t.modelType === "base")?.id); - setLoraThreadId(threads.find((t) => t.modelType === "lora")?.id); - }); + listStoredChatThreads({ pairId }).then((threads) => { + if (!isActive) return; + setBaseThreadId(threads.find((t) => t.modelType === "base")?.id); + setLoraThreadId(threads.find((t) => t.modelType === "lora")?.id); + }).catch((error) => { + if (!isExpectedBackgroundChatStorageError(error)) { + throw error; + } + }); return () => { isActive = false; }; @@ -471,23 +476,21 @@ const GeneralCompareContent = memo(function GeneralCompareContent({ useEffect(() => { let isActive = true; - db.threads - .where("pairId") - .equals(pairId) - .toArray() - .then((threads) => { - if (!isActive) return; - setModel1ThreadId( - threads.find( - (t) => t.modelType === "model1" || t.modelType === "base", - )?.id, - ); - setModel2ThreadId( - threads.find( - (t) => t.modelType === "model2" || t.modelType === "lora", - )?.id, - ); - }); + listStoredChatThreads({ pairId }).then((threads) => { + if (!isActive) return; + setModel1ThreadId( + threads.find((t) => t.modelType === "model1" || t.modelType === "base") + ?.id, + ); + setModel2ThreadId( + threads.find((t) => t.modelType === "model2" || t.modelType === "lora") + ?.id, + ); + }).catch((error) => { + if (!isExpectedBackgroundChatStorageError(error)) { + throw error; + } + }); return () => { isActive = false; }; @@ -565,6 +568,9 @@ export function ChatPage(): ReactElement { const settingsOpen = useChatRuntimeStore((s) => s.settingsPanelOpen); const setSettingsOpen = useChatRuntimeStore((s) => s.setSettingsPanelOpen); + const hydratePersistedSettings = useChatRuntimeStore( + (s) => s.hydratePersistedSettings, + ); const externalProviders = useExternalProvidersStore((s) => s.providers); const connectionsEnabled = useExternalProvidersStore( (s) => s.connectionsEnabled, @@ -572,13 +578,16 @@ export function ChatPage(): ReactElement { const setExternalProviders = useExternalProvidersStore((s) => s.setProviders); const externalProvidersForChat = connectionsEnabled ? externalProviders : []; + useEffect(() => { + void hydratePersistedSettings(); + }, [hydratePersistedSettings]); + useEffect(() => { const threadId = search.thread; if (!threadId) return; let canceled = false; - void db.threads - .get(threadId) + void getStoredChatThread(threadId) .then((thread) => { if (canceled || thread) return; useChatRuntimeStore.getState().setActiveThreadId(null); @@ -1041,12 +1050,9 @@ export function ChatPage(): ReactElement { void (async () => { let showImageCompatibilityWarning = false; if (view.mode === "single" && activeThreadId) { - const thread = await db.threads.get(activeThreadId); + const thread = await getStoredChatThread(activeThreadId); if (thread?.modelId && thread.modelId !== value) { - const messages = await db.messages - .where("threadId") - .equals(activeThreadId) - .toArray(); + const messages = await listStoredChatMessages(activeThreadId); if (messages.length > 0) { const hasImage = messages.some(messageHasImage); const targetModel = modelsFromStore.find( @@ -1131,17 +1137,22 @@ export function ChatPage(): ReactElement { const threadId = saved.thread ?? useChatRuntimeStore.getState().activeThreadId; if (threadId) { - void db.messages - .where("threadId") - .equals(threadId) - .reverse() - .first() + void listStoredChatMessages(threadId) + .then( + (messages) => + [...messages].sort((a, b) => b.createdAt - a.createdAt)[0], + ) .then((msg) => { const metadata = msg?.metadata as Record | undefined; const usage = metadata?.contextUsage as ReturnType< typeof useChatRuntimeStore.getState >["contextUsage"]; if (usage) useChatRuntimeStore.getState().setContextUsage(usage); + }) + .catch((error) => { + if (!isExpectedBackgroundChatStorageError(error)) { + throw error; + } }); } }, [navigate]); @@ -1464,7 +1475,7 @@ export function ChatPage(): ReactElement { ) : null} {!settingsOpen && ( - + diff --git a/studio/frontend/src/features/settings/tabs/general-tab.tsx b/studio/frontend/src/features/settings/tabs/general-tab.tsx index 9ea89f55cf..6ce9ad63b3 100644 --- a/studio/frontend/src/features/settings/tabs/general-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/general-tab.tsx @@ -199,7 +199,7 @@ export function GeneralTab() {