From 30472b8f0c4ceb3fd2938e22fc32d0b05c0b097d Mon Sep 17 00:00:00 2001 From: Fredrik Ekholdt Date: Fri, 5 Jun 2026 16:40:47 +0000 Subject: [PATCH] Improve AI session handling --- .../shared/src/internal/sessionStorage.ts | 1 + packages/ui/spa/components/AIChat.tsx | 12 +- packages/ui/spa/components/ToolsMenu.tsx | 30 +- packages/ui/spa/components/ValOverlay.tsx | 271 +++++++++++------- packages/ui/spa/components/ValRouter.tsx | 56 +++- packages/ui/spa/hooks/useAI.ts | 120 ++++++-- 6 files changed, 353 insertions(+), 137 deletions(-) diff --git a/packages/shared/src/internal/sessionStorage.ts b/packages/shared/src/internal/sessionStorage.ts index f7aa2f798..7861f9d79 100644 --- a/packages/shared/src/internal/sessionStorage.ts +++ b/packages/shared/src/internal/sessionStorage.ts @@ -1,2 +1,3 @@ export const VAL_CONFIG_SESSION_STORAGE_KEY = "val-config"; export const VAL_THEME_SESSION_STORAGE_KEY = "val-theme"; +export const VAL_AI_SESSION_STORAGE_KEY = "val-ai-session"; diff --git a/packages/ui/spa/components/AIChat.tsx b/packages/ui/spa/components/AIChat.tsx index b5d302c2c..bf0238ec0 100644 --- a/packages/ui/spa/components/AIChat.tsx +++ b/packages/ui/spa/components/AIChat.tsx @@ -134,14 +134,16 @@ export type AIChatProps = { mode: "http" | "fs" | "unknown"; /** List of past sessions (fetched on demand) */ sessions?: AISession[]; - /** The currently active session ID */ - currentSessionId?: string; + /** The currently active session ID; null when the session is unborn (no message sent yet). */ + currentSessionId?: string | null; /** Called to load a previous session */ onLoadSession?: (sessionId: string) => void; /** Called to trigger a sessions fetch */ onFetchSessions?: () => void; /** Called to rename a session */ onSetSessionName?: (sessionId: string, name: string) => void; + /** True while a previous session's messages are being fetched from the server. */ + isLoadingSession?: boolean; /** * @internal – seed messages for Storybook / testing only. * Not part of the public API. @@ -212,6 +214,7 @@ export const AIChat = forwardRef(function AIChat( onLoadSession, onFetchSessions, onSetSessionName, + isLoadingSession, initialMessages, }, ref, @@ -723,6 +726,11 @@ export const AIChat = forwardRef(function AIChat(
{authError ? ( + ) : isLoadingSession && isEmpty ? ( +
+ + Loading conversation… +
) : isEmpty ? ( (null); + const { sessionParam, setSessionParam } = useSessionParam(); + // Capture the URL session id once on first render — later URL changes (e.g. + // navigations that re-write the URL) must not retrigger session loading. + const initialSessionIdRef = useRef(sessionParam); const { sendMessage, uploadAiImage, @@ -78,7 +82,26 @@ export function ToolsMenu() { getSessions, setSessionName, loadSession, - } = useAI(chatRef); + isLoadingSession, + } = useAI(chatRef, { + initialSessionId: initialSessionIdRef.current, + onSessionBorn: (id) => { + setSessionParam(id, { replace: true }); + try { + sessionStorage.setItem(VAL_AI_SESSION_STORAGE_KEY, id); + } catch { + // sessionStorage may be disabled — URL remains the source of truth. + } + }, + onSessionCleared: () => { + setSessionParam(null, { replace: true }); + try { + sessionStorage.removeItem(VAL_AI_SESSION_STORAGE_KEY); + } catch { + // see above + } + }, + }); return (
)} diff --git a/packages/ui/spa/components/ValOverlay.tsx b/packages/ui/spa/components/ValOverlay.tsx index 1bdfd56d5..7e3b4497f 100644 --- a/packages/ui/spa/components/ValOverlay.tsx +++ b/packages/ui/spa/components/ValOverlay.tsx @@ -20,6 +20,9 @@ import { import React, { Dispatch, SetStateAction, + createContext, + useCallback, + useContext, useEffect, useMemo, useRef, @@ -45,7 +48,7 @@ import { useSchemaAtPath, useValConfig, useSchemas } from "./ValFieldProvider"; import { useTheme } from "./ValThemeProvider"; import { useValPortal } from "./ValPortalProvider"; import { FieldLoading } from "./FieldLoading"; -import { urlOf } from "@valbuild/shared/internal"; +import { urlOf, VAL_AI_SESSION_STORAGE_KEY } from "@valbuild/shared/internal"; import { Popover, PopoverContent } from "./designSystem/popover"; import { PopoverTrigger } from "@radix-ui/react-popover"; import { Switch } from "./designSystem/switch"; @@ -76,6 +79,23 @@ export type ValOverlayProps = { disableOverlay: () => void; }; +/** + * Holds the overlay's "born" AI chat session id so that the studio-link + * sites (top menu, WindowField) can append `?session=` to their hrefs. + * Null until the user has actually used the chat (sent a message or + * uploaded an attachment) — matches the unborn/born model in useAI. + */ +const OverlaySessionContext = createContext<{ + bornSessionId: string | null; + setBornSessionId: (id: string | null) => void; +}>({ bornSessionId: null, setBornSessionId: () => {} }); + +function appendSessionParam(base: string, sid: string | null): string { + if (!sid) return base; + const sep = base.includes("?") ? "&" : "?"; + return `${base}${sep}session=${encodeURIComponent(sid)}`; +} + type ValMenuProps = ValOverlayProps & { setMode: Dispatch>; mode: OverlayModes; @@ -151,6 +171,26 @@ export function ValOverlay(props: ValOverlayProps) { }> >([]); const [isChatOpen, setIsChatOpen] = useState(false); + // The overlay's "born" chat session id is mirrored to sessionStorage so it + // survives a host page reload and so other consumers (e.g. studio links + // generated from inside the overlay) keep working after reload. + const [bornSessionId, setBornSessionIdState] = useState(() => { + try { + return sessionStorage.getItem(VAL_AI_SESSION_STORAGE_KEY); + } catch { + return null; + } + }); + const setBornSessionId = useCallback((id: string | null) => { + setBornSessionIdState(id); + try { + if (id == null) sessionStorage.removeItem(VAL_AI_SESSION_STORAGE_KEY); + else sessionStorage.setItem(VAL_AI_SESSION_STORAGE_KEY, id); + } catch { + // sessionStorage may be disabled (private mode, quota) — fall back to + // React-state-only behavior. + } + }, []); useEffect(() => { if (!props.draftMode) { setMode(null); @@ -341,67 +381,31 @@ export function ValOverlay(props: ValOverlayProps) { }; return ( -
- - {boundingBox && ( -
{ - ev.preventDefault(); - ev.stopPropagation(); - setBoundingBox(null); - setEditMode({ - joinedPaths: boundingBox.joinedPaths, - clientX: boundingBox.left, - clientY: boundingBox.top, - boundingBox: { - top: boundingBox.top, - left: boundingBox.left, - width: boundingBox.width, - height: boundingBox.height, - }, - }); - }} - >
- )} - {showAllBoundingBoxes && - allBoundingBoxes.map((box, index) => ( + +
+ + {boundingBox && (
- ))} - setIsChatOpen(false)} - /> - {editMode === null && ( - { + ev.preventDefault(); + ev.stopPropagation(); + setBoundingBox(null); + setEditMode({ + joinedPaths: boundingBox.joinedPaths, + clientX: boundingBox.left, + clientY: boundingBox.top, + boundingBox: { + top: boundingBox.top, + left: boundingBox.left, + width: boundingBox.width, + height: boundingBox.height, + }, + }); + }} + >
+ )} + {showAllBoundingBoxes && + allBoundingBoxes.map((box, index) => ( +
+ ))} + setIsChatOpen(false)} /> - )} -
+ {editMode === null && ( + + )} +
+ ); } @@ -446,6 +495,7 @@ function Window({ setMode: Dispatch>; setEditMode: Dispatch>; }) { + const { bornSessionId } = useContext(OverlaySessionContext); // place outside viewport initially const [windowPos, setWindowPos] = useState({ x: window.innerWidth, @@ -757,11 +807,14 @@ function Window({ Internal.splitJoinedSourcePaths(editMode.joinedPaths).length === 1 && ( @@ -832,6 +885,12 @@ function ChatWindow({ onClose: () => void; }) { const chatRef = useRef(null); + const { bornSessionId, setBornSessionId } = useContext(OverlaySessionContext); + // Capture the seeded-from-sessionStorage value once so later state changes + // don't re-trigger useAI's mount-load effect. This is how a user coming + // back from the studio (with an active session) sees the chat rehydrate + // in the overlay. + const initialBornSessionIdRef = useRef(bornSessionId); const { sendMessage, uploadAiImage, @@ -843,7 +902,12 @@ function ChatWindow({ getSessions, setSessionName, loadSession, - } = useAI(chatRef); + isLoadingSession, + } = useAI(chatRef, { + initialSessionId: initialBornSessionIdRef.current, + onSessionBorn: setBornSessionId, + onSessionCleared: () => setBornSessionId(null), + }); const mode = useValMode(); const [windowPos, setWindowPos] = useState({ x: Math.max(20, window.innerWidth - 570), @@ -1013,6 +1077,7 @@ function ChatWindow({ onLoadSession={loadSession} onFetchSessions={getSessions} onSetSessionName={setSessionName} + isLoadingSession={isLoadingSession} />
{!isMobile && ( @@ -1073,7 +1138,11 @@ function WindowField({ showInlineStudioLink: boolean; }) { const schemaAtPath = useSchemaAtPath(path); - const studioUrl = window.origin + "/val/~" + path; + const { bornSessionId } = useContext(OverlaySessionContext); + const studioUrl = appendSessionParam( + window.origin + "/val/~" + path, + bornSessionId, + ); if (schemaAtPath.status === "error") { return ( @@ -1211,6 +1280,7 @@ function ValMenu({ dropZone: DropZones; ghost?: boolean; } & ValMenuProps) { + const { bornSessionId } = useContext(OverlaySessionContext); const dir = dropZone === "val-menu-right-center" || dropZone === "val-menu-left-center" ? "vertical" @@ -1381,7 +1451,10 @@ function ValMenu({ {patchIds.length > 0 && ( @@ -1455,14 +1528,15 @@ function ValMenu({ ) } - href={ + href={appendSessionParam( window.origin + - "/val/~" + - (sourcePathResult.status === "success" && - sourcePathResult.data - ? sourcePathResult.data - : "") - } + "/val/~" + + (sourcePathResult.status === "success" && + sourcePathResult.data + ? sourcePathResult.data + : ""), + bornSessionId, + )} /> ) } - href={ + href={appendSessionParam( window.origin + - "/val/~" + - (sourcePathResult.status === "success" && sourcePathResult.data - ? sourcePathResult.data - : "") - } + "/val/~" + + (sourcePathResult.status === "success" && sourcePathResult.data + ? sourcePathResult.data + : ""), + bornSessionId, + )} /> void; }; const ValRouterContext = React.createContext( new Proxy( @@ -88,9 +93,11 @@ export function ValRouter({ const [isCompareView, setIsCompareView] = useState(false); const [isErrorsView, setIsErrorsView] = useState(false); const [errorFields, setErrorFields] = useState([]); + const [sessionParam, setSessionParamState] = useState(null); const historyState = useRef([]); useEffect(() => { const listener = () => { + setSessionParamState(new URLSearchParams(location.search).get("session")); if ( location.pathname === VAL_COMPARE_ROUTE || location.pathname === VAL_COMPARE_ROUTE + "/" @@ -208,6 +215,25 @@ export function ValRouter({ : isErrors ? VAL_ERRORS_ROUTE + errorFieldsQuery : `${VAL_CONTENT_VIEW_ROUTE}${path}`; + // Preserve `?session=` across in-studio navigations. Without this every + // sidebar/compare/errors click would strip the AI chat session id from + // the URL. URL only — useAI does not re-read this after mount, so chat + // state is intentionally not affected by navigations. + // + // In overlay mode the host page URL has no `?session=`, so fall back to + // sessionStorage so AI navigate_to (and overlay→studio nav generally) + // brings the active chat along to the studio. + let sid: string | null = sessionParam; + if (overlay && sid == null) { + try { + sid = sessionStorage.getItem(VAL_AI_SESSION_STORAGE_KEY); + } catch { + sid = null; + } + } + const finalTo = sid + ? `${navigateTo}${navigateTo.includes("?") ? "&" : "?"}session=${encodeURIComponent(sid)}` + : navigateTo; setIsCompareView(isCompare); setIsErrorsView(isErrors); setErrorFields(isErrors ? (params?.errorFields ?? []) : []); @@ -240,18 +266,35 @@ export function ValRouter({ historyState.current.push(prevScrollPos); } if (params?.replace) { - window.history.replaceState(null, "", navigateTo); + window.history.replaceState(null, "", finalTo); } else { - window.history.pushState(null, "", navigateTo); + window.history.pushState(null, "", finalTo); } } else { window.location.href = - navigateTo + + finalTo + (params?.scrollToPath ? `#${encodeURIComponent(params.scrollToPath)}` : ""); } }, + [overlay, sessionParam], + ); + const setSessionParam = useCallback( + (id: string | null, opts?: { replace?: boolean }) => { + // Overlay runs on the host page — never mutate that URL. + if (overlay) return; + const url = new URL(window.location.href); + if (id == null) url.searchParams.delete("session"); + else url.searchParams.set("session", id); + const target = url.pathname + url.search + url.hash; + if (opts?.replace) { + window.history.replaceState(null, "", target); + } else { + window.history.pushState(null, "", target); + } + setSessionParamState(id); + }, [overlay], ); return ( @@ -264,6 +307,8 @@ export function ValRouter({ isCompareView, isErrorsView, errorFields, + sessionParam, + setSessionParam, }} > {children} @@ -298,3 +343,8 @@ export function useParams(): { sourcePath: ctx.currentSourcePath, }; } + +export function useSessionParam() { + const { sessionParam, setSessionParam } = useContext(ValRouterContext); + return { sessionParam, setSessionParam }; +} diff --git a/packages/ui/spa/hooks/useAI.ts b/packages/ui/spa/hooks/useAI.ts index 072825844..29a38da22 100644 --- a/packages/ui/spa/hooks/useAI.ts +++ b/packages/ui/spa/hooks/useAI.ts @@ -18,7 +18,6 @@ import type { AIMessageContentBlock, AIPromptMessage, } from "./useAIWebSocket"; -import { getRecentSession } from "./useAIWebSocket"; import { useAISearch } from "./useAISearch"; import { useAIValidation } from "./useAIValidation"; import type { @@ -338,7 +337,19 @@ const ALL_TOOLS: AITool[] = [ SHOW_COMPARE_VIEW_TOOL, ]; -export function useAI(chatRef: React.RefObject) { +export type UseAIOptions = { + /** Initial session id to load on mount (e.g. read from URL). When null, the chat starts in an unborn state. */ + initialSessionId?: string | null; + /** Called when the session transitions from unborn → born (first send) or when an existing session is loaded. */ + onSessionBorn?: (id: string) => void; + /** Called when the session returns to the unborn state (newSession, or loadSession failure). */ + onSessionCleared?: () => void; +}; + +export function useAI( + chatRef: React.RefObject, + opts?: UseAIOptions, +) { const { subscribeToWsMessages, sendWsMessage, @@ -359,11 +370,17 @@ export function useAI(chatRef: React.RefObject) { const config = useValConfig(); const isChatEnabled = config?.ai?.chat?.experimental?.enable === true; const [isStreaming, setIsStreaming] = useState(false); + const [isLoadingSession, setIsLoadingSession] = useState(false); const [sessions, setSessions] = useState([]); - const [currentSessionId, setCurrentSessionId] = useState(() => - crypto.randomUUID(), - ); - const sessionIdRef = useRef(currentSessionId); + const [currentSessionId, setCurrentSessionId] = useState(null); + const sessionIdRef = useRef(null); + // Keep callbacks in a ref so we don't have to thread them through every useCallback dep array. + const onSessionBornRef = useRef(opts?.onSessionBorn); + const onSessionClearedRef = useRef(opts?.onSessionCleared); + useEffect(() => { + onSessionBornRef.current = opts?.onSessionBorn; + onSessionClearedRef.current = opts?.onSessionCleared; + }); // Track active streaming ID — startAssistantMessage always appends a new // message (NOT idempotent), so we must only call it once per message ID. const activeIdRef = useRef(null); @@ -1154,12 +1171,22 @@ export function useAI(chatRef: React.RefObject) { } else if (message.name === "set_session_name") { const args = message.arguments as { name: string }; const name = String(args.name ?? "").slice(0, 60); - aiSetSessionName(sessionIdRef.current, name) + const sid = sessionIdRef.current; + if (sid == null) { + // Defensive: tool calls only run for an active session. + sendWsMessage({ + type: "ai_tool_result", + toolCallId: message.toolCallId, + result: { success: false, error: "No active session." }, + isError: true, + }); + chatRef.current?.errorToolCall(message.id, message.toolCallId); + return; + } + aiSetSessionName(sid, name) .then(() => { setSessions((prev) => - prev.map((s) => - s.id === sessionIdRef.current ? { ...s, name } : s, - ), + prev.map((s) => (s.id === sid ? { ...s, name } : s)), ); sendWsMessage({ type: "ai_tool_result", @@ -1249,6 +1276,15 @@ export function useAI(chatRef: React.RefObject) { "Cannot upload AI image: content host is not configured. Set the `project` option in val.config (or VAL_PROJECT) and ensure a personal access token is available.", ); } + // Uploading an image attaches binary data to the session on the server, + // so it counts as "using" the session — mint the id now if still unborn. + let sid = sessionIdRef.current; + const wasUnborn = sid == null; + if (sid == null) { + sid = crypto.randomUUID(); + sessionIdRef.current = sid; + setCurrentSessionId(sid); + } const headers: Record = { "Content-Type": file.type || "application/octet-stream", "Content-Length": String(file.size), @@ -1257,7 +1293,7 @@ export function useAI(chatRef: React.RefObject) { headers["x-val-auth-nonce"] = contentAuthNonce; } const queryParams = new URLSearchParams(); - queryParams.set("sessionid", encodeURIComponent(sessionIdRef.current)); + queryParams.set("sessionid", encodeURIComponent(sid)); queryParams.set("width", encodeURIComponent(readRes.width || 0)); queryParams.set("height", encodeURIComponent(readRes.height || 0)); queryParams.set( @@ -1278,6 +1314,9 @@ export function useAI(chatRef: React.RefObject) { throw new Error(`Upload failed: ${res.status} ${res.statusText}`); } const body = (await res.json()) as { key: string }; + if (wasUnborn) { + onSessionBornRef.current?.(sid); + } return body; }, [getDirectFileUploadSettings], @@ -1285,6 +1324,15 @@ export function useAI(chatRef: React.RefObject) { const sendMessage = useCallback( (text: string, attachments?: ChatMessageAttachment[]): boolean => { + // Lazily mint the session id on the first send so unborn sessions don't + // appear in the URL or on the server until the user actually says something. + let sid = sessionIdRef.current; + const wasUnborn = sid == null; + if (sid == null) { + sid = crypto.randomUUID(); + sessionIdRef.current = sid; + setCurrentSessionId(sid); + } let augmentedText = text; if (attachments && attachments.length > 0) { const lines = attachments.map( @@ -1307,7 +1355,7 @@ export function useAI(chatRef: React.RefObject) { const message: AIPromptMessage = { type: "ai_prompt", message: contentBlocks, - sessionId: sessionIdRef.current, + sessionId: sid, id: crypto.randomUUID(), agents: [ { @@ -1389,16 +1437,22 @@ Do not describe what you will do unless you do it for clarification — just do }, ], }; - return sendWsMessage(message); + const sent = sendWsMessage(message); + // Notify the session was "born" only after a successful send so a failed + // first send doesn't leak an empty session id into the URL. + if (sent && wasUnborn) { + onSessionBornRef.current?.(sid); + } + return sent; }, [sendWsMessage], ); const newSession = useCallback(() => { - const id = crypto.randomUUID(); - sessionIdRef.current = id; - setCurrentSessionId(id); + sessionIdRef.current = null; + setCurrentSessionId(null); chatRef.current?.clearMessages(); + onSessionClearedRef.current?.(); }, [chatRef]); const getSessions = useCallback( @@ -1428,6 +1482,7 @@ Do not describe what you will do unless you do it for clarification — just do sessionIdRef.current = sessionId; setCurrentSessionId(sessionId); chatRef.current?.clearMessages(); + setIsLoadingSession(true); try { const res = await aiGetSessionMessages(sessionId); const messages = res.messages @@ -1439,35 +1494,38 @@ Do not describe what you will do unless you do it for clarification — just do status: "complete" as const, })); chatRef.current?.loadMessages(messages); + onSessionBornRef.current?.(sessionId); } catch (err) { - console.error("Failed to load session messages:", err); + console.warn("Failed to load session messages, clearing:", err); + sessionIdRef.current = null; + setCurrentSessionId(null); + chatRef.current?.clearMessages(); + onSessionClearedRef.current?.(); + } finally { + setIsLoadingSession(false); } }, [chatRef, aiGetSessionMessages], ); - // On mount, restore the most recent session if it was used within the last 24 hours + // On mount, populate the sessions dropdown and (if an initial session id was + // passed in, e.g. from ?session= in the URL) load it. Intentional: this runs + // only once — popstate / URL changes after mount must NOT hijack the user's + // open chat, so we capture initialSessionId via opts only on the first render. + const initialSessionIdRef = useRef(opts?.initialSessionId ?? null); useEffect(() => { if (!isChatEnabled) return; - let cancelled = false; - getSessions({ limit: 1 }) - .then((fetchedSessions) => { - if (cancelled) return; - const session = getRecentSession(fetchedSessions); - if (session) return loadSession(session.id); - }) - .catch((err) => { - console.error("Failed to restore last session:", err); - }); - return () => { - cancelled = true; - }; + getSessions({ limit: 1 }).catch(() => {}); + if (initialSessionIdRef.current != null) { + loadSession(initialSessionIdRef.current); + } }, [isChatEnabled]); return { sendMessage, uploadAiImage, isStreaming, + isLoadingSession, isConnected: isWsConnected, authError: aiAuthError, newSession,