diff --git a/api/services/inference.py b/api/services/inference.py index ebd999c..9ea0e27 100755 --- a/api/services/inference.py +++ b/api/services/inference.py @@ -559,12 +559,9 @@ async def generate_stream_explanation(topic: str, level: str, model: str | None telemetry_sink=stream_telemetry, ): socratic_chunks.append(chunk) - yield chunk # yield immediately - do not buffer - - # NOTE: Socratic constraint enforcement (question capping) is now - # the caller's responsibility. The full response is available in - # the router's accumulated `full_content` after streaming ends. - # This trade-off is intentional: streaming health > post-processing. + constrained_response = _enforce_socratic_response_constraints("".join(socratic_chunks)) + for index in range(0, len(constrained_response), 400): + yield constrained_response[index : index + 400] else: async for chunk in stream_chat_completion( model=alias, diff --git a/api/tests/test_inference.py b/api/tests/test_inference.py index 5cf804c..d395222 100644 --- a/api/tests/test_inference.py +++ b/api/tests/test_inference.py @@ -82,9 +82,40 @@ async def fake_stream(*_args, **_kwargs): combined = "".join(streamed) assert combined.count("?") <= 3 + assert combined.count("How does it change in this process?") == 1 assert "Share your answer, and I will guide the next step." in combined +@pytest.mark.asyncio +async def test_generate_stream_explanation_socratic_dedupes_and_caps_questions(monkeypatch): + async def fake_stream(*_args, **_kwargs): + chunks = [ + "What is entropy? ", + "How does entropy change here? ", + "How does entropy change here? ", + "Why does entropy matter? ", + "How would you measure entropy in practice?", + ] + for chunk in chunks: + yield chunk + + monkeypatch.setattr(inference_module, "stream_chat_completion", fake_stream) + + streamed = [] + async for chunk in inference_module.generate_stream_explanation( + "entropy", + "eli15", + mode="socratic", + ): + streamed.append(chunk) + + combined = "".join(streamed) + assert combined.count("?") == 3 + assert combined.count("How does entropy change here?") == 1 + assert "How would you measure entropy in practice?" not in combined + assert combined.endswith("Share your answer, and I will guide the next step.") + + @pytest.mark.asyncio async def test_technical_mode_handler_uses_safe_defaults_when_classification_fails(monkeypatch): captured = {} diff --git a/src/components/chat/DepthDropdown.tsx b/src/components/chat/DepthDropdown.tsx index 3a8d83c..c5512a3 100644 --- a/src/components/chat/DepthDropdown.tsx +++ b/src/components/chat/DepthDropdown.tsx @@ -1,6 +1,6 @@ import { useEffect, useId, useMemo, useRef, useState } from "react"; import { Bolt, ChevronDown } from "lucide-react"; -import type { DepthLevel } from "../../stores/useChatStore"; +import type { DepthLevel } from "../../lib/chatStoreUtils"; interface DepthOption { id: DepthLevel; diff --git a/src/components/chat/WelcomeEmptyState.tsx b/src/components/chat/WelcomeEmptyState.tsx index 6f0d7e6..ebcb2c0 100644 --- a/src/components/chat/WelcomeEmptyState.tsx +++ b/src/components/chat/WelcomeEmptyState.tsx @@ -1,4 +1,4 @@ -import type { Workspace } from "../../stores/useChatStore"; +import type { Workspace } from "../../lib/chatStoreUtils"; import { WORKSPACE_CONTENT, WORKSPACE_ICONS, diff --git a/src/components/chat/WorkspaceInput.tsx b/src/components/chat/WorkspaceInput.tsx index 695e318..719ac3f 100644 --- a/src/components/chat/WorkspaceInput.tsx +++ b/src/components/chat/WorkspaceInput.tsx @@ -2,7 +2,7 @@ import { useEffect, useMemo, useRef, useState } from "react"; import { ArrowRight, Paperclip } from "lucide-react"; import DepthDropdown from "./DepthDropdown"; import { useChatStore } from "../../stores/useChatStore"; -import type { DepthLevel, Workspace } from "../../stores/useChatStore"; +import type { DepthLevel, Workspace } from "../../lib/chatStoreUtils"; import type { ChatMode, PromptMode } from "../../types/chat"; interface WorkspaceInputProps { diff --git a/src/components/chat/WorkspaceSidebar.tsx b/src/components/chat/WorkspaceSidebar.tsx index 088e4da..ea2be01 100644 --- a/src/components/chat/WorkspaceSidebar.tsx +++ b/src/components/chat/WorkspaceSidebar.tsx @@ -10,7 +10,7 @@ import { Trash2, } from "lucide-react"; import type { Conversation } from "../../types/chat"; -import type { Workspace } from "../../stores/useChatStore"; +import type { Workspace } from "../../lib/chatStoreUtils"; interface WorkspaceSidebarProps { workspace: Workspace; diff --git a/src/lib/chatStoreUtils.ts b/src/lib/chatStoreUtils.ts new file mode 100644 index 0000000..000b9d3 --- /dev/null +++ b/src/lib/chatStoreUtils.ts @@ -0,0 +1,293 @@ +import type { Level } from "../types"; +import type { ChatMode, Conversation, Message, PromptMode } from "../types/chat"; +import { CHAT_PREMIUM_MODES, isPromptMode, resolveChatMode } from "../lib/chatModes"; + +export type Workspace = "learn" | "socratic" | "technical"; +export type ThemeMode = "dark" | "light"; +export const DEPTH_LEVELS = [ + "eli5", + "eli10", + "eli12", + "eli15", + "meme", +] as const; +export type DepthLevel = (typeof DEPTH_LEVELS)[number]; +export type StoreLevel = Level; +export type StoreConversation = Conversation; +export const CHAT_STORE_PREMIUM_MODES = CHAT_PREMIUM_MODES; + +export const THEME_STORAGE_KEY = "kb_theme_v1"; +export const DEFAULT_WORKSPACE: Workspace = "learn"; +export const DEFAULT_DEPTH_LEVEL: DepthLevel = "eli12"; +export const PENDING_SYNC_KEY = "kb_pending_sync_v1"; + +export const supabaseConfigured = + Boolean(import.meta.env.VITE_SUPABASE_URL) && + Boolean(import.meta.env.VITE_SUPABASE_ANON_KEY); + +const defaultIsProEnv = import.meta.env.VITE_DEFAULT_IS_PRO; +export const defaultIsPro = defaultIsProEnv ? defaultIsProEnv === "true" : false; +export const API_URL = import.meta.env.VITE_API_URL || ""; + +export const createUuid = () => { + const webCrypto: Crypto | undefined = + typeof globalThis !== "undefined" ? globalThis.crypto : undefined; + + if (webCrypto?.randomUUID) { + return webCrypto.randomUUID(); + } + + const getRandomValues = webCrypto?.getRandomValues + ? webCrypto.getRandomValues.bind(webCrypto) + : null; + const rnd = (size: number) => { + if (getRandomValues) { + const arr = new Uint8Array(size); + getRandomValues(arr); + return arr; + } + return Uint8Array.from({ length: size }, () => + Math.floor(Math.random() * 256), + ); + }; + const bytes = rnd(16); + bytes[6] = (bytes[6] & 0x0f) | 0x40; + bytes[8] = (bytes[8] & 0x3f) | 0x80; + const hex = Array.from(bytes, (b) => b.toString(16).padStart(2, "0")); + return `${hex.slice(0, 4).join("")}-${hex.slice(4, 6).join("")}-${hex + .slice(6, 8) + .join("")}-${hex.slice(8, 10).join("")}-${hex.slice(10, 16).join("")}`; +}; + +export const makeLocalId = () => `local-${createUuid()}`; + +export const makeClientId = () => createUuid(); + +export const truncateTitle = (content: string) => { + const trimmed = content.trim().replace(/\s+/g, " "); + if (trimmed.length <= 64) return trimmed; + return `${trimmed.slice(0, 61)}...`; +}; + +export const isDepthLevel = (mode: string | null | undefined): mode is DepthLevel => { + return DEPTH_LEVELS.includes(mode as DepthLevel); +}; + +export const resolveDepthLevel = ( + mode: string | null | undefined, + fallback: DepthLevel = DEFAULT_DEPTH_LEVEL, +): DepthLevel => { + if (isDepthLevel(mode)) return mode; + return fallback; +}; + +export const resolveWorkspaceFromMode = (mode: ChatMode): Workspace => { + if (mode === "socratic") return "socratic"; + if (mode === "technical") return "technical"; + return "learn"; +}; + +export const resolveWorkspaceState = ( + mode: string | null | undefined, + promptMode: string | null | undefined, + fallbackDepth: DepthLevel, +) => { + const resolvedMode = resolveChatMode(mode); + + if (resolvedMode === "socratic") { + return { + workspace: "socratic" as Workspace, + mode: "socratic" as ChatMode, + promptMode: resolveDepthLevel(promptMode, fallbackDepth) as PromptMode, + depthLevel: resolveDepthLevel(promptMode, fallbackDepth), + }; + } + + if (resolvedMode === "technical") { + return { + workspace: "technical" as Workspace, + mode: "technical" as ChatMode, + promptMode: resolveDepthLevel(promptMode, fallbackDepth) as PromptMode, + depthLevel: resolveDepthLevel(promptMode, fallbackDepth), + }; + } + + const nextDepth = resolveDepthLevel( + promptMode || (isPromptMode(resolvedMode) ? resolvedMode : undefined), + fallbackDepth, + ); + + return { + workspace: "learn" as Workspace, + mode: "learning" as ChatMode, + promptMode: nextDepth as PromptMode, + depthLevel: nextDepth, + }; +}; + +export const loadTheme = (): ThemeMode => { + if (typeof window === "undefined") return "light"; + + const cachedTheme = window.localStorage.getItem(THEME_STORAGE_KEY); + if (cachedTheme === "light" || cachedTheme === "dark") { + return cachedTheme; + } + + return window.matchMedia?.("(prefers-color-scheme: dark)").matches + ? "dark" + : "light"; +}; + +export const applyThemeClass = (theme: ThemeMode) => { + if (typeof document === "undefined") return; + document.documentElement.classList.toggle("dark", theme === "dark"); +}; + +export const persistTheme = (theme: ThemeMode) => { + if (typeof window === "undefined") return; + window.localStorage.setItem(THEME_STORAGE_KEY, theme); +}; + +export const getModeForWorkspace = (workspace: Workspace): ChatMode => { + if (workspace === "socratic") return "socratic"; + if (workspace === "technical") return "technical"; + return "learning"; +}; + +export const asString = (value: unknown): string | undefined => { + return typeof value === "string" ? value : undefined; +}; + +export const isAbortError = (error: unknown): boolean => { + return ( + typeof error === "object" && + error !== null && + "name" in error && + (error as { name?: string }).name === "AbortError" + ); +}; + +export const getErrorMessage = (error: unknown, fallback: string): string => { + if (error instanceof Error && error.message) return error.message; + return fallback; +}; + +export const notifyError = (message: string) => { + console.error(message); + if (typeof window !== "undefined") { + window.dispatchEvent( + new CustomEvent("kb-toast", { detail: { type: "error", message } }), + ); + } +}; + +export const resolveMessageKey = (message: Message) => { + return ( + message.clientGeneratedId || + message.metadata?.assistant_client_id || + message.metadata?.client_id || + message.serverMessageId || + message.id + ); +}; + +export const messagesMatch = (existing: Message, incoming: Message) => { + if (existing.id === incoming.id) return true; + if ( + existing.clientGeneratedId && + incoming.clientGeneratedId && + existing.clientGeneratedId === incoming.clientGeneratedId + ) { + return true; + } + if ( + incoming.metadata?.assistant_client_id && + existing.clientGeneratedId === incoming.metadata.assistant_client_id + ) { + return true; + } + if ( + existing.metadata?.assistant_client_id && + incoming.clientGeneratedId && + existing.metadata.assistant_client_id === incoming.clientGeneratedId + ) { + return true; + } + if ( + existing.serverMessageId && + incoming.id && + existing.serverMessageId === incoming.id + ) + return true; + if ( + incoming.serverMessageId && + existing.id && + incoming.serverMessageId === existing.id + ) + return true; + if ( + incoming.metadata?.client_id && + existing.id === incoming.metadata.client_id + ) + return true; + if ( + existing.metadata?.client_id && + existing.metadata.client_id === incoming.id + ) + return true; + if ( + incoming.metadata?.client_id && + existing.clientGeneratedId === incoming.metadata.client_id + ) + return true; + if ( + existing.metadata?.client_id && + incoming.clientGeneratedId && + existing.metadata.client_id === incoming.clientGeneratedId + ) { + return true; + } + if ( + incoming.metadata?.client_id && + existing.metadata?.client_id && + existing.metadata.client_id === incoming.metadata.client_id + ) { + return true; + } + return false; +}; + +export const findExistingMessageKey = ( + state: { messagesById: Record; messageIds: string[] }, + incoming: Message, +) => { + for (const messageKey of state.messageIds) { + const existing = state.messagesById[messageKey]; + if (!existing) continue; + if (messagesMatch(existing, incoming)) { + return messageKey; + } + } + return null; +}; + +export const buildMessageRegistry = (messages: Message[]) => { + const messagesById: Record = {}; + const messageIds: string[] = []; + + for (const message of messages) { + const key = resolveMessageKey(message); + if (!key) { + console.warn("Message has no identifiable key, skipping:", message); + continue; + } + if (messagesById[key]) { + messagesById[key] = { ...messagesById[key], ...message }; + continue; + } + messagesById[key] = message; + messageIds.push(key); + } + + return { messagesById, messageIds }; +}; diff --git a/src/pages/ChatPage.tsx b/src/pages/ChatPage.tsx index a8de901..629d161 100644 --- a/src/pages/ChatPage.tsx +++ b/src/pages/ChatPage.tsx @@ -14,7 +14,7 @@ import { getHealth } from "../api"; import { useConversations } from "../hooks/useConversations"; import { useMessages } from "../hooks/useMessages"; import { useChatStore } from "../stores/useChatStore"; -import type { Workspace } from "../stores/useChatStore"; +import type { Workspace } from "../lib/chatStoreUtils"; const WORKSPACE_LABELS: Record = { learn: "Learn", diff --git a/src/services/chatService.ts b/src/services/chatService.ts new file mode 100644 index 0000000..f9985e1 --- /dev/null +++ b/src/services/chatService.ts @@ -0,0 +1,255 @@ +import { supabase } from "../lib/supabase"; +import { splitSseEvents, parseSseEvent } from "../lib/sse"; +import { ChatStreamChunkSchema } from "../lib/sseSchemas"; +import { getTracePropagationHeaders } from "../lib/monitoring"; +import { toQueryLevel } from "../lib/chatModes"; +import type { ChatMode, PromptMode } from "../types/chat"; +import { API_URL, createUuid, supabaseConfigured } from "../lib/chatStoreUtils"; + +interface SendChatParams { + conversationId: string; + content: string; + mode: ChatMode; + promptMode: PromptMode; + temperature: number; + isPro: boolean; + isRegeneration?: boolean; + clientMessageId: string; + assistantClientId: string; + signal: AbortSignal; + onChunk: (chunk: string) => void; + onServerMessageId?: (id: string) => void; + onError: (error: Error) => void; + onDone: () => void; +} + +const getSupabaseSession = async () => { + if (!supabaseConfigured) return null; + const { data } = await supabase.auth.getSession(); + return data.session; +}; + +const buildHeaders = async (): Promise> => { + const headers: Record = { "Content-Type": "application/json" }; + const session = await getSupabaseSession(); + if (session?.access_token) { + headers.Authorization = `Bearer ${session.access_token}`; + } + Object.assign(headers, getTracePropagationHeaders()); + headers["x-request-id"] = createUuid(); + return headers; +}; + +const buildHttpError = async (response: Response): Promise => { + let message = ""; + try { + const payload = (await response.json()) as Record; + const detail = payload.detail; + const error = payload.error; + if (typeof detail === "string" && detail.trim()) { + message = detail.trim(); + } else if (typeof error === "string" && error.trim()) { + message = error.trim(); + } + } catch { + // ignore non-json error payloads + } + + const err = new Error( + message || `Request failed with status ${response.status}`, + ) as Error & { status?: number }; + err.status = response.status; + return err; +}; + +const handlePayload = ( + rawPayload: unknown, + chunkKey: "delta" | "chunk", + onChunk: (chunk: string) => void, + onServerMessageId?: (id: string) => void, +) => { + const parsed = ChatStreamChunkSchema.safeParse(rawPayload); + if (!parsed.success) return; + + const payload = parsed.data; + const chunk = payload?.[chunkKey] ?? payload?.delta ?? payload?.chunk; + if (chunk) { + onChunk(chunk); + } + + const serverMessageId = payload?.assistant_message_id || payload?.message_id; + if (serverMessageId && onServerMessageId) { + onServerMessageId(serverMessageId); + } + + if (payload?.error) { + throw new Error(payload.error); + } +}; + +async function streamSSE( + response: Response, + signal: AbortSignal, + onPayload: (payload: unknown) => void, +): Promise { + if (!response.body) throw new Error("Streaming not supported"); + const contentType = response.headers.get("content-type"); + if (contentType && !contentType.includes("text/event-stream")) { + throw new Error(`Unexpected content-type: ${contentType}`); + } + + const reader = response.body.getReader(); + const decoder = new TextDecoder(); + let buffer = ""; + const READ_TIMEOUT_MS = 20_000; + let doneReceived = false; + let timeoutId: ReturnType | undefined; + + const abortHandler = () => { + clearTimeout(timeoutId); + reader.cancel().catch(() => {}); + }; + signal.addEventListener("abort", abortHandler, { once: true }); + + try { + while (true) { + if (signal.aborted) break; + + let result: ReadableStreamReadResult; + try { + result = await Promise.race([ + reader.read(), + new Promise>((_, reject) => { + timeoutId = setTimeout( + () => reject(new Error("Stream read timed out")), + READ_TIMEOUT_MS, + ); + }), + ]); + } finally { + clearTimeout(timeoutId); + } + + const { value, done } = result; + buffer += done + ? decoder.decode(undefined, { stream: false }) + : decoder.decode(value, { stream: true }); + + const { events, remainder } = splitSseEvents(buffer); + buffer = remainder; + + for (const eventBlock of events) { + const parsed = parseSseEvent(eventBlock); + if (!parsed) continue; + if (parsed.event === "heartbeat") continue; + if (parsed.event === "done" || parsed.data === "[DONE]") { + doneReceived = true; + break; + } + let payload: unknown; + try { + payload = JSON.parse(parsed.data); + } catch { + payload = { delta: parsed.data }; + } + onPayload(payload); + } + + if (done || doneReceived) break; + } + } finally { + signal.removeEventListener("abort", abortHandler); + reader.cancel().catch(() => {}); + } +} + +export async function sendChat(params: SendChatParams): Promise { + try { + const headers = await buildHeaders(); + const session = await getSupabaseSession(); + + const fallbackToQueryStream = async () => { + const fallbackLevel = toQueryLevel(params.promptMode); + const fallbackResponse = await fetch(`${API_URL}/api/query/stream`, { + method: "POST", + headers, + signal: params.signal, + body: JSON.stringify({ + topic: params.content, + levels: [fallbackLevel], + mode: params.mode, + premium: params.isPro, + regenerate: Boolean(params.isRegeneration), + bypass_cache: Boolean(params.isRegeneration), + temperature: params.temperature, + message_id: params.clientMessageId, + }), + }); + + if (!fallbackResponse.ok) { + throw await buildHttpError(fallbackResponse); + } + + await streamSSE(fallbackResponse, params.signal, (payload) => + handlePayload( + payload, + "chunk", + params.onChunk, + params.onServerMessageId, + ), + ); + }; + + const shouldUseMessagesEndpoint = + Boolean(session?.access_token) && + supabaseConfigured && + !params.conversationId.startsWith("local-"); + + if (!shouldUseMessagesEndpoint) { + await fallbackToQueryStream(); + params.onDone(); + return; + } + + const response = await fetch(`${API_URL}/api/messages`, { + method: "POST", + headers, + signal: params.signal, + body: JSON.stringify({ + conversation_id: params.conversationId, + content: params.content, + client_generated_id: params.clientMessageId, + assistant_client_id: params.assistantClientId, + mode: params.mode, + prompt_mode: params.promptMode, + regenerate: Boolean(params.isRegeneration), + temperature: params.temperature, + }), + }); + + if (response.status === 404 || response.status === 405) { + await fallbackToQueryStream(); + params.onDone(); + return; + } + + if (!response.ok) { + throw await buildHttpError(response); + } + + await streamSSE(response, params.signal, (payload) => + handlePayload( + payload, + "delta", + params.onChunk, + params.onServerMessageId, + ), + ); + + params.onDone(); + } catch (error) { + params.onError( + error instanceof Error ? error : new Error(String(error)), + ); + } +} diff --git a/src/stores/useChatStore.test.ts b/src/stores/useChatStore.test.ts index 9053811..8a79eca 100644 --- a/src/stores/useChatStore.test.ts +++ b/src/stores/useChatStore.test.ts @@ -1,4 +1,6 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; +import { useConversationStore } from "../stores/useConversationStore"; +import { useMessageStore } from "../stores/useMessageStore"; type MockBuilder = { select: () => MockBuilder; @@ -43,11 +45,27 @@ vi.mock("../lib/supabase", () => ({ import { useChatStore } from "./useChatStore"; -const initialState = useChatStore.getInitialState(); +const resetStore = () => { + useChatStore.setState( + { + theme: "dark", + isSidebarOpen: false, + isPro: false, + gatedModes: ["technical"], + upgradeModalOpen: false, + regenerationModalOpen: false, + regenerationTargetId: null, + regeneratingMessageId: null, + streamControllers: {}, + } as never, + ); + useMessageStore.setState(useMessageStore.getInitialState(), true); + useConversationStore.setState(useConversationStore.getInitialState(), true); +}; describe("useChatStore", () => { beforeEach(() => { - useChatStore.setState(initialState, true); + resetStore(); vi.restoreAllMocks(); }); @@ -58,8 +76,10 @@ describe("useChatStore", () => { it("aborts active streams when starting a new thread", () => { const controller = new AbortController(); - useChatStore.setState({ + useConversationStore.setState({ currentConversationId: "local-conv-1", + }); + useMessageStore.setState({ messageIds: ["assistant-1"], messagesById: { "assistant-1": { @@ -71,8 +91,8 @@ describe("useChatStore", () => { isStreaming: true, }, }, - streamControllers: { "assistant-client-1": controller }, }); + useChatStore.setState({ streamControllers: { "assistant-client-1": controller } }); useChatStore.getState().startNewThread(); @@ -85,7 +105,7 @@ describe("useChatStore", () => { }); it("deletes active conversation and switches to newest remaining conversation", async () => { - useChatStore.setState({ + useConversationStore.setState({ conversations: [ { id: "local-old", @@ -116,7 +136,7 @@ describe("useChatStore", () => { }); it("deletes the last active conversation and starts a new draft thread", async () => { - useChatStore.setState({ + useConversationStore.setState({ conversations: [ { id: "local-only", @@ -152,7 +172,7 @@ describe("useChatStore", () => { ), ); - useChatStore.setState({ + useConversationStore.setState({ currentConversationId: "local-conv-1", conversations: [ { @@ -164,6 +184,8 @@ describe("useChatStore", () => { updated_at: "2026-01-01T00:00:00.000Z", }, ], + }); + useMessageStore.setState({ messageIds: ["user-1", "assistant-1"], messagesById: { "user-1": { @@ -225,7 +247,7 @@ describe("useChatStore", () => { }), ); - useChatStore.setState({ + useConversationStore.setState({ currentConversationId: "local-conv-2", conversations: [ { @@ -237,6 +259,8 @@ describe("useChatStore", () => { updated_at: "2026-01-01T00:00:00.000Z", }, ], + }); + useMessageStore.setState({ messageIds: ["user-2", "assistant-2"], messagesById: { "user-2": { @@ -273,6 +297,8 @@ describe("useChatStore", () => { useChatStore.setState({ regeneratingMessageId: "assistant-locked", + }); + useMessageStore.setState({ messageIds: ["user-3", "assistant-locked"], messagesById: { "user-3": { @@ -321,7 +347,7 @@ const findAssistantMessage = () => { describe("useChatStore streaming", () => { beforeEach(() => { - useChatStore.setState(initialState, true); + resetStore(); vi.restoreAllMocks(); vi.useRealTimers(); }); diff --git a/src/stores/useChatStore.ts b/src/stores/useChatStore.ts index 7b187aa..5b8241e 100644 --- a/src/stores/useChatStore.ts +++ b/src/stores/useChatStore.ts @@ -1,285 +1,45 @@ import { create } from "zustand"; import { supabase } from "../lib/supabase"; -import type { Session } from "@supabase/supabase-js"; -import { splitSseEvents, parseSseEvent } from "../lib/sse"; -import { ChatStreamChunkSchema } from "../lib/sseSchemas"; import type { Level } from "../types"; -import type { - ChatMode, - Conversation, - Message, - PromptMode, -} from "../types/chat"; +import type { ChatMode, Conversation, Message, PromptMode } from "../types/chat"; import { CHAT_PREMIUM_MODES, - resolveChatMode, isModeGated, isPromptMode, - toQueryLevel, } from "../lib/chatModes"; import { captureFrontendError, - getTracePropagationHeaders, trackTelemetry, } from "../lib/monitoring"; - -export type Workspace = "learn" | "socratic" | "technical"; -export type ThemeMode = "dark" | "light"; -export const DEPTH_LEVELS = [ - "eli5", - "eli10", - "eli12", - "eli15", - "meme", -] as const; -export type DepthLevel = (typeof DEPTH_LEVELS)[number]; - -interface ChatState { - conversations: Conversation[]; - currentConversationId: string | null; - isDraftThread: boolean; - workspace: Workspace; - depthLevel: DepthLevel; - theme: ThemeMode; - currentMode: ChatMode; - currentPromptMode: PromptMode; - selectedLevel: Level; - isSidebarOpen: boolean; - messagesById: Record; - messageIds: string[]; - streamControllers: Record; - isLoading: boolean; - isPro: boolean; - gatedModes: ChatMode[]; - upgradeModalOpen: boolean; - regenerationModalOpen: boolean; - regenerationTargetId: string | null; - regeneratingMessageId: string | null; - syncConversations: (conversations: Conversation[]) => void; - selectConversation: (id: string) => Promise; - renameConversation: (id: string, title: string) => Promise; - deleteConversation: (id: string) => Promise; - sendMessage: ( - content: string, - options?: { - mode?: ChatMode; - promptMode?: PromptMode; - isRegeneration?: boolean; - temperature?: number; - clientMessageId?: string; - assistantClientId?: string; - skipUserMessage?: boolean; - replaceMessageId?: string; - }, - ) => Promise; - regenerateMessage: (messageId: string, mode?: ChatMode) => Promise; - retrySync: (messageId: string) => Promise; - setMode: (mode: ChatMode) => void; - setPromptMode: (mode: PromptMode) => void; - setWorkspace: (workspace: Workspace) => void; - setDepthLevel: (level: DepthLevel) => void; - setSelectedLevel: (level: Level) => void; - setTheme: (theme: ThemeMode) => void; - toggleTheme: () => void; - setIsSidebarOpen: (open: boolean) => void; - setIsPro: (isPro: boolean) => void; - startNewThread: () => void; - openUpgradeModal: () => void; - closeUpgradeModal: () => void; - openRegenerationModal: (messageId: string) => void; - closeRegenerationModal: () => void; - abortStream: (clientId: string) => void; - abortAllStreams: () => void; - addMessage: (msg: Message) => void; - updateMessageByClientId: ( - clientId: string, - updater: (msg: Message) => Message, - ) => void; - removeMessageByClientId: (clientId: string) => void; -} - -const supabaseConfigured = - Boolean(import.meta.env.VITE_SUPABASE_URL) && - Boolean(import.meta.env.VITE_SUPABASE_ANON_KEY); -const defaultIsProEnv = import.meta.env.VITE_DEFAULT_IS_PRO; -const defaultIsPro = defaultIsProEnv ? defaultIsProEnv === "true" : false; -const API_URL = import.meta.env.VITE_API_URL || ""; -const THEME_STORAGE_KEY = "kb_theme_v1"; -const DEFAULT_WORKSPACE: Workspace = "learn"; -const DEFAULT_DEPTH_LEVEL: DepthLevel = "eli12"; - -const getSupabaseSession = async (): Promise => { - if (!supabaseConfigured) return null; - const { data } = await supabase.auth.getSession(); - return data.session; -}; - -const isAbortError = (error: unknown): boolean => { - return ( - typeof error === "object" && - error !== null && - "name" in error && - (error as { name?: string }).name === "AbortError" - ); -}; - -const getErrorMessage = (error: unknown, fallback: string): string => { - if (error instanceof Error && error.message) return error.message; - return fallback; -}; - -const isDepthLevel = (mode: string | null | undefined): mode is DepthLevel => { - return DEPTH_LEVELS.includes(mode as DepthLevel); -}; - -const resolveDepthLevel = ( - mode: string | null | undefined, - fallback: DepthLevel = DEFAULT_DEPTH_LEVEL, -): DepthLevel => { - if (isDepthLevel(mode)) return mode; - return fallback; -}; - -const resolveWorkspaceFromMode = (mode: ChatMode): Workspace => { - if (mode === "socratic") return "socratic"; - if (mode === "technical") return "technical"; - return "learn"; -}; - -const resolveWorkspaceState = ( - mode: string | null | undefined, - promptMode: string | null | undefined, - fallbackDepth: DepthLevel, -) => { - const resolvedMode = resolveChatMode(mode); - - if (resolvedMode === "socratic") { - return { - workspace: "socratic" as Workspace, - mode: "socratic" as ChatMode, - promptMode: resolveDepthLevel(promptMode, fallbackDepth) as PromptMode, - depthLevel: resolveDepthLevel(promptMode, fallbackDepth), - }; - } - - if (resolvedMode === "technical") { - return { - workspace: "technical" as Workspace, - mode: "technical" as ChatMode, - promptMode: resolveDepthLevel(promptMode, fallbackDepth) as PromptMode, - depthLevel: resolveDepthLevel(promptMode, fallbackDepth), - }; - } - - const nextDepth = resolveDepthLevel( - promptMode || (isPromptMode(resolvedMode) ? resolvedMode : undefined), - fallbackDepth, - ); - - return { - workspace: "learn" as Workspace, - mode: "learning" as ChatMode, - promptMode: nextDepth as PromptMode, - depthLevel: nextDepth, - }; -}; - -const loadTheme = (): ThemeMode => { - if (typeof window === "undefined") return "light"; - - const cachedTheme = window.localStorage.getItem(THEME_STORAGE_KEY); - if (cachedTheme === "light" || cachedTheme === "dark") { - return cachedTheme; - } - - return window.matchMedia?.("(prefers-color-scheme: dark)").matches - ? "dark" - : "light"; -}; - -const applyThemeClass = (theme: ThemeMode) => { - if (typeof document === "undefined") return; - document.documentElement.classList.toggle("dark", theme === "dark"); -}; - -const persistTheme = (theme: ThemeMode) => { - if (typeof window === "undefined") return; - window.localStorage.setItem(THEME_STORAGE_KEY, theme); -}; - -const getModeForWorkspace = (workspace: Workspace): ChatMode => { - if (workspace === "socratic") return "socratic"; - if (workspace === "technical") return "technical"; - return "learning"; -}; - -const toPersistedConversationMode = ( - mode: ChatMode, - promptMode: PromptMode, -): string => { - // Database constraints in some deployments do not include "learning". - // Persist the concrete prompt mode for learn workspace while keeping - // canonical mode in settings.mode for UI/runtime behavior. - return mode === "learning" ? promptMode : mode; -}; - -const asString = (value: unknown): string | undefined => { - return typeof value === "string" ? value : undefined; -}; - -const initialTheme = loadTheme(); -applyThemeClass(initialTheme); - -const createUuid = () => { - const webCrypto: Crypto | undefined = - typeof globalThis !== "undefined" ? globalThis.crypto : undefined; - - if (webCrypto?.randomUUID) { - return webCrypto.randomUUID(); - } - - const getRandomValues = webCrypto?.getRandomValues - ? webCrypto.getRandomValues.bind(webCrypto) - : null; - const rnd = (size: number) => { - if (getRandomValues) { - const arr = new Uint8Array(size); - getRandomValues(arr); - return arr; - } - return Uint8Array.from({ length: size }, () => - Math.floor(Math.random() * 256), - ); - }; - const bytes = rnd(16); - bytes[6] = (bytes[6] & 0x0f) | 0x40; - bytes[8] = (bytes[8] & 0x3f) | 0x80; - const hex = Array.from(bytes, (b) => b.toString(16).padStart(2, "0")); - return `${hex.slice(0, 4).join("")}-${hex.slice(4, 6).join("")}-${hex - .slice(6, 8) - .join("")}-${hex.slice(8, 10).join("")}-${hex.slice(10, 16).join("")}`; -}; - -const makeLocalId = () => `local-${createUuid()}`; - -const makeClientId = () => createUuid(); - -const truncateTitle = (content: string) => { - const trimmed = content.trim().replace(/\s+/g, " "); - if (trimmed.length <= 64) return trimmed; - return `${trimmed.slice(0, 61)}...`; -}; - -const notifyError = (message: string) => { - console.error(message); - if (typeof window !== "undefined") { - window.dispatchEvent( - new CustomEvent("kb-toast", { detail: { type: "error", message } }), - ); - } -}; - -const PENDING_SYNC_KEY = "kb_pending_sync_v1"; +import { sendChat } from "../services/chatService"; +import { + makeLocalId, + makeClientId, + truncateTitle, + notifyError, + isAbortError, + getErrorMessage, + resolveDepthLevel, + resolveWorkspaceFromMode, + getModeForWorkspace, + supabaseConfigured, + defaultIsPro, + DEPTH_LEVELS, + PENDING_SYNC_KEY, + loadTheme, + applyThemeClass, + persistTheme, + type Workspace, + type ThemeMode, + type DepthLevel, +} from "../lib/chatStoreUtils"; +import { useMessageStore } from "./useMessageStore"; +import { useConversationStore } from "./useConversationStore"; + +export type { Workspace, ThemeMode, DepthLevel }; +export { DEPTH_LEVELS }; + +// ─── Pending sync helpers (unchanged from original) ──────────────────────── interface PendingSyncEntry { id: string; @@ -317,253 +77,123 @@ const savePendingSyncs = (entries: PendingSyncEntry[]) => { const cachePendingSync = (entry: PendingSyncEntry) => { const existing = loadPendingSyncs(); - const next = [ - entry, - ...existing.filter((item) => item.id !== entry.id), - ].slice(0, 50); + const next = [entry, ...existing.filter((i) => i.id !== entry.id)].slice(0, 50); savePendingSyncs(next); }; const removePendingSync = (id: string) => { - const existing = loadPendingSyncs(); - const next = existing.filter((item) => item.id !== id); - savePendingSyncs(next); + savePendingSyncs(loadPendingSyncs().filter((i) => i.id !== id)); }; -const resolveMessageKey = (message: Message) => { - return ( - message.clientGeneratedId || - message.metadata?.assistant_client_id || - message.metadata?.client_id || - message.serverMessageId || - message.id - ); -}; +// ─── Initial theme ───────────────────────────────────────────────────────── -const messagesMatch = (existing: Message, incoming: Message) => { - if (existing.id === incoming.id) return true; - if ( - existing.clientGeneratedId && - incoming.clientGeneratedId && - existing.clientGeneratedId === incoming.clientGeneratedId - ) { - return true; - } - if ( - incoming.metadata?.assistant_client_id && - existing.clientGeneratedId === incoming.metadata.assistant_client_id - ) { - return true; - } - if ( - existing.metadata?.assistant_client_id && - incoming.clientGeneratedId && - existing.metadata.assistant_client_id === incoming.clientGeneratedId - ) { - return true; - } - if ( - existing.serverMessageId && - incoming.id && - existing.serverMessageId === incoming.id - ) - return true; - if ( - incoming.serverMessageId && - existing.id && - incoming.serverMessageId === existing.id - ) - return true; - if ( - incoming.metadata?.client_id && - existing.id === incoming.metadata.client_id - ) - return true; - if ( - existing.metadata?.client_id && - existing.metadata.client_id === incoming.id - ) - return true; - if ( - incoming.metadata?.client_id && - existing.clientGeneratedId === incoming.metadata.client_id - ) - return true; - if ( - existing.metadata?.client_id && - incoming.clientGeneratedId && - existing.metadata.client_id === incoming.clientGeneratedId - ) { - return true; - } - if ( - incoming.metadata?.client_id && - existing.metadata?.client_id && - existing.metadata.client_id === incoming.metadata.client_id - ) { - return true; - } - return false; -}; +const initialTheme = loadTheme(); +applyThemeClass(initialTheme); -const findExistingMessageKey = ( - state: Pick, - incoming: Message, -) => { - for (const messageKey of state.messageIds) { - const existing = state.messagesById[messageKey]; - if (!existing) continue; - if (messagesMatch(existing, incoming)) { - return messageKey; - } - } - return null; -}; +// ─── Store interface ──────────────────────────────────────────────────────── + +interface ChatState { + // UI state + theme: ThemeMode; + isSidebarOpen: boolean; + isPro: boolean; + gatedModes: ChatMode[]; + upgradeModalOpen: boolean; + regenerationModalOpen: boolean; + regenerationTargetId: string | null; + regeneratingMessageId: string | null; + streamControllers: Record; -const buildMessageRegistry = (messages: Message[]) => { - const messagesById: Record = {}; - const messageIds: string[] = []; + // Proxy selectors (read-through to sub-stores for backwards compatibility) + readonly conversations: Conversation[]; + readonly currentConversationId: string | null; + readonly isDraftThread: boolean; + readonly isLoading: boolean; + readonly workspace: Workspace; + readonly depthLevel: DepthLevel; + readonly currentMode: ChatMode; + readonly currentPromptMode: PromptMode; + readonly selectedLevel: Level; + readonly messagesById: Record; + readonly messageIds: string[]; + + // Actions — UI + setTheme: (theme: ThemeMode) => void; + toggleTheme: () => void; + setIsSidebarOpen: (open: boolean) => void; + setIsPro: (isPro: boolean) => void; + openUpgradeModal: () => void; + closeUpgradeModal: () => void; + openRegenerationModal: (messageId: string) => void; + closeRegenerationModal: () => void; + abortStream: (clientId: string) => void; + abortAllStreams: () => void; - for (const message of messages) { - const key = resolveMessageKey(message); - if (messagesById[key]) { - messagesById[key] = { ...messagesById[key], ...message }; - continue; - } - messagesById[key] = message; - messageIds.push(key); - } + // Actions — workspace/mode (delegate to conversation store) + setMode: (mode: ChatMode) => void; + setPromptMode: (mode: PromptMode) => void; + setWorkspace: (workspace: Workspace) => void; + setDepthLevel: (level: DepthLevel) => void; + setSelectedLevel: (level: Level) => void; - return { messagesById, messageIds }; -}; + // Actions — conversations (delegate to conversation store) + syncConversations: (conversations: Conversation[]) => void; + selectConversation: (id: string) => Promise; + renameConversation: (id: string, title: string) => Promise; + deleteConversation: (id: string) => Promise; + startNewThread: () => void; + + // Actions — messages (delegate to message store) + addMessage: (msg: Message) => void; + updateMessageByClientId: (clientId: string, updater: (msg: Message) => Message) => void; + removeMessageByClientId: (clientId: string) => void; + + // Actions — streaming + sendMessage: ( + content: string, + options?: { + mode?: ChatMode; + promptMode?: PromptMode; + isRegeneration?: boolean; + temperature?: number; + clientMessageId?: string; + assistantClientId?: string; + skipUserMessage?: boolean; + replaceMessageId?: string; + }, + ) => Promise; + regenerateMessage: (messageId: string, mode?: ChatMode) => Promise; + retrySync: (messageId: string) => Promise; +} + +// ─── Store implementation ─────────────────────────────────────────────────── export const useChatStore = create((set, get) => ({ - conversations: [], - currentConversationId: null, - isDraftThread: false, - workspace: DEFAULT_WORKSPACE, - depthLevel: DEFAULT_DEPTH_LEVEL, + // UI state theme: initialTheme, - currentMode: "learning", - currentPromptMode: DEFAULT_DEPTH_LEVEL as PromptMode, - selectedLevel: DEFAULT_DEPTH_LEVEL as Level, isSidebarOpen: false, - messagesById: {}, - messageIds: [], - streamControllers: {}, - isLoading: false, isPro: defaultIsPro, gatedModes: [...CHAT_PREMIUM_MODES], upgradeModalOpen: false, regenerationModalOpen: false, regenerationTargetId: null, regeneratingMessageId: null, + streamControllers: {}, - setMode: (mode: ChatMode) => { - const { currentConversationId, conversations, depthLevel } = get(); - const conversation = conversations.find( - (item) => item.id === currentConversationId, - ); - const nextPromptMode = isPromptMode(mode) ? mode : get().currentPromptMode; - const nextDepthLevel = resolveDepthLevel(nextPromptMode, depthLevel); - const nextWorkspace = resolveWorkspaceFromMode(mode); - const nextSettings = conversation?.settings - ? { ...conversation.settings, mode, prompt_mode: nextPromptMode } - : conversation - ? { mode, prompt_mode: nextPromptMode } - : undefined; - const persistedMode = toPersistedConversationMode(mode, nextPromptMode); - - set((state) => ({ - workspace: nextWorkspace, - depthLevel: nextDepthLevel, - currentMode: mode, - currentPromptMode: nextPromptMode, - selectedLevel: nextDepthLevel as Level, - conversations: currentConversationId - ? state.conversations.map((item) => - item.id === currentConversationId - ? { ...item, mode, settings: nextSettings ?? item.settings } - : item, - ) - : state.conversations, - })); - - if ( - !supabaseConfigured || - !currentConversationId || - !conversation || - currentConversationId.startsWith("local-") - ) { - return; - } - - void supabase - .from("conversations") - .update({ mode: persistedMode, settings: nextSettings ?? conversation.settings }) - .eq("id", currentConversationId); - }, - - setPromptMode: (mode: PromptMode) => { - const { currentConversationId, conversations } = get(); - const conversation = conversations.find( - (item) => item.id === currentConversationId, - ); - const nextDepthLevel = resolveDepthLevel(mode, get().depthLevel); - const nextSettings = conversation?.settings - ? { ...conversation.settings, prompt_mode: mode } - : conversation - ? { prompt_mode: mode } - : undefined; - - set((state) => ({ - currentPromptMode: mode, - depthLevel: nextDepthLevel, - selectedLevel: nextDepthLevel as Level, - conversations: currentConversationId - ? state.conversations.map((item) => - item.id === currentConversationId - ? { ...item, settings: nextSettings ?? item.settings } - : item, - ) - : state.conversations, - })); - - if ( - !supabaseConfigured || - !currentConversationId || - !conversation || - currentConversationId.startsWith("local-") - ) { - return; - } - - void supabase - .from("conversations") - .update({ settings: nextSettings ?? conversation.settings }) - .eq("id", currentConversationId); - }, - - setWorkspace: (workspace: Workspace) => { - get().setMode(getModeForWorkspace(workspace)); - if (workspace === "learn") { - const nextDepth = get().depthLevel; - get().setPromptMode(nextDepth as PromptMode); - } - }, - - setDepthLevel: (level: DepthLevel) => { - set({ - depthLevel: level, - selectedLevel: level as Level, - }); - get().setPromptMode(level as PromptMode); - if (get().workspace === "learn") { - get().setMode("learning"); - } - }, - - setSelectedLevel: (selectedLevel: Level) => set({ selectedLevel }), + // Proxy selectors — read directly from sub-stores + get conversations() { return useConversationStore.getState().conversations; }, + get currentConversationId() { return useConversationStore.getState().currentConversationId; }, + get isDraftThread() { return useConversationStore.getState().isDraftThread; }, + get isLoading() { return useConversationStore.getState().isLoading; }, + get workspace() { return useConversationStore.getState().workspace; }, + get depthLevel() { return useConversationStore.getState().depthLevel; }, + get currentMode() { return useConversationStore.getState().currentMode; }, + get currentPromptMode() { return useConversationStore.getState().currentPromptMode; }, + get selectedLevel() { return useConversationStore.getState().selectedLevel; }, + get messagesById() { return useMessageStore.getState().messagesById; }, + get messageIds() { return useMessageStore.getState().messageIds; }, + + // ── UI actions ──────────────────────────────────────────────────────────── setTheme: (theme: ThemeMode) => { applyThemeClass(theme); @@ -572,360 +202,205 @@ export const useChatStore = create((set, get) => ({ }, toggleTheme: () => { - const nextTheme: ThemeMode = get().theme === "dark" ? "light" : "dark"; - get().setTheme(nextTheme); - }, - - setIsSidebarOpen: (isSidebarOpen: boolean) => set({ isSidebarOpen }), - - setIsPro: (isPro: boolean) => set({ isPro }), - - startNewThread: () => { - const { workspace, depthLevel } = get(); - get().abortAllStreams(); - set({ - currentConversationId: null, - isDraftThread: true, - currentMode: getModeForWorkspace(workspace), - currentPromptMode: depthLevel as PromptMode, - selectedLevel: depthLevel as Level, - messagesById: {}, - messageIds: [], - isLoading: false, - }); + const next: ThemeMode = get().theme === "dark" ? "light" : "dark"; + get().setTheme(next); }, + setIsSidebarOpen: (isSidebarOpen) => set({ isSidebarOpen }), + setIsPro: (isPro) => set({ isPro }), openUpgradeModal: () => set({ upgradeModalOpen: true }), closeUpgradeModal: () => set({ upgradeModalOpen: false }), - openRegenerationModal: (messageId: string) => + openRegenerationModal: (messageId) => set({ regenerationModalOpen: true, regenerationTargetId: messageId }), closeRegenerationModal: () => set({ regenerationModalOpen: false, regenerationTargetId: null }), + abortStream: (clientId: string) => { const controller = get().streamControllers[clientId]; if (controller) controller.abort(); set((state) => { - const { [clientId]: removedController, ...rest } = - state.streamControllers; - void removedController; - return { streamControllers: rest }; + return { + streamControllers: Object.fromEntries( + Object.entries(state.streamControllers).filter(([key]) => key !== clientId), + ), + }; }); - get().updateMessageByClientId(clientId, (message) => ({ - ...message, + useMessageStore.getState().updateMessageByClientId(clientId, (msg) => ({ + ...msg, isStreaming: false, error: "Canceled", })); - const stillStreaming = get().messageIds.some( - (id) => get().messagesById[id]?.isStreaming, - ); - set({ isLoading: stillStreaming }); + const stillStreaming = useMessageStore + .getState() + .messageIds.some((id) => useMessageStore.getState().messagesById[id]?.isStreaming); + useConversationStore.getState().setIsLoading(stillStreaming); }, + abortAllStreams: () => { const controllers = get().streamControllers; - Object.values(controllers).forEach((controller) => controller.abort()); + Object.values(controllers).forEach((c) => c.abort()); set({ streamControllers: {} }); - set((state) => { - const messagesById = { ...state.messagesById }; - for (const id of state.messageIds) { - const message = messagesById[id]; - if (message?.isStreaming) { - messagesById[id] = { - ...message, - isStreaming: false, - error: "Canceled", - }; - } + const { messagesById, messageIds } = useMessageStore.getState(); + const updated = { ...messagesById }; + for (const id of messageIds) { + if (updated[id]?.isStreaming) { + updated[id] = { ...updated[id], isStreaming: false, error: "Canceled" }; } - return { messagesById }; - }); - set({ isLoading: false }); - }, - - syncConversations: (conversations: Conversation[]) => { - set((state) => { - if (state.isDraftThread && state.currentConversationId === null) { - return { conversations }; - } - - const preferredId = state.currentConversationId; - const hasPreferred = preferredId - ? conversations.some((item) => item.id === preferredId) - : false; - const nextConversationId = hasPreferred - ? preferredId - : (conversations[0]?.id ?? null); - const activeConversation = conversations.find( - (item) => item.id === nextConversationId, - ); - const conversationMode = - asString(activeConversation?.mode) || - asString(activeConversation?.settings?.mode); - const conversationPrompt = - asString(activeConversation?.settings?.prompt_mode) || - asString(activeConversation?.settings?.mode) || - asString(activeConversation?.mode) || - state.currentPromptMode; - const nextWorkspaceState = resolveWorkspaceState( - conversationMode, - conversationPrompt, - state.depthLevel, - ); - return { - conversations, - currentConversationId: nextConversationId, - isDraftThread: false, - workspace: nextWorkspaceState.workspace, - depthLevel: nextWorkspaceState.depthLevel, - currentMode: nextWorkspaceState.mode, - currentPromptMode: nextWorkspaceState.promptMode, - selectedLevel: nextWorkspaceState.depthLevel as Level, - }; - }); - }, - - selectConversation: async (id: string) => { - if (!id) return; - const state = get(); - - if ( - state.currentConversationId === id && - (state.isLoading || state.messageIds.length > 0) - ) { - return; - } - - const activeConversation = state.conversations.find( - (item) => item.id === id, - ); - const conversationMode = - asString(activeConversation?.mode) || - asString(activeConversation?.settings?.mode) || - state.currentMode; - const conversationPrompt = - asString(activeConversation?.settings?.prompt_mode) || - asString(activeConversation?.settings?.mode) || - asString(activeConversation?.mode) || - state.currentPromptMode; - const nextWorkspaceState = resolveWorkspaceState( - conversationMode, - conversationPrompt, - state.depthLevel, - ); - set({ - currentConversationId: id, - isDraftThread: false, - messagesById: {}, - messageIds: [], - isLoading: true, - workspace: nextWorkspaceState.workspace, - depthLevel: nextWorkspaceState.depthLevel, - currentMode: nextWorkspaceState.mode, - currentPromptMode: nextWorkspaceState.promptMode, - selectedLevel: nextWorkspaceState.depthLevel as Level, - }); - - if (!supabaseConfigured) { - set({ isLoading: false }); - return; - } - - try { - const { data, error } = await supabase - .from("messages") - .select("id, role, content, attachments, metadata, created_at") - .eq("conversation_id", id) - .order("created_at", { ascending: true }); - - if (error) throw error; - - const { messagesById, messageIds } = buildMessageRegistry( - (data ?? []) as Message[], - ); - set({ messagesById, messageIds }); - } catch (error) { - console.error("Failed to fetch messages:", error); - } finally { - set({ isLoading: false }); } + useMessageStore.setState({ messagesById: updated }); + useConversationStore.getState().setIsLoading(false); }, - renameConversation: async (id: string, title: string) => { - if (!id) return; - const trimmed = title.trim(); - if (!trimmed) return; + // ── Workspace/mode actions ──────────────────────────────────────────────── - const now = new Date().toISOString(); - set((state) => ({ - conversations: state.conversations - .map((item) => - item.id === id ? { ...item, title: trimmed, updated_at: now } : item, - ) - .sort((a, b) => (a.updated_at < b.updated_at ? 1 : -1)), + setMode: (mode: ChatMode) => { + const convStore = useConversationStore.getState(); + const { currentConversationId, conversations, depthLevel } = convStore; + const conversation = conversations.find((c) => c.id === currentConversationId); + const nextPromptMode = isPromptMode(mode) ? mode : convStore.currentPromptMode; + const nextDepthLevel = resolveDepthLevel(nextPromptMode, depthLevel); + const nextWorkspace = resolveWorkspaceFromMode(mode); + const nextSettings = conversation?.settings + ? { ...conversation.settings, mode, prompt_mode: nextPromptMode } + : conversation + ? { mode, prompt_mode: nextPromptMode } + : undefined; + + convStore.setWorkspaceState(nextWorkspace, mode, nextPromptMode, nextDepthLevel); + useConversationStore.setState((state) => ({ + ...state, + selectedLevel: nextDepthLevel as Level, })); - if (!supabaseConfigured || id.startsWith("local-")) return; + if (currentConversationId && conversation && nextSettings) { + useConversationStore.setState((state) => ({ + conversations: state.conversations.map((c) => + c.id === currentConversationId + ? { ...c, mode, settings: nextSettings } + : c, + ), + })); - try { - await supabase - .from("conversations") - .update({ title: trimmed, updated_at: now }) - .eq("id", id); - } catch (error) { - console.error("Failed to rename conversation:", error); + if (supabaseConfigured && !currentConversationId.startsWith("local-")) { + void supabase + .from("conversations") + .update({ mode, settings: nextSettings }) + .eq("id", currentConversationId); + } } }, - deleteConversation: async (id: string) => { - if (!id) return; + setPromptMode: (mode: PromptMode) => { + const convStore = useConversationStore.getState(); + const { currentConversationId, conversations, depthLevel, workspace, currentMode } = convStore; + const conversation = conversations.find((c) => c.id === currentConversationId); + const nextDepthLevel = resolveDepthLevel(mode, depthLevel); + const nextSettings = conversation?.settings + ? { ...conversation.settings, prompt_mode: mode } + : conversation + ? { prompt_mode: mode } + : undefined; - const initialState = get(); - const targetConversation = initialState.conversations.find( - (conversation) => conversation.id === id, - ); - if (!targetConversation) return; - const isActiveConversation = initialState.currentConversationId === id; + convStore.setWorkspaceState(workspace, currentMode, mode, nextDepthLevel); - if (isActiveConversation) { - get().abortAllStreams(); - } + if (currentConversationId && conversation && nextSettings) { + useConversationStore.setState((state) => ({ + conversations: state.conversations.map((c) => + c.id === currentConversationId + ? { ...c, settings: nextSettings } + : c, + ), + })); - if (!id.startsWith("local-") && supabaseConfigured) { - try { - const { error } = await supabase + if (supabaseConfigured && !currentConversationId.startsWith("local-")) { + void supabase .from("conversations") - .delete() - .eq("id", id); - if (error) throw error; - } catch (error) { - console.error("Failed to delete conversation:", error); - notifyError("Failed to delete conversation."); - return; + .update({ settings: nextSettings }) + .eq("id", currentConversationId); } } + }, - const state = get(); - const remainingConversations = state.conversations - .filter((conversation) => conversation.id !== id) - .sort((a, b) => (a.updated_at < b.updated_at ? 1 : -1)); - - if (!isActiveConversation) { - set({ conversations: remainingConversations }); - return; + setWorkspace: (workspace: Workspace) => { + get().setMode(getModeForWorkspace(workspace)); + if (workspace === "learn") { + get().setPromptMode(useConversationStore.getState().depthLevel as PromptMode); } + }, - if (remainingConversations.length === 0) { - set({ conversations: remainingConversations }); - get().startNewThread(); - return; + setDepthLevel: (level: DepthLevel) => { + useConversationStore.getState().setWorkspaceState( + useConversationStore.getState().workspace, + useConversationStore.getState().currentMode, + level as PromptMode, + level, + ); + get().setPromptMode(level as PromptMode); + if (useConversationStore.getState().workspace === "learn") { + get().setMode("learning"); } + }, - const nextConversationId = remainingConversations[0].id; - set({ - conversations: remainingConversations, - currentConversationId: nextConversationId, - isDraftThread: false, - messagesById: {}, - messageIds: [], - isLoading: false, - }); - await get().selectConversation(nextConversationId); + setSelectedLevel: (selectedLevel: Level) => { + useConversationStore.setState((state) => ({ ...state, selectedLevel })); }, - addMessage: (msg: Message) => { - set((state) => { - const resolvedKey = resolveMessageKey(msg); - const directMatch = state.messagesById[resolvedKey] ? resolvedKey : null; - const existingKey = directMatch || findExistingMessageKey(state, msg); + // ── Conversation delegates ──────────────────────────────────────────────── - if (existingKey) { - const nextMessagesById = { - ...state.messagesById, - [existingKey]: { ...state.messagesById[existingKey], ...msg }, - }; - if (!state.messageIds.includes(existingKey)) { - return { - messagesById: nextMessagesById, - messageIds: [...state.messageIds, existingKey], - }; - } - return { messagesById: nextMessagesById }; - } + syncConversations: (conversations) => + useConversationStore.getState().syncConversations(conversations), - return { - messagesById: { ...state.messagesById, [resolvedKey]: msg }, - messageIds: [...state.messageIds, resolvedKey], - }; - }); - }, + selectConversation: (id) => + useConversationStore.getState().selectConversation(id), - updateMessageByClientId: ( - clientId: string, - updater: (msg: Message) => Message, - ) => { - set((state) => { - const messageKey = state.messageIds.find((id) => { - const message = state.messagesById[id]; - if (!message) return false; - return ( - message.clientGeneratedId === clientId || - message.metadata?.assistant_client_id === clientId || - message.metadata?.client_id === clientId - ); - }); + renameConversation: (id, title) => + useConversationStore.getState().renameConversation(id, title), - if (!messageKey) return state; - const message = state.messagesById[messageKey]; - return { - messagesById: { - ...state.messagesById, - [messageKey]: updater(message), - }, - }; + deleteConversation: async (id) => { + get().abortAllStreams(); + await useConversationStore.getState().deleteConversation(id); + }, + + startNewThread: () => { + const { workspace, depthLevel } = useConversationStore.getState(); + get().abortAllStreams(); + useMessageStore.getState().clearMessages(); + useConversationStore.setState({ + currentConversationId: null, + isDraftThread: true, + currentMode: getModeForWorkspace(workspace), + currentPromptMode: depthLevel as PromptMode, + selectedLevel: depthLevel as Level, + isLoading: false, }); }, - removeMessageByClientId: (clientId: string) => { - set((state) => { - const messageKey = state.messageIds.find((id) => { - const message = state.messagesById[id]; - if (!message) return false; - return ( - message.clientGeneratedId === clientId || - message.metadata?.assistant_client_id === clientId || - message.metadata?.client_id === clientId - ); - }); + // ── Message delegates ───────────────────────────────────────────────────── - if (!messageKey) return state; + addMessage: (msg) => useMessageStore.getState().addMessage(msg), + updateMessageByClientId: (clientId, updater) => + useMessageStore.getState().updateMessageByClientId(clientId, updater), + removeMessageByClientId: (clientId) => + useMessageStore.getState().removeMessageByClientId(clientId), - const { [messageKey]: removedMessage, ...rest } = state.messagesById; - void removedMessage; - return { - messagesById: rest, - messageIds: state.messageIds.filter((id) => id !== messageKey), - }; - }); - }, + // ── Streaming ───────────────────────────────────────────────────────────── - sendMessage: async ( - content: string, - options?: { - mode?: ChatMode; - promptMode?: PromptMode; - isRegeneration?: boolean; - temperature?: number; - clientMessageId?: string; - assistantClientId?: string; - skipUserMessage?: boolean; - replaceMessageId?: string; - }, - ) => { + sendMessage: async (content, options) => { const trimmed = content.trim(); if (!trimmed) return; - const { currentMode, currentPromptMode, isPro, gatedModes } = get(); + const convStore = useConversationStore.getState(); + const msgStore = useMessageStore.getState(); + const { currentMode, currentPromptMode, isPro, gatedModes } = { + currentMode: convStore.currentMode, + currentPromptMode: convStore.currentPromptMode, + isPro: get().isPro, + gatedModes: get().gatedModes, + }; + const requestedMode = options?.mode ?? currentMode; const requestedPromptMode = options?.promptMode ?? currentPromptMode; + if (isModeGated(requestedMode, isPro, gatedModes)) { get().openUpgradeModal(); return; @@ -936,24 +411,18 @@ export const useChatStore = create((set, get) => ({ const clientMessageId = options?.clientMessageId ?? makeClientId(); const assistantClientId = options?.assistantClientId ?? makeClientId(); const skipUserMessage = Boolean(options?.skipUserMessage); - const requestTemperature = Math.min( - Math.max(options?.temperature ?? 0.7, 0), - 1, - ); - let conversationId = get().currentConversationId; - let conversation = get().conversations.find( - (item) => item.id === conversationId, - ); + const requestTemperature = Math.min(Math.max(options?.temperature ?? 0.7, 0), 1); + + let conversationId = convStore.currentConversationId; + let conversation = convStore.conversations.find((c) => c.id === conversationId); const effectivePromptMode = isPromptMode(requestedMode) ? requestedMode : requestedPromptMode; - const persistedConversationMode = toPersistedConversationMode( - requestedMode, - effectivePromptMode, - ); - set({ isLoading: true, isDraftThread: false }); + convStore.setIsLoading(true); + convStore.setIsDraftThread(false); + // Create conversation if needed if (!conversationId && !skipUserMessage) { const title = truncateTitle(trimmed); if (supabaseConfigured) { @@ -965,31 +434,21 @@ export const useChatStore = create((set, get) => ({ .insert({ user_id: authData.user.id, title, - mode: persistedConversationMode, - settings: { - mode: requestedMode, - prompt_mode: effectivePromptMode, - }, + mode: requestedMode, + settings: { mode: requestedMode, prompt_mode: effectivePromptMode }, }) .select("id, title, mode, settings, created_at, updated_at") .single(); - if (error) throw error; - if (data) { conversation = data as Conversation; conversationId = data.id; - set((state) => ({ - conversations: [ - conversation as Conversation, - ...state.conversations, - ], - currentConversationId: conversationId, - })); + useConversationStore.getState().upsertConversation(conversation); + convStore.setCurrentConversationId(conversationId); } } - } catch (error) { - console.error("Failed to create conversation:", error); + } catch (err) { + console.error("Failed to create conversation:", err); } } @@ -997,82 +456,78 @@ export const useChatStore = create((set, get) => ({ conversationId = makeLocalId(); conversation = { id: conversationId, - title, + title: truncateTitle(trimmed), mode: requestedMode, settings: { mode: requestedMode, prompt_mode: effectivePromptMode }, created_at: now, updated_at: now, }; - set((state) => ({ - conversations: [conversation as Conversation, ...state.conversations], - currentConversationId: conversationId, - })); + useConversationStore.getState().upsertConversation(conversation); + convStore.setCurrentConversationId(conversationId); } } if (!conversationId) { notifyError("No active conversation available."); - set({ isLoading: false }); + convStore.setIsLoading(false); return; } - const existingUserMessageId = get().messageIds.find((id) => { - const message = get().messagesById[id]; - if (!message) return false; - return ( - message.clientGeneratedId === clientMessageId || - message.metadata?.client_id === clientMessageId - ); - }); - - if (!skipUserMessage && !existingUserMessageId) { - const optimisticUserMessage: Message = { - id: localUserId, - role: "user", - content: trimmed, - metadata: { - client_id: clientMessageId, - mode: requestedMode, - prompt_mode: effectivePromptMode, - }, - created_at: now, - clientGeneratedId: clientMessageId, - }; - - get().addMessage(optimisticUserMessage); + // Optimistic user message + if (!skipUserMessage) { + const existingUserMessageId = msgStore.messageIds.find((id) => { + const msg = msgStore.messagesById[id]; + return ( + msg?.clientGeneratedId === clientMessageId || + msg?.metadata?.client_id === clientMessageId + ); + }); + if (!existingUserMessageId) { + msgStore.addMessage({ + id: localUserId, + role: "user", + content: trimmed, + metadata: { + client_id: clientMessageId, + mode: requestedMode, + prompt_mode: effectivePromptMode, + }, + created_at: now, + clientGeneratedId: clientMessageId, + }); + } } - set((state) => ({ + + // Update conversation updated_at + const updatedNow = new Date().toISOString(); + useConversationStore.setState((state) => ({ conversations: state.conversations - .map((item) => - item.id === conversationId - ? { - ...item, - title: item.title || truncateTitle(trimmed), - updated_at: now, - } - : item, + .map((c) => + c.id === conversationId + ? { ...c, title: c.title || truncateTitle(trimmed), updated_at: updatedNow } + : c, ) .sort((a, b) => (a.updated_at < b.updated_at ? 1 : -1)), })); - const existingAssistantMessageId = - options?.replaceMessageId && get().messagesById[options.replaceMessageId] + // Optimistic assistant placeholder + const existingAssistantId = + options?.replaceMessageId && msgStore.messagesById[options.replaceMessageId] ? options.replaceMessageId - : get().messageIds.find((id) => { - const message = get().messagesById[id]; - if (!message) return false; + : msgStore.messageIds.find((id) => { + const msg = msgStore.messagesById[id]; return ( - message.clientGeneratedId === assistantClientId || - message.metadata?.assistant_client_id === assistantClientId + msg?.clientGeneratedId === assistantClientId || + msg?.metadata?.assistant_client_id === assistantClientId ); }); - if (existingAssistantMessageId) { - set((state) => ({ + if (existingAssistantId) { + useMessageStore.setState((state) => ({ messagesById: { ...state.messagesById, - [existingAssistantMessageId]: { - ...state.messagesById[existingAssistantMessageId], + [existingAssistantId]: { + ...state.messagesById[existingAssistantId], content: "", isStreaming: true, isRegenerating: Boolean(options?.isRegeneration), @@ -1080,7 +535,7 @@ export const useChatStore = create((set, get) => ({ syncStatus: "pending", clientGeneratedId: assistantClientId, metadata: { - ...state.messagesById[existingAssistantMessageId]?.metadata, + ...state.messagesById[existingAssistantId]?.metadata, mode: requestedMode, prompt_mode: effectivePromptMode, temperature: requestTemperature, @@ -1090,7 +545,7 @@ export const useChatStore = create((set, get) => ({ }, })); } else { - const assistantPlaceholder: Message = { + msgStore.addMessage({ id: makeLocalId(), role: "assistant", content: "", @@ -1105,472 +560,77 @@ export const useChatStore = create((set, get) => ({ temperature: requestTemperature, assistant_client_id: assistantClientId, }, - }; - - get().addMessage(assistantPlaceholder); + }); } const controller = new AbortController(); const streamStartedAt = Date.now(); - let streamStarted = false; - let usedQueryFallback = false; + set((state) => ({ - streamControllers: { - ...state.streamControllers, - [assistantClientId]: controller, - }, + streamControllers: { ...state.streamControllers, [assistantClientId]: controller }, })); trackTelemetry("message_send", { mode: requestedMode, prompt_mode: effectivePromptMode, regenerate: Boolean(options?.isRegeneration), - retry: Boolean(options?.clientMessageId), }); - try { - const session = await getSupabaseSession(); - const headers: Record = { - "Content-Type": "application/json", - }; - if (session?.access_token) { - headers["Authorization"] = `Bearer ${session.access_token}`; - } - Object.assign(headers, getTracePropagationHeaders()); - headers["x-request-id"] = createUuid(); - - const streamFromResponse = async ( - response: Response, - handler: (payload: unknown) => void, - signal: AbortSignal, - ): Promise => { - if (!response.body) { - throw new Error("Streaming not supported in this environment"); - } - - const contentType = response.headers.get("content-type"); - if (contentType && !contentType.includes("text/event-stream")) { - throw new Error(`Unexpected content-type: ${contentType}`); - } - - const reader = response.body.getReader(); - const decoder = new TextDecoder(); - let buffer = ""; - const READ_TIMEOUT_MS = 20_000; - let doneReceived = false; - let timeoutId: ReturnType | undefined; - - const abortHandler = () => { - clearTimeout(timeoutId); - reader.cancel().catch(() => {}); - }; - signal.addEventListener("abort", abortHandler, { once: true }); - - try { - while (true) { - if (signal.aborted) break; - - let result: ReadableStreamReadResult; - try { - result = await Promise.race([ - reader.read(), - new Promise>((_, reject) => { - timeoutId = setTimeout( - () => reject(new Error("Stream read timed out")), - READ_TIMEOUT_MS, - ); - }), - ]); - } finally { - clearTimeout(timeoutId); - } - - const { value, done } = result; - - buffer += done - ? decoder.decode(undefined, { stream: false }) - : decoder.decode(value, { stream: true }); - - const { events, remainder } = splitSseEvents(buffer); - buffer = remainder; - - for (const eventBlock of events) { - const parsed = parseSseEvent(eventBlock); - if (!parsed) { - console.warn("[sse] Skipping invalid SSE event:", eventBlock); - continue; - } - if (parsed.event === "heartbeat") continue; - if (parsed.event === "done" || parsed.data === "[DONE]") { - doneReceived = true; - break; - } - - let payload: unknown; - try { - payload = JSON.parse(parsed.data); - } catch { - payload = { delta: parsed.data }; - } - handler(payload); - } - - if (done || doneReceived) break; - } - } finally { - signal.removeEventListener("abort", abortHandler); - reader.cancel().catch(() => {}); - } - }; - - const buildHttpError = async (response: Response) => { - let message = ""; - try { - const payload = (await response.json()) as Record; - if (payload && typeof payload === "object") { - const detail = payload.detail; - const error = payload.error; - if (typeof detail === "string" && detail.trim()) { - message = detail.trim(); - } else if (typeof error === "string" && error.trim()) { - message = error.trim(); - } - } - } catch { - // ignored: non-JSON error responses - } - - if (response.status === 429) { - const retryAfter = response.headers.get("retry-after"); - const suffix = retryAfter - ? ` Retry after ${retryAfter} seconds.` - : ""; - const err = new Error( - `${message || "You are sending requests too quickly."}${suffix}`, - ) as Error & { status?: number }; - err.status = 429; - return err; - } - - const err = new Error( - message || `Request failed with status ${response.status}`, - ) as Error & { status?: number }; - err.status = response.status; - return err; - }; - - const handleStreamingPayload = ( - rawPayload: unknown, - chunkKey: "delta" | "chunk", - ) => { - const parsed = ChatStreamChunkSchema.safeParse(rawPayload); - if (!parsed.success) { - console.warn("Skipping invalid SSE payload:", parsed.error); - return; - } - - const payload = parsed.data; - const chunk = payload?.[chunkKey] ?? payload?.delta ?? payload?.chunk; - if (chunk) { - get().updateMessageByClientId(assistantClientId, (message) => ({ - ...message, - content: `${message.content}${chunk}`, - })); - } - - const serverMessageId = - payload?.assistant_message_id || payload?.message_id; - if (serverMessageId) { - get().updateMessageByClientId(assistantClientId, (message) => ({ - ...message, - serverMessageId, - })); - } - - if (payload?.error) { - throw new Error(payload.error); - } - }; - - const executeStream = async () => { - streamStarted = true; - trackTelemetry("stream_start", { - endpoint: "/api/messages", - mode: requestedMode, - regenerate: Boolean(options?.isRegeneration), - }); + // ── Execute stream ─────────────────────────────────────────────────────── - const fallbackToQueryStream = async (reason: "local" | "fallback") => { - usedQueryFallback = true; - const fallbackLevel = toQueryLevel(effectivePromptMode); - trackTelemetry("stream_start", { - endpoint: "/api/query/stream", - mode: requestedMode, - regenerate: Boolean(options?.isRegeneration), - fallback: true, - reason, - }); - const fallbackResponse = await fetch(`${API_URL}/api/query/stream`, { - method: "POST", - headers, - signal: controller.signal, - body: JSON.stringify({ - topic: trimmed, - levels: [fallbackLevel], - mode: requestedMode, - premium: isPro, - regenerate: Boolean(options?.isRegeneration), - bypass_cache: Boolean(options?.isRegeneration), - temperature: requestTemperature, - message_id: clientMessageId, - }), - }); - - if (!fallbackResponse.ok) { - throw await buildHttpError(fallbackResponse); - } + try { + trackTelemetry("stream_start", { mode: requestedMode }); + let streamError: Error | null = null; - await streamFromResponse( - fallbackResponse, - (payload) => handleStreamingPayload(payload, "chunk"), - controller.signal, + await sendChat({ + conversationId, + content: trimmed, + mode: requestedMode, + promptMode: effectivePromptMode, + temperature: requestTemperature, + isPro, + isRegeneration: Boolean(options?.isRegeneration), + clientMessageId, + assistantClientId, + signal: controller.signal, + onChunk: (chunk) => { + useMessageStore.getState().updateMessageByClientId( + assistantClientId, + (msg) => ({ ...msg, content: msg.content + chunk }), ); - return; - }; - - const hasPersistedConversation = - typeof conversationId === "string" && !conversationId.startsWith("local-"); - const shouldUseMessagesEndpoint = - Boolean(session?.access_token) && - supabaseConfigured && - hasPersistedConversation; - - if (!shouldUseMessagesEndpoint) { - await fallbackToQueryStream("local"); - return; - } - - const response = await fetch(`${API_URL}/api/messages`, { - method: "POST", - headers, - signal: controller.signal, - body: JSON.stringify({ - conversation_id: conversationId, - content: trimmed, - client_generated_id: clientMessageId, - assistant_client_id: assistantClientId, - mode: requestedMode, - prompt_mode: effectivePromptMode, - regenerate: Boolean(options?.isRegeneration), - temperature: requestTemperature, - }), - }); - - const shouldFallback = - response.status === 404 || - response.status === 405 || - response.status >= 500; - if (shouldFallback) { - await fallbackToQueryStream("fallback"); - return; - } - - if (!response.ok) { - throw await buildHttpError(response); - } - - await streamFromResponse( - response, - (payload) => handleStreamingPayload(payload, "delta"), - controller.signal, - ); - }; - - const persistFallbackConversationState = async () => { - if (!supabaseConfigured || !conversationId || conversationId.startsWith("local-")) return; - - const { data: authData } = await supabase.auth.getUser(); - if (!authData?.user) return; - - const assistantMessage = get().messageIds - .map((id) => get().messagesById[id]) - .find( - (message) => - message?.clientGeneratedId === assistantClientId || - message?.metadata?.assistant_client_id === assistantClientId, + }, + onServerMessageId: (id) => { + useMessageStore.getState().updateMessageByClientId( + assistantClientId, + (msg) => ({ ...msg, serverMessageId: id }), ); - - const assistantContent = assistantMessage?.content?.trim() ?? ""; - if (!assistantContent) return; - - const modeMetadata = { - mode: requestedMode, - prompt_mode: effectivePromptMode, - temperature: requestTemperature, - }; - - const messageRows = [ - ...(!skipUserMessage - ? [ - { - conversation_id: conversationId, - role: "user", - content: trimmed, - metadata: { - client_id: clientMessageId, - ...modeMetadata, - }, - }, - ] - : []), - { - conversation_id: conversationId, - role: "assistant", - content: assistantContent, - metadata: { - assistant_client_id: assistantClientId, - ...modeMetadata, - }, - }, - ]; - - const { error: insertError } = await supabase - .from("messages") - .insert(messageRows); - if (insertError) throw insertError; - - const existingConversation = get().conversations.find( - (item) => item.id === conversationId, - ); - const conversationSettings = { - ...(existingConversation?.settings || {}), - mode: requestedMode, - prompt_mode: effectivePromptMode, - }; - const { error: updateError } = await supabase - .from("conversations") - .update({ - mode: persistedConversationMode, - settings: conversationSettings, - updated_at: new Date().toISOString(), - }) - .eq("id", conversationId); - if (updateError) throw updateError; - }; - - const promoteLocalConversationState = async () => { - if (!supabaseConfigured || !conversationId || !conversationId.startsWith("local-")) return; - - const { data: authData } = await supabase.auth.getUser(); - if (!authData?.user) return; - - const state = get(); - const localConversation = state.conversations.find( - (item) => item.id === conversationId, - ); - if (!localConversation) return; - - const { data: remoteConversation, error: createConversationError } = await supabase - .from("conversations") - .insert({ - user_id: authData.user.id, - title: localConversation.title || truncateTitle(trimmed), - mode: persistedConversationMode, - settings: { - ...(localConversation.settings || {}), - mode: requestedMode, - prompt_mode: effectivePromptMode, - }, - }) - .select("id, title, mode, settings, created_at, updated_at") - .single(); - - if (createConversationError || !remoteConversation) { - throw createConversationError || new Error("Failed to create remote conversation"); - } - - const messageRows = state.messageIds - .map((id) => state.messagesById[id]) - .filter((message): message is Message => Boolean(message)) - .filter((message) => message.content.trim().length > 0) - .map((message) => ({ - conversation_id: remoteConversation.id, - role: message.role, - content: message.content, - attachments: message.attachments || [], - metadata: { - ...(message.metadata || {}), - ...(message.role === "assistant" - ? { - assistant_client_id: - message.metadata?.assistant_client_id || message.clientGeneratedId, - } - : { - client_id: message.metadata?.client_id || message.clientGeneratedId, - }), - }, + }, + onError: (error) => { + streamError = error; + }, + onDone: () => { + useMessageStore.getState().updateMessageByClientId(assistantClientId, (msg) => ({ + ...msg, + isStreaming: false, + isRegenerating: false, + syncStatus: "synced", })); + }, + }); - if (messageRows.length > 0) { - const { error: insertMessagesError } = await supabase - .from("messages") - .insert(messageRows); - if (insertMessagesError) throw insertMessagesError; - } - - set((currentState) => ({ - conversations: currentState.conversations - .map((item) => - item.id === conversationId - ? ({ - ...(remoteConversation as Conversation), - settings: - (remoteConversation as Conversation).settings || - item.settings || - null, - } as Conversation) - : item, - ) - .sort((a, b) => (a.updated_at < b.updated_at ? 1 : -1)), - currentConversationId: - currentState.currentConversationId === conversationId - ? remoteConversation.id - : currentState.currentConversationId, - })); - - conversationId = remoteConversation.id; - }; - - await executeStream(); - if (usedQueryFallback) { - await persistFallbackConversationState(); - } - await promoteLocalConversationState(); - if (streamStarted) { - trackTelemetry("stream_end", { - status: "success", - mode: requestedMode, - duration_ms: Math.max(Date.now() - streamStartedAt, 0), - }); + if (streamError) { + throw streamError; } - get().updateMessageByClientId(assistantClientId, (message) => ({ - ...message, - isStreaming: false, - isRegenerating: false, - syncStatus: "synced", - })); + trackTelemetry("stream_end", { + status: "success", + mode: requestedMode, + duration_ms: Math.max(Date.now() - streamStartedAt, 0), + }); } catch (error) { if (isAbortError(error) || controller.signal.aborted) { - if (streamStarted) { - trackTelemetry("stream_end", { - status: "aborted", - mode: requestedMode, - duration_ms: Math.max(Date.now() - streamStartedAt, 0), - }); - } - get().updateMessageByClientId(assistantClientId, (message) => ({ - ...message, + useMessageStore.getState().updateMessageByClientId(assistantClientId, (msg) => ({ + ...msg, isStreaming: false, isRegenerating: false, error: "Canceled", @@ -1579,44 +639,17 @@ export const useChatStore = create((set, get) => ({ } let errorMessage = getErrorMessage(error, "Failed to send message"); - const errorStatus = - typeof error === "object" && error !== null && "status" in error - ? (error as { status?: number }).status - : undefined; - if (errorStatus === 409) { - errorMessage = - "Previous request is still in progress. Retry will send a new request."; - } - if (/timed out/i.test(errorMessage)) { - errorMessage = "Streaming timed out. Retry."; - } - if (streamStarted) { - trackTelemetry("stream_end", { - status: "error", - mode: requestedMode, - duration_ms: Math.max(Date.now() - streamStartedAt, 0), - error_type: (error as { name?: string })?.name || "Error", - }); + if (/timed out/i.test(errorMessage)) errorMessage = "Streaming timed out. Retry."; + if (/duplicate request already in progress/i.test(errorMessage)) { + errorMessage = "Retry will send a new request."; } + captureFrontendError( - error instanceof Error - ? error - : new Error(String(error || "Unknown chat send error")), - { - source: "chat.send_message", - mode: requestedMode, - regenerate: Boolean(options?.isRegeneration), - }, + error instanceof Error ? error : new Error(String(error)), + { source: "chat.send_message", mode: requestedMode }, ); notifyError(errorMessage); - const retryPayload = { - content: trimmed, - mode: requestedMode, - promptMode: effectivePromptMode, - temperature: requestTemperature, - clientMessageId, - assistantClientId, - }; + cachePendingSync({ id: assistantClientId, content: trimmed, @@ -1627,77 +660,70 @@ export const useChatStore = create((set, get) => ({ assistantClientId, }); - get().updateMessageByClientId(assistantClientId, (message) => ({ - ...message, + useMessageStore.getState().updateMessageByClientId(assistantClientId, (msg) => ({ + ...msg, isStreaming: false, isRegenerating: false, - error: errorMessage || getErrorMessage(error, "Failed to sync message"), + error: errorMessage, syncStatus: "failed", - retryPayload, + retryPayload: { + content: trimmed, + mode: requestedMode, + promptMode: effectivePromptMode, + temperature: requestTemperature, + clientMessageId, + assistantClientId, + }, })); } finally { controller.abort(); set((state) => { - const { [assistantClientId]: removedController, ...rest } = - state.streamControllers; - void removedController; - return { streamControllers: rest }; + return { + streamControllers: Object.fromEntries( + Object.entries(state.streamControllers).filter( + ([key]) => key !== assistantClientId, + ), + ), + }; }); - const stillStreaming = get().messageIds.some( - (id) => get().messagesById[id]?.isStreaming, - ); - set({ isLoading: stillStreaming }); + const stillStreaming = useMessageStore + .getState() + .messageIds.some((id) => useMessageStore.getState().messagesById[id]?.isStreaming); + useConversationStore.getState().setIsLoading(stillStreaming); } }, regenerateMessage: async (messageId: string, mode?: ChatMode) => { - if (get().regeneratingMessageId) { - return; - } + if (get().regeneratingMessageId) return; - const { messageIds, messagesById, currentMode, currentPromptMode } = get(); + const { messageIds, messagesById } = useMessageStore.getState(); + const { currentMode, currentPromptMode } = useConversationStore.getState(); const targetIndex = messageIds.indexOf(messageId); - if (targetIndex < 0) { - notifyError("Unable to find the selected message."); - return; - } + if (targetIndex < 0) { notifyError("Unable to find the selected message."); return; } let userMessage: Message | undefined; - for (let i = targetIndex - 1; i >= 0; i -= 1) { + for (let i = targetIndex - 1; i >= 0; i--) { const candidate = messagesById[messageIds[i]]; - if (candidate?.role === "user") { - userMessage = candidate; - break; - } - } - - if (!userMessage) { - notifyError("No user prompt found to regenerate."); - return; + if (candidate?.role === "user") { userMessage = candidate; break; } } + if (!userMessage) { notifyError("No user prompt found to regenerate."); return; } get().abortAllStreams(); - const targetAssistant = messagesById[messageId]; - const nextMode = - (targetAssistant?.metadata?.mode as ChatMode | undefined) ?? - mode ?? - currentMode; + const target = messagesById[messageId]; + const nextMode = (target?.metadata?.mode as ChatMode | undefined) ?? mode ?? currentMode; const nextPromptMode = - (targetAssistant?.metadata?.prompt_mode as PromptMode | undefined) ?? + (target?.metadata?.prompt_mode as PromptMode | undefined) ?? (isPromptMode(nextMode) ? nextMode : currentPromptMode); - const originalTemperature = - typeof targetAssistant?.metadata?.temperature === "number" - ? targetAssistant.metadata.temperature - : 0.7; - const nextTemperature = Math.min(originalTemperature + 0.1, 1.0); + const originalTemp = + typeof target?.metadata?.temperature === "number" ? target.metadata.temperature : 0.7; + const nextTemperature = Math.min(originalTemp + 0.1, 1.0); const originalClientId = typeof userMessage.metadata?.client_id === "string" ? userMessage.metadata.client_id : makeClientId(); set({ regeneratingMessageId: messageId }); - try { await get().sendMessage(userMessage.content, { mode: nextMode, @@ -1715,46 +741,140 @@ export const useChatStore = create((set, get) => ({ }, retrySync: async (messageId: string) => { - const state = get(); - const messageKey = state.messageIds.find((id) => { - const msg = state.messagesById[id]; + const { messageIds, messagesById } = useMessageStore.getState(); + const messageKey = messageIds.find((id) => { + const msg = messagesById[id]; return msg?.clientGeneratedId === messageId || id === messageId; }); - const message = messageKey ? state.messagesById[messageKey] : undefined; + const message = messageKey ? messagesById[messageKey] : undefined; if (!message?.retryPayload) return; - const retryPayloadBase = message.retryPayload; - - // Suspend any active streams before sending a fresh retry request. - get().abortAllStreams(); - - const nextClientMessageId = makeClientId(); - const nextAssistantClientId = makeClientId(); removePendingSync(messageId); - get().updateMessageByClientId(message.clientGeneratedId || messageId, (current) => ({ - ...current, - clientGeneratedId: nextAssistantClientId, - syncStatus: "pending", - error: undefined, - retryPayload: { - ...retryPayloadBase, - clientMessageId: nextClientMessageId, - assistantClientId: nextAssistantClientId, - }, - metadata: { - ...current.metadata, - assistant_client_id: nextAssistantClientId, - }, - })); + useMessageStore.getState().updateMessageByClientId( + message.clientGeneratedId || messageId, + (current) => ({ ...current, syncStatus: "pending", error: undefined }), + ); - await get().sendMessage(retryPayloadBase.content, { - mode: retryPayloadBase.mode as ChatMode, - promptMode: retryPayloadBase.promptMode, - temperature: retryPayloadBase.temperature, - clientMessageId: nextClientMessageId, - assistantClientId: nextAssistantClientId, + await get().sendMessage(message.retryPayload.content, { + mode: message.retryPayload.mode as ChatMode, + promptMode: message.retryPayload.promptMode, + temperature: message.retryPayload.temperature, + clientMessageId: makeClientId(), + assistantClientId: makeClientId(), skipUserMessage: true, replaceMessageId: messageKey, }); }, })); + +const syncLegacyInjectedSlices = (candidate: Partial) => { + const conversationPatch: Partial> = {}; + if ("conversations" in candidate && candidate.conversations !== undefined) { + conversationPatch.conversations = candidate.conversations; + } + if ("currentConversationId" in candidate && candidate.currentConversationId !== undefined) { + conversationPatch.currentConversationId = candidate.currentConversationId; + } + if ("isDraftThread" in candidate && candidate.isDraftThread !== undefined) { + conversationPatch.isDraftThread = candidate.isDraftThread; + } + if ("isLoading" in candidate && candidate.isLoading !== undefined) { + conversationPatch.isLoading = candidate.isLoading; + } + if ("workspace" in candidate && candidate.workspace !== undefined) { + conversationPatch.workspace = candidate.workspace; + } + if ("depthLevel" in candidate && candidate.depthLevel !== undefined) { + conversationPatch.depthLevel = candidate.depthLevel; + } + if ("currentMode" in candidate && candidate.currentMode !== undefined) { + conversationPatch.currentMode = candidate.currentMode; + } + if ("currentPromptMode" in candidate && candidate.currentPromptMode !== undefined) { + conversationPatch.currentPromptMode = candidate.currentPromptMode; + } + if ("selectedLevel" in candidate && candidate.selectedLevel !== undefined) { + conversationPatch.selectedLevel = candidate.selectedLevel; + } + if (Object.keys(conversationPatch).length > 0) { + useConversationStore.setState(conversationPatch); + } + + const messagePatch: Partial> = {}; + if ("messagesById" in candidate && candidate.messagesById !== undefined) { + messagePatch.messagesById = candidate.messagesById; + } + if ("messageIds" in candidate && candidate.messageIds !== undefined) { + messagePatch.messageIds = candidate.messageIds; + } + if (Object.keys(messagePatch).length > 0) { + useMessageStore.setState(messagePatch); + } +}; + +const restoreProxyGetters = () => { + const state = useChatStore.getState() as unknown as Record; + Object.defineProperties(state, { + conversations: { + configurable: true, + enumerable: true, + get: () => useConversationStore.getState().conversations, + }, + currentConversationId: { + configurable: true, + enumerable: true, + get: () => useConversationStore.getState().currentConversationId, + }, + isDraftThread: { + configurable: true, + enumerable: true, + get: () => useConversationStore.getState().isDraftThread, + }, + isLoading: { + configurable: true, + enumerable: true, + get: () => useConversationStore.getState().isLoading, + }, + workspace: { + configurable: true, + enumerable: true, + get: () => useConversationStore.getState().workspace, + }, + depthLevel: { + configurable: true, + enumerable: true, + get: () => useConversationStore.getState().depthLevel, + }, + currentMode: { + configurable: true, + enumerable: true, + get: () => useConversationStore.getState().currentMode, + }, + currentPromptMode: { + configurable: true, + enumerable: true, + get: () => useConversationStore.getState().currentPromptMode, + }, + selectedLevel: { + configurable: true, + enumerable: true, + get: () => useConversationStore.getState().selectedLevel, + }, + messagesById: { + configurable: true, + enumerable: true, + get: () => useMessageStore.getState().messagesById, + }, + messageIds: { + configurable: true, + enumerable: true, + get: () => useMessageStore.getState().messageIds, + }, + }); +}; + +restoreProxyGetters(); +useChatStore.subscribe((state) => { + syncLegacyInjectedSlices(state); + restoreProxyGetters(); +}); diff --git a/src/stores/useConversationStore.ts b/src/stores/useConversationStore.ts new file mode 100644 index 0000000..1f37b3c --- /dev/null +++ b/src/stores/useConversationStore.ts @@ -0,0 +1,295 @@ +import { create } from "zustand"; +import { supabase } from "../lib/supabase"; +import type { ChatMode, Conversation, Message, PromptMode } from "../types/chat"; +import type { Level } from "../types"; +import { + asString, + notifyError, + resolveWorkspaceState, + type Workspace, + type DepthLevel, + DEFAULT_DEPTH_LEVEL, +} from "../lib/chatStoreUtils"; +import { useMessageStore } from "./useMessageStore"; + +const supabaseConfigured = + Boolean(import.meta.env.VITE_SUPABASE_URL) && + Boolean(import.meta.env.VITE_SUPABASE_ANON_KEY); + +interface ConversationState { + conversations: Conversation[]; + currentConversationId: string | null; + isDraftThread: boolean; + isLoading: boolean; + + // Derived workspace state (kept here because it derives from active conversation) + workspace: Workspace; + depthLevel: DepthLevel; + currentMode: ChatMode; + currentPromptMode: PromptMode; + selectedLevel: Level; + + // Actions + syncConversations: (conversations: Conversation[]) => void; + selectConversation: (id: string) => Promise; + renameConversation: (id: string, title: string) => Promise; + deleteConversation: (id: string) => Promise; + setWorkspaceState: ( + workspace: Workspace, + mode: ChatMode, + promptMode: PromptMode, + depthLevel: DepthLevel, + ) => void; + setCurrentConversationId: (id: string | null) => void; + setIsDraftThread: (draft: boolean) => void; + setIsLoading: (loading: boolean) => void; + upsertConversation: (conversation: Conversation) => void; +} + +export const useConversationStore = create((set, get) => ({ + conversations: [], + currentConversationId: null, + isDraftThread: false, + isLoading: false, + workspace: "learn", + depthLevel: DEFAULT_DEPTH_LEVEL, + currentMode: "learning", + currentPromptMode: DEFAULT_DEPTH_LEVEL as PromptMode, + selectedLevel: DEFAULT_DEPTH_LEVEL as Level, + + setWorkspaceState: (workspace, mode, promptMode, depthLevel) => + set({ workspace, currentMode: mode, currentPromptMode: promptMode, depthLevel, selectedLevel: depthLevel as Level }), + + setCurrentConversationId: (id) => set({ currentConversationId: id }), + setIsDraftThread: (draft) => set({ isDraftThread: draft }), + setIsLoading: (loading) => set({ isLoading: loading }), + + upsertConversation: (conversation: Conversation) => { + set((state) => { + const exists = state.conversations.some((c) => c.id === conversation.id); + const next = exists + ? state.conversations.map((c) => + c.id === conversation.id ? conversation : c, + ) + : [conversation, ...state.conversations]; + return { + conversations: next.sort((a, b) => + a.updated_at < b.updated_at ? 1 : -1, + ), + }; + }); + }, + + syncConversations: (conversations: Conversation[]) => { + set((state) => { + if (state.isDraftThread && state.currentConversationId === null) { + return { conversations }; + } + + const preferredId = state.currentConversationId; + const hasPreferred = preferredId + ? conversations.some((item) => item.id === preferredId) + : false; + const nextConversationId = hasPreferred + ? preferredId + : (conversations[0]?.id ?? null); + const activeConversation = conversations.find( + (item) => item.id === nextConversationId, + ); + const conversationMode = + asString(activeConversation?.mode) || + asString(activeConversation?.settings?.mode); + const conversationPrompt = + asString(activeConversation?.settings?.prompt_mode) || + asString(activeConversation?.settings?.mode) || + asString(activeConversation?.mode) || + state.currentPromptMode; + const nextWorkspaceState = resolveWorkspaceState( + conversationMode, + conversationPrompt, + state.depthLevel, + ); + return { + conversations, + currentConversationId: nextConversationId, + isDraftThread: false, + workspace: nextWorkspaceState.workspace, + depthLevel: nextWorkspaceState.depthLevel, + currentMode: nextWorkspaceState.mode, + currentPromptMode: nextWorkspaceState.promptMode, + selectedLevel: nextWorkspaceState.depthLevel as Level, + }; + }); + }, + + selectConversation: async (id: string) => { + if (!id) return; + const state = get(); + + if ( + state.currentConversationId === id && + (state.isLoading || useMessageStore.getState().messageIds.length > 0) + ) { + return; + } + + const activeConversation = state.conversations.find( + (item) => item.id === id, + ); + const conversationMode = + asString(activeConversation?.mode) || + asString(activeConversation?.settings?.mode) || + state.currentMode; + const conversationPrompt = + asString(activeConversation?.settings?.prompt_mode) || + asString(activeConversation?.settings?.mode) || + asString(activeConversation?.mode) || + state.currentPromptMode; + const nextWorkspaceState = resolveWorkspaceState( + conversationMode, + conversationPrompt, + state.depthLevel, + ); + + useMessageStore.getState().clearMessages(); + set({ + currentConversationId: id, + isDraftThread: false, + isLoading: true, + workspace: nextWorkspaceState.workspace, + depthLevel: nextWorkspaceState.depthLevel, + currentMode: nextWorkspaceState.mode, + currentPromptMode: nextWorkspaceState.promptMode, + selectedLevel: nextWorkspaceState.depthLevel as Level, + }); + + if (!supabaseConfigured) { + set({ isLoading: false }); + return; + } + + try { + const { data, error } = await supabase + .from("messages") + .select("id, role, content, attachments, metadata, created_at") + .eq("conversation_id", id) + .order("created_at", { ascending: true }); + + if (error) throw error; + useMessageStore.getState().setMessages((data ?? []) as Message[]); + } catch (error) { + console.error("Failed to fetch messages:", error); + } finally { + set({ isLoading: false }); + } + }, + + renameConversation: async (id: string, title: string) => { + if (!id) return; + const trimmed = title.trim(); + if (!trimmed) return; + + const existingConversation = get().conversations.find((item) => item.id === id); + if (!existingConversation) return; + const previousTitle = existingConversation.title; + const previousUpdatedAt = existingConversation.updated_at; + + const now = new Date().toISOString(); + set((state) => ({ + conversations: state.conversations + .map((item) => + item.id === id ? { ...item, title: trimmed, updated_at: now } : item, + ) + .sort((a, b) => (a.updated_at < b.updated_at ? 1 : -1)), + })); + + if (!supabaseConfigured || id.startsWith("local-")) return; + + const rollbackRename = () => { + set((state) => ({ + conversations: state.conversations + .map((item) => + item.id === id + ? { ...item, title: previousTitle, updated_at: previousUpdatedAt } + : item, + ) + .sort((a, b) => (a.updated_at < b.updated_at ? 1 : -1)), + })); + }; + + try { + const { data, error } = await supabase + .from("conversations") + .update({ title: trimmed, updated_at: now }) + .eq("id", id); + + void data; + if (error) { + console.error("Failed to rename conversation:", { + id, + title: trimmed, + error, + }); + rollbackRename(); + notifyError("Failed to rename conversation."); + } + } catch (error) { + console.error("Failed to rename conversation:", error); + rollbackRename(); + notifyError("Failed to rename conversation."); + } + }, + + deleteConversation: async (id: string) => { + if (!id) return; + + const state = get(); + const targetConversation = state.conversations.find((c) => c.id === id); + if (!targetConversation) return; + const isActive = state.currentConversationId === id; + + if (!id.startsWith("local-") && supabaseConfigured) { + try { + const { error } = await supabase + .from("conversations") + .delete() + .eq("id", id); + if (error) throw error; + } catch (error) { + console.error("Failed to delete conversation:", error); + notifyError("Failed to delete conversation."); + return; + } + } + + const remaining = get() + .conversations.filter((c) => c.id !== id) + .sort((a, b) => (a.updated_at < b.updated_at ? 1 : -1)); + + if (!isActive) { + set({ conversations: remaining }); + return; + } + + if (remaining.length === 0) { + set({ conversations: remaining }); + useMessageStore.getState().clearMessages(); + set({ + currentConversationId: null, + isDraftThread: true, + isLoading: false, + }); + return; + } + + const nextId = remaining[0].id; + set({ + conversations: remaining, + currentConversationId: nextId, + isDraftThread: false, + isLoading: false, + }); + useMessageStore.getState().clearMessages(); + await get().selectConversation(nextId); + }, +})); diff --git a/src/stores/useMessageStore.ts b/src/stores/useMessageStore.ts new file mode 100644 index 0000000..495eb60 --- /dev/null +++ b/src/stores/useMessageStore.ts @@ -0,0 +1,114 @@ +import { create } from "zustand"; +import type { Message } from "../types/chat"; +import { + resolveMessageKey, + findExistingMessageKey, + buildMessageRegistry, +} from "../lib/chatStoreUtils"; + +interface MessageState { + messagesById: Record; + messageIds: string[]; + + // Actions + setMessages: (messages: Message[]) => void; + addMessage: (msg: Message) => void; + updateMessageByClientId: ( + clientId: string, + updater: (msg: Message) => Message, + ) => void; + removeMessageByClientId: (clientId: string) => void; + clearMessages: () => void; +} + +export const useMessageStore = create((set) => ({ + messagesById: {}, + messageIds: [], + + setMessages: (messages: Message[]) => { + const { messagesById, messageIds } = buildMessageRegistry(messages); + set({ messagesById, messageIds }); + }, + + addMessage: (msg: Message) => { + set((state) => { + const resolvedKey = resolveMessageKey(msg); + const directMatch = state.messagesById[resolvedKey] ? resolvedKey : null; + const existingKey = directMatch || findExistingMessageKey(state, msg); + + if (existingKey) { + const nextMessagesById = { + ...state.messagesById, + [existingKey]: { + ...state.messagesById[existingKey], + ...msg, + metadata: { + ...state.messagesById[existingKey].metadata, + ...msg.metadata, + }, + }, + }; + if (!state.messageIds.includes(existingKey)) { + return { + messagesById: nextMessagesById, + messageIds: [...state.messageIds, existingKey], + }; + } + return { messagesById: nextMessagesById }; + } + + return { + messagesById: { ...state.messagesById, [resolvedKey]: msg }, + messageIds: [...state.messageIds, resolvedKey], + }; + }); + }, + + updateMessageByClientId: ( + clientId: string, + updater: (msg: Message) => Message, + ) => { + set((state) => { + const messageKey = state.messageIds.find((id) => { + const message = state.messagesById[id]; + if (!message) return false; + return ( + message.clientGeneratedId === clientId || + message.metadata?.assistant_client_id === clientId || + message.metadata?.client_id === clientId + ); + }); + if (!messageKey) return state; + const message = state.messagesById[messageKey]; + return { + messagesById: { + ...state.messagesById, + [messageKey]: updater(message), + }, + }; + }); + }, + + removeMessageByClientId: (clientId: string) => { + set((state) => { + const messageKey = state.messageIds.find((id) => { + const message = state.messagesById[id]; + if (!message) return false; + return ( + message.clientGeneratedId === clientId || + message.metadata?.assistant_client_id === clientId || + message.metadata?.client_id === clientId + ); + }); + if (!messageKey) return state; + const { [messageKey]: removed, ...rest } = state.messagesById; + void removed; + return { + messagesById: rest, + messageIds: state.messageIds.filter((id) => id !== messageKey), + }; + }); + }, + + clearMessages: () => set({ messagesById: {}, messageIds: [] }), +}));