Improve AI session handling

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";

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<AIChatHandle, AIChatProps>(function AIChat(
     onLoadSession,
     onFetchSessions,
     onSetSessionName,
+    isLoadingSession,
     initialMessages,
   },
   ref,
@@ -723,6 +726,11 @@ export const AIChat = forwardRef<AIChatHandle, AIChatProps>(function AIChat(
         <div className="flex flex-col gap-4 p-4">
           {authError ? (
             <AuthPrompt mode={mode} />
+          ) : isLoadingSession && isEmpty ? (
+            <div className="flex flex-col items-center justify-center gap-2 py-12 text-fg-secondary">
+              <Loader2 className="h-5 w-5 animate-spin" />
+              <span className="text-sm">Loading conversation…</span>
+            </div>
           ) : isEmpty ? (
             <EmptyState
               suggestions={effectiveSuggestions}

packages/ui/spa/components/ToolsMenu.tsx

@@ -24,15 +24,15 @@ import {
   PanelsTopLeft,
 } from "lucide-react";
 import { Button } from "./designSystem/button";
-import { urlOf } from "@valbuild/shared/internal";
+import { urlOf, VAL_AI_SESSION_STORAGE_KEY } from "@valbuild/shared/internal";
 import { Fragment, useMemo, useRef, useState } from "react";
 import { cn } from "./designSystem/cn";
 import { AIChat } from "./AIChat";
 import type { AIChatHandle } from "./AIChat";
 import { useAI } from "../hooks/useAI";
 import { Internal, ModuleFilePath, SourcePath } from "@valbuild/core";
 import { prettifyFilename } from "../utils/prettifyFilename";
-import { useNavigation, VAL_ERRORS_ROUTE } from "./ValRouter";
+import { useNavigation, useSessionParam, VAL_ERRORS_ROUTE } from "./ValRouter";
 import { PublishButton } from "./PublishButton";
 import { Checkbox } from "./designSystem/checkbox";
 import {
@@ -67,6 +67,10 @@ export function ToolsMenu() {
   const committedPatchIds = useCommittedPatches();
   const pendingChanges = currentPatchIds.length - committedPatchIds.size;
   const chatRef = useRef<AIChatHandle | null>(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 (
     <div
       className="flex flex-col h-[100svh] bg-bg-primary"
@@ -160,6 +183,7 @@ export function ToolsMenu() {
             onLoadSession={loadSession}
             onFetchSessions={getSessions}
             onSetSessionName={setSessionName}
+            isLoadingSession={isLoadingSession}
           />
         </div>
       )}