diff --git a/studio/frontend/src/components/ui/alert-dialog.tsx b/studio/frontend/src/components/ui/alert-dialog.tsx index b39a310c49..61327e525e 100644 --- a/studio/frontend/src/components/ui/alert-dialog.tsx +++ b/studio/frontend/src/components/ui/alert-dialog.tsx @@ -42,19 +42,21 @@ function AlertDialogOverlay({ ); } -function AlertDialogContent({ - className, - size = "default", - ...props -}: React.ComponentProps & { - size?: "default" | "sm"; -}) { - return ( - - - & { + size?: "default" | "sm"; + overlayClassName?: string; +}) { + return ( + + + -
+ -
+ ); } diff --git a/studio/frontend/src/features/auth/signup-page.tsx b/studio/frontend/src/features/auth/signup-page.tsx index 5b71fe52ee..3d9b09ff7b 100644 --- a/studio/frontend/src/features/auth/signup-page.tsx +++ b/studio/frontend/src/features/auth/signup-page.tsx @@ -1,4 +1,5 @@ import { LightRays } from "@/components/ui/light-rays"; +import { Card } from "@/components/ui/card"; import { AuthForm } from "./components/auth-form"; export function SignupPage() { @@ -12,9 +13,9 @@ export function SignupPage() { length="70vh" style={{ opacity: 0.4 }} /> -
+ -
+ ); } diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index c79f4b387c..5c1e61b432 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -29,6 +29,10 @@ import { GuidedTour, useGuidedTourController } from "@/features/tour"; import { ChatSettingsPanel } from "./chat-settings-sheet"; import { db } from "./db"; import { useChatModelRuntime } from "./hooks/use-chat-model-runtime"; +import { + clearTrainingCompareHandoff, + getTrainingCompareHandoff, +} from "./lib/training-compare-handoff"; import { ChatRuntimeProvider } from "./runtime-provider"; import { useChatRuntimeStore } from "./stores/chat-runtime-store"; import { @@ -41,6 +45,43 @@ import { ThreadSidebar } from "./thread-sidebar"; import type { ChatView } from "./types"; import { buildChatTourSteps } from "./tour"; +type LoraCandidate = { + id: string; + baseModel: string; + updatedAt?: number; +}; + +function normalizeModelRef(value: string | null | undefined): string { + return value?.trim().toLowerCase() ?? ""; +} + +function pickBestLoraForBase( + loras: LoraCandidate[], + baseModel: string | null, +): LoraCandidate | null { + if (loras.length === 0) return null; + const sorted = [...loras].sort( + (a, b) => (b.updatedAt ?? -1) - (a.updatedAt ?? -1), + ); + const normalizedBase = normalizeModelRef(baseModel); + if (!normalizedBase) return sorted[0]; + + const exact = sorted.find( + (lora) => normalizeModelRef(lora.baseModel) === normalizedBase, + ); + if (exact) return exact; + + const partial = sorted.find((lora) => { + const normalizedLoraBase = normalizeModelRef(lora.baseModel); + if (!normalizedLoraBase) return false; + return ( + normalizedLoraBase.includes(normalizedBase) || + normalizedBase.includes(normalizedLoraBase) + ); + }); + return partial ?? sorted[0]; +} + const SingleContent = memo(function SingleContent({ threadId, newThreadNonce, @@ -207,7 +248,9 @@ export function ChatPage(): ReactElement { const [modelSelectorOpen, setModelSelectorOpen] = useState(false); const [modelSelectorLocked, setModelSelectorLocked] = useState(false); const [sidebarOpen, setSidebarOpen] = useState(true); - const viewBeforeCompareRef = useRef(null); + const [viewBeforeCompare, setViewBeforeCompare] = useState( + null, + ); const inferenceParams = useChatRuntimeStore((state) => state.params); const setInferenceParams = useChatRuntimeStore((state) => state.setParams); const autoTitle = useChatRuntimeStore((state) => state.autoTitle); @@ -216,6 +259,13 @@ export function ChatPage(): ReactElement { const lorasFromStore = useChatRuntimeStore((state) => state.loras); const modelsError = useChatRuntimeStore((state) => state.modelsError); const { refresh, selectModel, ejectModel } = useChatModelRuntime(); + const refreshRef = useRef(refresh); + const selectModelRef = useRef(selectModel); + + useEffect(() => { + refreshRef.current = refresh; + selectModelRef.current = selectModel; + }, [refresh, selectModel]); const canCompare = useMemo(() => { const selected = inferenceParams.checkpoint; if (!selected) return false; @@ -262,18 +312,15 @@ export function ChatPage(): ReactElement { const openSidebar = useCallback(() => setSidebarOpen(true), []); const enterCompare = useCallback(() => { - if (viewBeforeCompareRef.current == null) { - viewBeforeCompareRef.current = view; - } + setViewBeforeCompare((prev) => prev ?? view); setView({ mode: "compare", pairId: crypto.randomUUID() }); }, [view]); const exitCompare = useCallback(() => { - const prev = viewBeforeCompareRef.current; - if (!prev) return; - viewBeforeCompareRef.current = null; - setView(prev); - }, []); + if (!viewBeforeCompare) return; + setView(viewBeforeCompare); + setViewBeforeCompare(null); + }, [viewBeforeCompare]); const models = useMemo( () => @@ -297,9 +344,69 @@ export function ChatPage(): ReactElement { ); useEffect(() => { + if (getTrainingCompareHandoff()) return; void refresh(); }, [refresh]); + useEffect(() => { + const handoff = getTrainingCompareHandoff(); + if (!handoff) return; + console.info("[chat-handoff] received", handoff); + function clearHandoff(): void { + clearTrainingCompareHandoff(); + } + + let canceled = false; + void (async () => { + try { + console.info("[chat-handoff] refreshing models+loras"); + await refreshRef.current(); + if (canceled) return; + + const state = useChatRuntimeStore.getState(); + const targetLora = pickBestLoraForBase(state.loras, handoff.baseModel); + if (targetLora) { + console.info("[chat-handoff] loading lora", { + id: targetLora.id, + baseModel: targetLora.baseModel, + }); + await selectModelRef.current({ id: targetLora.id, isLora: true }); + if (canceled) return; + setView({ mode: "compare", pairId: crypto.randomUUID() }); + clearHandoff(); + console.info("[chat-handoff] loaded lora + opened compare"); + return; + } + + if ( + handoff.baseModel && + state.models.some((model) => model.id === handoff.baseModel) + ) { + console.info("[chat-handoff] no lora match, loading base", { + id: handoff.baseModel, + }); + await selectModelRef.current({ id: handoff.baseModel, isLora: false }); + if (canceled) return; + } else { + console.warn("[chat-handoff] no lora/base match found", { + requestedBaseModel: handoff.baseModel, + loraCount: state.loras.length, + modelCount: state.models.length, + }); + } + clearHandoff(); + console.info("[chat-handoff] completed"); + } catch (error) { + console.error("[chat-handoff] failed", error); + clearHandoff(); + } + })(); + + return () => { + canceled = true; + }; + }, []); + const tourSteps = useMemo( () => buildChatTourSteps({ @@ -332,8 +439,11 @@ export function ChatPage(): ReactElement { useEffect(() => { if (tour.open) return; if (!modelSelectorLocked) return; - setModelSelectorLocked(false); - setModelSelectorOpen(false); + const timeoutId = window.setTimeout(() => { + setModelSelectorLocked(false); + setModelSelectorOpen(false); + }, 0); + return () => window.clearTimeout(timeoutId); }, [modelSelectorLocked, tour.open]); return ( diff --git a/studio/frontend/src/features/chat/index.ts b/studio/frontend/src/features/chat/index.ts index b7eaaf83ae..521a0fe27f 100644 --- a/studio/frontend/src/features/chat/index.ts +++ b/studio/frontend/src/features/chat/index.ts @@ -7,3 +7,4 @@ export { } from "./chat-settings-sheet"; export { useChatRuntimeStore } from "./stores/chat-runtime-store"; export { useChatModelRuntime } from "./hooks/use-chat-model-runtime"; +export { setTrainingCompareHandoff } from "./lib/training-compare-handoff"; diff --git a/studio/frontend/src/features/chat/lib/training-compare-handoff.ts b/studio/frontend/src/features/chat/lib/training-compare-handoff.ts new file mode 100644 index 0000000000..1ceaf7a00b --- /dev/null +++ b/studio/frontend/src/features/chat/lib/training-compare-handoff.ts @@ -0,0 +1,53 @@ +const TRAINING_COMPARE_HANDOFF_KEY = "chat:training-compare-handoff:v1"; +const HANDOFF_MAX_AGE_MS = 15 * 60 * 1000; + +export type TrainingCompareHandoff = { + intent: "compare"; + baseModel: string | null; + requestedAt: number; +}; + +export function setTrainingCompareHandoff(baseModel: string | null): void { + if (typeof window === "undefined") return; + + const payload: TrainingCompareHandoff = { + intent: "compare", + baseModel, + requestedAt: Date.now(), + }; + window.sessionStorage.setItem( + TRAINING_COMPARE_HANDOFF_KEY, + JSON.stringify(payload), + ); +} + +export function getTrainingCompareHandoff(): TrainingCompareHandoff | null { + if (typeof window === "undefined") return null; + + const raw = window.sessionStorage.getItem(TRAINING_COMPARE_HANDOFF_KEY); + if (!raw) return null; + + try { + const parsed = JSON.parse(raw) as Partial; + if (parsed.intent !== "compare") return null; + if (typeof parsed.requestedAt !== "number") return null; + if (Date.now() - parsed.requestedAt > HANDOFF_MAX_AGE_MS) { + clearTrainingCompareHandoff(); + return null; + } + return { + intent: "compare", + baseModel: + typeof parsed.baseModel === "string" ? parsed.baseModel : null, + requestedAt: parsed.requestedAt, + }; + } catch { + clearTrainingCompareHandoff(); + return null; + } +} + +export function clearTrainingCompareHandoff(): void { + if (typeof window === "undefined") return; + window.sessionStorage.removeItem(TRAINING_COMPARE_HANDOFF_KEY); +} diff --git a/studio/frontend/src/features/export/export-page.tsx b/studio/frontend/src/features/export/export-page.tsx index 5c11a5914b..af670aae16 100644 --- a/studio/frontend/src/features/export/export-page.tsx +++ b/studio/frontend/src/features/export/export-page.tsx @@ -278,7 +278,7 @@ export function ExportPage() {
{/* Training run dropdown */} -
+
{/* Checkpoint dropdown */} -
+