From b0e7479699b4dd163f48c4bce9809af560174406 Mon Sep 17 00:00:00 2001 From: Shine1i Date: Tue, 17 Feb 2026 20:42:43 +0100 Subject: [PATCH] fix chat compare handoff: auto-load trained lora, stop refresh loop, add debug logs --- .../frontend/src/features/chat/chat-page.tsx | 132 ++++++++++++++++-- studio/frontend/src/features/chat/index.ts | 1 + .../studio/sections/progress-section.tsx | 12 +- 3 files changed, 131 insertions(+), 14 deletions(-) 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/studio/sections/progress-section.tsx b/studio/frontend/src/features/studio/sections/progress-section.tsx index 2d075e870b..3752a0fb25 100644 --- a/studio/frontend/src/features/studio/sections/progress-section.tsx +++ b/studio/frontend/src/features/studio/sections/progress-section.tsx @@ -31,12 +31,14 @@ import { } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { useState, type ReactElement, type ReactNode } from "react"; -import { Link } from "@tanstack/react-router"; +import { Link, useNavigate } from "@tanstack/react-router"; import { useShallow } from "zustand/react/shallow"; import { useGpuUtilization } from "@/hooks"; +import { setTrainingCompareHandoff } from "@/features/chat"; import { formatDuration, formatNumber, phaseColors, phaseLabel } from "./progress-section-lib"; export function ProgressSection(): ReactElement { + const navigate = useNavigate(); const runtime = useTrainingRuntimeStore( useShallow((state) => ({ phase: state.phase, @@ -105,6 +107,10 @@ export function ProgressSection(): ReactElement { const showHalfwayHint = runtime.phase === "training" && pct >= 50 && pct < 100; const showCompletedHint = runtime.phase === "completed"; + const handleCompareInChat = () => { + setTrainingCompareHandoff(config.selectedModel); + void navigate({ to: "/chat" }); + }; const stoppedLoss = getDisplayMetric( runtime.isTrainingRunning, @@ -265,8 +271,8 @@ export function ProgressSection(): ReactElement {

{showCompletedHint && (
-