From 4fd4d2fe760301058629a749d1a71d5f41ea9c31 Mon Sep 17 00:00:00 2001
From: Shine1i <wasimysdev@gmail.com>
Date: Mon, 16 Feb 2026 01:44:26 +0100
Subject: [PATCH 1/4] feat: implement dataset mapping UI and preview dialog

---
 .../dataset-preview-dialog-mapping.tsx        | 162 +++++++++++
 .../sections/dataset-preview-dialog-utils.ts  |  55 ++++
 .../sections/dataset-preview-dialog.tsx       | 260 ++++++++++--------
 .../studio/sections/dataset-section.tsx       |  14 +-
 .../src/features/studio/studio-page.tsx       |  24 ++
 .../src/features/training/api/datasets-api.ts |  37 +++
 .../src/features/training/api/mappers.ts      |  10 +
 .../training/hooks/use-training-actions.ts    |  53 ++++
 .../frontend/src/features/training/index.ts   |   1 +
 .../stores/dataset-preview-dialog-store.ts    |  33 +++
 .../training/stores/training-config-store.ts  |  19 +-
 .../src/features/training/types/api.ts        |   1 +
 .../src/features/training/types/config.ts     |   7 +
 .../src/features/training/types/datasets.ts   |  13 +
 studio/frontend/src/main.tsx                  |  24 +-
 15 files changed, 576 insertions(+), 137 deletions(-)
 create mode 100644 studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx
 create mode 100644 studio/frontend/src/features/studio/sections/dataset-preview-dialog-utils.ts
 create mode 100644 studio/frontend/src/features/training/api/datasets-api.ts
 create mode 100644 studio/frontend/src/features/training/stores/dataset-preview-dialog-store.ts
 create mode 100644 studio/frontend/src/features/training/types/datasets.ts

diff --git a/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx b/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx
new file mode 100644
index 0000000000..5f8833b742
--- /dev/null
+++ b/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx
@@ -0,0 +1,162 @@
+import { Badge } from "@/components/ui/badge";
+import { Button } from "@/components/ui/button";
+import { Checkbox } from "@/components/ui/checkbox";
+import { AlertCircleIcon } from "@hugeicons/core-free-icons";
+import { HugeiconsIcon } from "@hugeicons/react";
+import type { CheckFormatResponse } from "@/features/training/types/datasets";
+
+export function HeaderPick({
+  label,
+  checked,
+  onCheckedChange,
+}: {
+  label: string;
+  checked: boolean;
+  onCheckedChange: (checked: boolean) => void;
+}) {
+  return (
+    <label className="flex items-center gap-1.5 text-[10px] text-muted-foreground cursor-pointer select-none">
+      <Checkbox
+        checked={checked}
+        onCheckedChange={(v) => onCheckedChange(v === true)}
+        aria-label={label}
+        className="h-3.5 w-3.5"
+      />
+      <span>{label}</span>
+    </label>
+  );
+}
+
+export function DatasetMappingCard({
+  leftLabel,
+  rightLabel,
+  mappingOk,
+  input,
+  output,
+}: {
+  leftLabel: string;
+  rightLabel: string;
+  mappingOk: boolean;
+  input: string | null;
+  output: string | null;
+}) {
+  return (
+    <div className="rounded-xl corner-squircle ring-1 ring-amber-200/70 bg-amber-50/70 px-5 py-4 mb-4 text-amber-950 dark:ring-amber-900/50 dark:bg-amber-950/30 dark:text-amber-50">
+      <div className="flex items-start gap-3">
+        <div className="rounded-xl corner-squircle bg-amber-500/15 p-2 shrink-0">
+          <HugeiconsIcon
+            icon={AlertCircleIcon}
+            className="size-4 text-amber-700 dark:text-amber-300"
+          />
+        </div>
+        <div className="min-w-0">
+          <p className="text-sm font-semibold tracking-tight">Map dataset columns</p>
+          <p className="text-xs text-amber-800/80 dark:text-amber-200/80 mt-0.5">
+            We couldn&apos;t auto-detect the format. Pick the{" "}
+            {leftLabel.toLowerCase()} column and the {rightLabel.toLowerCase()}{" "}
+            column. We&apos;ll convert it to a supported format automatically.
+          </p>
+          <div className="mt-3 flex flex-wrap items-center gap-2">
+            <Badge
+              variant="outline"
+              className="h-6 text-[11px] bg-white/60 dark:bg-transparent"
+            >
+              {leftLabel}: <span className="font-mono">{input ?? "--"}</span>
+            </Badge>
+            <Badge
+              variant="outline"
+              className="h-6 text-[11px] bg-white/60 dark:bg-transparent"
+            >
+              {rightLabel}: <span className="font-mono">{output ?? "--"}</span>
+            </Badge>
+          </div>
+          {!mappingOk && (
+            <p className="mt-2 text-xs text-amber-800/80 dark:text-amber-200/80">
+              Select exactly 1 {leftLabel.toLowerCase()} and 1{" "}
+              {rightLabel.toLowerCase()} column to continue.
+            </p>
+          )}
+        </div>
+      </div>
+    </div>
+  );
+}
+
+export function DatasetMappingFooter({
+  leftLabel,
+  rightLabel,
+  mappingOk,
+  isStarting,
+  startError,
+  onCancel,
+  onStartTraining,
+}: {
+  leftLabel: string;
+  rightLabel: string;
+  mappingOk: boolean;
+  isStarting: boolean;
+  startError: string | null;
+  onCancel: () => void;
+  onStartTraining: () => Promise<void>;
+}) {
+  return (
+    <div className="mt-3 flex flex-col gap-2">
+      <div className="flex items-center justify-between gap-3">
+        <p className="text-[11px] text-muted-foreground/70 leading-relaxed">
+          Tip: you can click {leftLabel} / {rightLabel} in the headers.
+        </p>
+        <div className="flex items-center gap-2">
+          <Button
+            variant="outline"
+            size="sm"
+            className="cursor-pointer"
+            onClick={onCancel}
+          >
+            Cancel
+          </Button>
+          <Button
+            size="sm"
+            className="cursor-pointer"
+            disabled={!mappingOk || isStarting}
+            onClick={() => void onStartTraining()}
+          >
+            {isStarting ? "Starting..." : "Continue"}
+          </Button>
+        </div>
+      </div>
+
+      {startError && (
+        <p className="text-xs text-red-500 leading-relaxed text-center">
+          {startError}
+        </p>
+      )}
+    </div>
+  );
+}
+
+export function deriveDefaultMapping(
+  data: CheckFormatResponse,
+  isVlm: boolean,
+): { input: string | null; output: string | null } {
+  if (isVlm) {
+    const input =
+      data.detected_image_column ?? pickRole(data.suggested_mapping, "image");
+    const output =
+      data.detected_text_column ?? pickRole(data.suggested_mapping, "text");
+    return { input: input ?? null, output: output ?? null };
+  }
+  const input = pickRole(data.suggested_mapping, "user");
+  const output = pickRole(data.suggested_mapping, "assistant");
+  return { input: input ?? null, output: output ?? null };
+}
+
+function pickRole(
+  mapping: Record<string, string> | null | undefined,
+  role: string,
+): string | null {
+  if (!mapping) return null;
+  for (const [col, mappedRole] of Object.entries(mapping)) {
+    if (mappedRole === role) return col;
+  }
+  return null;
+}
diff --git a/studio/frontend/src/features/studio/sections/dataset-preview-dialog-utils.ts b/studio/frontend/src/features/studio/sections/dataset-preview-dialog-utils.ts
new file mode 100644
index 0000000000..735a3abbff
--- /dev/null
+++ b/studio/frontend/src/features/studio/sections/dataset-preview-dialog-utils.ts
@@ -0,0 +1,55 @@
+type PreviewImagePayload = {
+  type: "image";
+  mime?: string;
+  width?: number;
+  height?: number;
+  data?: string;
+};
+
+export function formatCell(value: unknown): string {
+  if (value == null) return "";
+  if (typeof value === "string") return value;
+  if (typeof value === "number" || typeof value === "boolean") return String(value);
+  if (Array.isArray(value) || typeof value === "object")
+    return JSON.stringify(value).slice(0, 500);
+  return String(value);
+}
+
+function isPreviewImagePayload(value: unknown): value is PreviewImagePayload {
+  if (!value || typeof value !== "object") return false;
+  const record = value as Record<string, unknown>;
+  return (
+    record.type === "image" &&
+    typeof record.data === "string" &&
+    record.data.length > 0
+  );
+}
+
+export function collectPreviewImages(value: unknown): PreviewImagePayload[] {
+  const images: PreviewImagePayload[] = [];
+  const stack: unknown[] = [value];
+  let steps = 0;
+
+  while (stack.length > 0 && steps < 200) {
+    steps += 1;
+    const current = stack.pop();
+    if (isPreviewImagePayload(current)) {
+      images.push(current);
+      continue;
+    }
+
+    if (Array.isArray(current)) {
+      for (const item of current) stack.push(item);
+      continue;
+    }
+
+    if (current && typeof current === "object") {
+      for (const nested of Object.values(current as Record<string, unknown>)) {
+        stack.push(nested);
+      }
+    }
+  }
+
+  return images;
+}
+
diff --git a/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx b/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx
index f145b17f61..11de320147 100644
--- a/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx
+++ b/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx
@@ -11,30 +11,21 @@ import { Spinner } from "@/components/ui/spinner";
 import { Database02Icon, AlertCircleIcon } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
 import { type ReactNode, useEffect, useMemo, useState } from "react";
+import { useTrainingActions, useTrainingConfigStore } from "@/features/training";
+import { checkDatasetFormat } from "@/features/training/api/datasets-api";
+import type { CheckFormatResponse } from "@/features/training/types/datasets";
+import { collectPreviewImages, formatCell } from "./dataset-preview-dialog-utils";
+import {
+  DatasetMappingCard,
+  DatasetMappingFooter,
+  HeaderPick,
+  deriveDefaultMapping,
+} from "./dataset-preview-dialog-mapping";
 
 // ---------------------------------------------------------------------------
-// Types (matches CheckFormatResponse from backend)
+// Types
 // ---------------------------------------------------------------------------
 
-type CheckFormatResponse = {
-  requires_manual_mapping: boolean;
-  detected_format: string;
-  columns: string[];
-  suggested_mapping?: Record<string, string> | null;
-  detected_image_column?: string | null;
-  detected_text_column?: string | null;
-  preview_samples?: Record<string, unknown>[] | null;
-  total_rows?: number | null;
-};
-
-type PreviewImagePayload = {
-  type: "image";
-  mime?: string;
-  width?: number;
-  height?: number;
-  data?: string;
-};
-
 type DatasetPreviewDialogProps = {
   open: boolean;
   onOpenChange: (open: boolean) => void;
@@ -42,37 +33,11 @@ type DatasetPreviewDialogProps = {
   hfToken: string | null;
   datasetSubset?: string | null;
   datasetSplit?: string | null;
+  mode?: "preview" | "mapping";
+  initialData?: CheckFormatResponse | null;
+  isVlm?: boolean;
 };
 
-// ---------------------------------------------------------------------------
-// API -- uses existing /check-format endpoint
-// ---------------------------------------------------------------------------
-
-// TODO(backend): Needs to accept `config` and `split` fields (see #37).
-// The frontend already sends them in the request below.
-async function fetchCheckFormat(
-  datasetName: string,
-  hfToken: string | null,
-  subset?: string | null,
-  split?: string | null,
-): Promise<CheckFormatResponse> {
-  const res = await fetch("/api/datasets/check-format", {
-    method: "POST",
-    headers: { "Content-Type": "application/json" },
-    body: JSON.stringify({
-      dataset_name: datasetName,
-      hf_token: hfToken || undefined,
-      config: subset || undefined,
-      split: split || "train",
-    }),
-  });
-  if (!res.ok) {
-    const body = await res.json().catch(() => null);
-    throw new Error(body?.detail || `Request failed (${res.status})`);
-  }
-  return res.json();
-}
-
 // ---------------------------------------------------------------------------
 // Component
 // ---------------------------------------------------------------------------
@@ -84,22 +49,47 @@ export function DatasetPreviewDialog({
   hfToken,
   datasetSubset,
   datasetSplit,
+  mode = "preview",
+  initialData,
+  isVlm = false,
 }: DatasetPreviewDialogProps) {
   const [data, setData] = useState<CheckFormatResponse | null>(null);
   const [loading, setLoading] = useState(false);
   const [error, setError] = useState<string | null>(null);
 
+  const manualMapping = useTrainingConfigStore((s) => s.datasetManualMapping);
+  const setManualMapping = useTrainingConfigStore((s) => s.setDatasetManualMapping);
+  const { isStarting, startError, startTrainingRun } = useTrainingActions();
+
+  const mappingEnabled = !!data?.requires_manual_mapping;
+  const showMappingFooter = mode === "mapping" && mappingEnabled;
+  const mappingOk = !!manualMapping.input && !!manualMapping.output;
+  const leftLabel = isVlm ? "Image" : "Input";
+  const rightLabel = isVlm ? "Text" : "Output";
+
   useEffect(() => {
     if (!open || !datasetName) {
       setData(null);
       setError(null);
       return;
     }
+    if (initialData) {
+      setData(initialData);
+      setError(null);
+      setLoading(false);
+      return;
+    }
     let cancelled = false;
     setLoading(true);
     setError(null);
 
-    fetchCheckFormat(datasetName, hfToken, datasetSubset, datasetSplit)
+    checkDatasetFormat({
+      datasetName,
+      hfToken,
+      subset: datasetSubset,
+      split: datasetSplit,
+      isVlm,
+    })
       .then((res) => {
         if (!cancelled) {
           setData(res);
@@ -116,7 +106,24 @@ export function DatasetPreviewDialog({
     return () => {
       cancelled = true;
     };
-  }, [open, datasetName, hfToken, datasetSubset, datasetSplit]);
+  }, [open, datasetName, hfToken, datasetSubset, datasetSplit, isVlm, initialData]);
+
+  useEffect(() => {
+    if (!open || !datasetName || !data?.requires_manual_mapping) return;
+    if (manualMapping.input || manualMapping.output) return;
+    const derived = deriveDefaultMapping(data, isVlm);
+    if (!derived.input && !derived.output) return;
+    setManualMapping(derived);
+  }, [
+    open,
+    datasetName,
+    data?.requires_manual_mapping,
+    isVlm,
+    manualMapping.input,
+    manualMapping.output,
+    setManualMapping,
+    data,
+  ]);
 
   const rows = data?.preview_samples ?? [];
   const columns = data?.columns ?? [];
@@ -140,9 +147,39 @@ export function DatasetPreviewDialog({
     return columns.map((colName) => ({
       accessorKey: colName,
       header: () => (
-        <span className="font-heading text-[13px] font-semibold tracking-tight text-foreground">
-          {colName}
-        </span>
+        <div className="flex flex-col gap-2">
+          <span className="font-heading text-[13px] font-semibold tracking-tight text-foreground">
+            {colName}
+          </span>
+          {mappingEnabled && (
+            <div className="flex items-center gap-3">
+              {(manualMapping.input == null || manualMapping.input === colName) && (
+                <HeaderPick
+                  label={leftLabel}
+                  checked={manualMapping.input === colName}
+                  onCheckedChange={(checked) => {
+                    setManualMapping({
+                      input: checked ? colName : null,
+                      output: manualMapping.output,
+                    });
+                  }}
+                />
+              )}
+              {(manualMapping.output == null || manualMapping.output === colName) && (
+                <HeaderPick
+                  label={rightLabel}
+                  checked={manualMapping.output === colName}
+                  onCheckedChange={(checked) => {
+                    setManualMapping({
+                      input: manualMapping.input,
+                      output: checked ? colName : null,
+                    });
+                  }}
+                />
+              )}
+            </div>
+          )}
+        </div>
       ),
       cell: ({ getValue }: { getValue: () => unknown }) => {
         const value = getValue();
@@ -196,7 +233,15 @@ export function DatasetPreviewDialog({
         );
       },
     }));
-  }, [columns]);
+  }, [
+    columns,
+    manualMapping.input,
+    manualMapping.output,
+    setManualMapping,
+    mappingEnabled,
+    leftLabel,
+    rightLabel,
+  ]);
 
   return (
     <Dialog open={open} onOpenChange={onOpenChange}>
@@ -266,10 +311,10 @@ export function DatasetPreviewDialog({
                       : "--"
                   }
                 />
-                <MetaRow
-                  label="Columns"
-                  value={
-                    <span className="flex items-center gap-1.5 flex-wrap">
+              <MetaRow
+                label="Columns"
+                value={
+                  <span className="flex items-center gap-1.5 flex-wrap">
                       {columns.map((col) => (
                         <Badge
                           key={col}
@@ -284,18 +329,51 @@ export function DatasetPreviewDialog({
                 />
               </div>
 
+              {mappingEnabled && (
+                <DatasetMappingCard
+                  leftLabel={leftLabel}
+                  rightLabel={rightLabel}
+                  mappingOk={mappingOk}
+                  input={manualMapping.input}
+                  output={manualMapping.output}
+                />
+              )}
+
               {/* Data table */}
               <div className="flex-1 min-h-0 rounded-xl corner-squircle ring-1 ring-border/60 overflow-auto">
                 <DataTable columns={tableColumns} data={rows} />
               </div>
 
               {/* Footer */}
-              <p className="text-[11px] text-muted-foreground/60 mt-3 text-center tabular-nums">
-                Showing {rows.length}
-                {data.total_rows != null &&
-                  ` of ${data.total_rows.toLocaleString()}`}{" "}
-                rows
-              </p>
+              <div className="mt-3">
+                <p className="text-[11px] text-muted-foreground/60 text-center tabular-nums">
+                  Showing {rows.length}
+                  {data.total_rows != null &&
+                    ` of ${data.total_rows.toLocaleString()}`}{" "}
+                  rows
+                </p>
+
+                {mode === "preview" && mappingEnabled && (
+                  <p className="mt-2 text-[11px] text-muted-foreground/70 text-center">
+                    Mapping is saved automatically. You can start training anytime.
+                  </p>
+                )}
+
+                {showMappingFooter && (
+                  <DatasetMappingFooter
+                    leftLabel={leftLabel}
+                    rightLabel={rightLabel}
+                    mappingOk={mappingOk}
+                    isStarting={isStarting}
+                    startError={startError}
+                    onCancel={() => onOpenChange(false)}
+                    onStartTraining={async () => {
+                      const ok = await startTrainingRun();
+                      if (ok) onOpenChange(false);
+                    }}
+                  />
+                )}
+              </div>
             </>
           )}
         </div>
@@ -325,54 +403,4 @@ function MetaRow({
   );
 }
 
-// ---------------------------------------------------------------------------
-// Helpers
-// ---------------------------------------------------------------------------
-
-function formatCell(value: unknown): string {
-  if (value == null) return "";
-  if (typeof value === "string") return value;
-  if (typeof value === "number" || typeof value === "boolean")
-    return String(value);
-  if (Array.isArray(value) || typeof value === "object")
-    return JSON.stringify(value).slice(0, 500);
-  return String(value);
-}
-
-function isPreviewImagePayload(value: unknown): value is PreviewImagePayload {
-  if (!value || typeof value !== "object") return false;
-  const record = value as Record<string, unknown>;
-  return (
-    record.type === "image" &&
-    typeof record.data === "string" &&
-    record.data.length > 0
-  );
-}
-
-function collectPreviewImages(value: unknown): PreviewImagePayload[] {
-  const images: PreviewImagePayload[] = [];
-  const stack: unknown[] = [value];
-  let steps = 0;
-
-  while (stack.length > 0 && steps < 200) {
-    steps += 1;
-    const current = stack.pop();
-    if (isPreviewImagePayload(current)) {
-      images.push(current);
-      continue;
-    }
-
-    if (Array.isArray(current)) {
-      for (const item of current) stack.push(item);
-      continue;
-    }
-
-    if (current && typeof current === "object") {
-      for (const nested of Object.values(current as Record<string, unknown>)) {
-        stack.push(nested);
-      }
-    }
-  }
-
-  return images;
-}
+// mapping UI extracted to ./dataset-preview-dialog-mapping.tsx
diff --git a/studio/frontend/src/features/studio/sections/dataset-section.tsx b/studio/frontend/src/features/studio/sections/dataset-section.tsx
index ba9b96373a..3d7e3f810f 100644
--- a/studio/frontend/src/features/studio/sections/dataset-section.tsx
+++ b/studio/frontend/src/features/studio/sections/dataset-section.tsx
@@ -30,6 +30,7 @@ import {
 import { formatCompact } from "@/lib/utils";
 import {
   HfDatasetSubsetSplitSelectors,
+  useDatasetPreviewDialogStore,
   useTrainingConfigStore,
 } from "@/features/training";
 import {
@@ -43,7 +44,6 @@ import {
 import { HugeiconsIcon } from "@hugeicons/react";
 import { useMemo, useRef, useState } from "react";
 import { useShallow } from "zustand/react/shallow";
-import { DatasetPreviewDialog } from "./dataset-preview-dialog";
 
 function isLikelyLocalDatasetRef(value: string) {
   return (
@@ -81,7 +81,7 @@ export function DatasetSection() {
   );
 
   const [inputValue, setInputValue] = useState("");
-  const [previewOpen, setPreviewOpen] = useState(false);
+  const openPreview = useDatasetPreviewDialogStore((s) => s.openPreview);
   const selectingRef = useRef(false);
   const debouncedQuery = useDebouncedValue(inputValue);
 
@@ -341,7 +341,7 @@ export function DatasetSection() {
             size="sm"
             className="cursor-pointer gap-1.5"
             disabled={!dataset}
-            onClick={() => setPreviewOpen(true)}
+            onClick={() => openPreview()}
           >
             <HugeiconsIcon icon={ViewIcon} className="size-3.5" />
             View dataset
@@ -349,14 +349,6 @@ export function DatasetSection() {
         </div>
       </div>
       </SectionCard>
-      <DatasetPreviewDialog
-        open={previewOpen}
-        onOpenChange={setPreviewOpen}
-        datasetName={dataset}
-        hfToken={hfToken}
-        datasetSubset={datasetSubset}
-        datasetSplit={datasetSplit}
-      />
     </div>
   );
 }
diff --git a/studio/frontend/src/features/studio/studio-page.tsx b/studio/frontend/src/features/studio/studio-page.tsx
index 6cdf5225ec..b7018dd843 100644
--- a/studio/frontend/src/features/studio/studio-page.tsx
+++ b/studio/frontend/src/features/studio/studio-page.tsx
@@ -2,6 +2,8 @@ import { Button } from "@/components/ui/button";
 import {
   shouldShowTrainingView,
   useTrainingActions,
+  useDatasetPreviewDialogStore,
+  useTrainingConfigStore,
   useTrainingRuntimeLifecycle,
   useTrainingRuntimeStore,
 } from "@/features/training";
@@ -10,6 +12,7 @@ import { studioTourSteps } from "@/features/studio/tour";
 import { ArrowLeft01Icon } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
 import { type ReactElement, useEffect, useState } from "react";
+import { DatasetPreviewDialog } from "./sections/dataset-preview-dialog";
 import { DatasetSection } from "./sections/dataset-section";
 import { ModelSection } from "./sections/model-section";
 import { ParamsSection } from "./sections/params-section";
@@ -26,6 +29,11 @@ export function StudioPage(): ReactElement {
   const isHydratingRuntime = useTrainingRuntimeStore((state) => state.isHydrating);
   const hasHydratedRuntime = useTrainingRuntimeStore((state) => state.hasHydrated);
   const { dismissTrainingRun } = useTrainingActions();
+  const config = useTrainingConfigStore();
+  const dialogOpen = useDatasetPreviewDialogStore((s) => s.open);
+  const dialogMode = useDatasetPreviewDialogStore((s) => s.mode);
+  const dialogInitial = useDatasetPreviewDialogStore((s) => s.initialData);
+  const closeDialog = useDatasetPreviewDialogStore((s) => s.close);
 
   const canGoBack = runtimePhase === "stopped" || runtimePhase === "error";
   const tourEnabled = hasHydratedRuntime && !isHydratingRuntime && !showTrainingView;
@@ -48,6 +56,22 @@ export function StudioPage(): ReactElement {
           onComplete={() => localStorage.setItem(STUDIO_TOUR_KEY, "done")}
         />
 
+        <DatasetPreviewDialog
+          open={dialogOpen}
+          onOpenChange={(open) => {
+            if (!open) closeDialog();
+          }}
+          datasetName={
+            config.datasetSource === "huggingface" ? config.dataset : config.uploadedFile
+          }
+          hfToken={config.hfToken.trim() || null}
+          datasetSubset={config.datasetSubset}
+          datasetSplit={config.datasetSplit}
+          mode={dialogMode}
+          initialData={dialogInitial}
+          isVlm={config.modelType === "vision"}
+        />
+
         {canGoBack && (
           <Button
             variant="ghost"
diff --git a/studio/frontend/src/features/training/api/datasets-api.ts b/studio/frontend/src/features/training/api/datasets-api.ts
new file mode 100644
index 0000000000..37f3744701
--- /dev/null
+++ b/studio/frontend/src/features/training/api/datasets-api.ts
@@ -0,0 +1,37 @@
+import type { CheckFormatResponse } from "../types/datasets";
+
+type CheckDatasetFormatArgs = {
+  datasetName: string;
+  hfToken: string | null;
+  subset?: string | null;
+  split?: string | null;
+  isVlm?: boolean;
+};
+
+export async function checkDatasetFormat({
+  datasetName,
+  hfToken,
+  subset,
+  split,
+  isVlm,
+}: CheckDatasetFormatArgs): Promise<CheckFormatResponse> {
+  const res = await fetch("/api/datasets/check-format", {
+    method: "POST",
+    headers: { "Content-Type": "application/json" },
+    body: JSON.stringify({
+      dataset_name: datasetName,
+      hf_token: hfToken || undefined,
+      config: subset || undefined, // backend currently ignores, safe to send
+      split: split || "train",
+      is_vlm: !!isVlm,
+    }),
+  });
+
+  if (!res.ok) {
+    const body = await res.json().catch(() => null);
+    throw new Error(body?.detail || `Request failed (${res.status})`);
+  }
+
+  return res.json();
+}
+
diff --git a/studio/frontend/src/features/training/api/mappers.ts b/studio/frontend/src/features/training/api/mappers.ts
index 990862a27e..ed2c3b752d 100644
--- a/studio/frontend/src/features/training/api/mappers.ts
+++ b/studio/frontend/src/features/training/api/mappers.ts
@@ -14,6 +14,15 @@ export function buildTrainingStartPayload(
   const adapterMethod = config.trainingMethod !== "full";
   const isQlorMethod = config.trainingMethod === "qlora";
   const hfDataset = config.datasetSource === "huggingface" ? config.dataset : null;
+  const manual = config.datasetManualMapping;
+  const isVlm = config.modelType === "vision";
+  const customFormatMapping =
+    manual.input && manual.output
+      ? {
+          [manual.input]: isVlm ? "image" : "user",
+          [manual.output]: isVlm ? "text" : "assistant",
+        }
+      : undefined;
 
   return {
     model_name: config.selectedModel ?? "",
@@ -26,6 +35,7 @@ export function buildTrainingStartPayload(
     hf_dataset_split: hfDataset ? config.datasetSplit : null,
     local_datasets: [],
     format_type: config.datasetFormat,
+    custom_format_mapping: customFormatMapping,
     num_epochs: config.epochs,
     learning_rate: String(config.learningRate),
     batch_size: config.batchSize,
diff --git a/studio/frontend/src/features/training/hooks/use-training-actions.ts b/studio/frontend/src/features/training/hooks/use-training-actions.ts
index 4ace9dd5ed..27e36f9765 100644
--- a/studio/frontend/src/features/training/hooks/use-training-actions.ts
+++ b/studio/frontend/src/features/training/hooks/use-training-actions.ts
@@ -1,8 +1,10 @@
 import { useCallback } from "react";
 import { useTrainingConfigStore } from "../stores/training-config-store";
 import { useTrainingRuntimeStore } from "../stores/training-runtime-store";
+import { useDatasetPreviewDialogStore } from "../stores/dataset-preview-dialog-store";
 import { startTraining, stopTraining, resetTraining } from "../api/train-api";
 import { buildTrainingStartPayload } from "../api/mappers";
+import { checkDatasetFormat } from "../api/datasets-api";
 import { syncTrainingRuntimeFromBackend } from "../lib/sync-runtime";
 import { validateTrainingConfig } from "../lib/validation";
 
@@ -13,6 +15,7 @@ export function useTrainingActions() {
   const startTrainingRun = useCallback(async (): Promise<boolean> => {
     const config = useTrainingConfigStore.getState();
     const runtimeStore = useTrainingRuntimeStore.getState();
+    const dialogStore = useDatasetPreviewDialogStore.getState();
 
     runtimeStore.setStartError(null);
     const validation = validateTrainingConfig(config);
@@ -24,6 +27,45 @@ export function useTrainingActions() {
     runtimeStore.setStarting(true);
 
     try {
+      const datasetName =
+        config.datasetSource === "huggingface" ? config.dataset : config.uploadedFile;
+      const isVlm = config.modelType === "vision";
+
+      if (datasetName) {
+        const check = await checkDatasetFormat({
+          datasetName,
+          hfToken: config.hfToken.trim() || null,
+          subset: config.datasetSubset,
+          split: config.datasetSplit,
+          isVlm,
+        });
+
+        if (check.requires_manual_mapping) {
+          const existing = useTrainingConfigStore.getState().datasetManualMapping;
+          const hasMapping = !!existing.input && !!existing.output;
+
+          if (!hasMapping) {
+            const hintInput = isVlm
+              ? check.detected_image_column
+              : pickRoleColumn(check.suggested_mapping, "user");
+            const hintOutput = isVlm
+              ? check.detected_text_column
+              : pickRoleColumn(check.suggested_mapping, "assistant");
+
+            if (hintInput || hintOutput) {
+              useTrainingConfigStore.getState().setDatasetManualMapping({
+                input: hintInput ?? null,
+                output: hintOutput ?? null,
+              });
+            }
+
+            runtimeStore.setStarting(false);
+            dialogStore.openMapping(check);
+            return false;
+          }
+        }
+      }
+
       const payload = buildTrainingStartPayload(config);
       const response = await startTraining(payload);
 
@@ -78,3 +120,14 @@ export function useTrainingActions() {
     dismissTrainingRun,
   };
 }
+
+function pickRoleColumn(
+  mapping: Record<string, string> | null | undefined,
+  role: string,
+): string | null {
+  if (!mapping) return null;
+  for (const [col, mappedRole] of Object.entries(mapping)) {
+    if (mappedRole === role) return col;
+  }
+  return null;
+}
diff --git a/studio/frontend/src/features/training/index.ts b/studio/frontend/src/features/training/index.ts
index 41d1aade75..77451140f5 100644
--- a/studio/frontend/src/features/training/index.ts
+++ b/studio/frontend/src/features/training/index.ts
@@ -6,4 +6,5 @@ export {
 export { useTrainingActions } from "./hooks/use-training-actions";
 export { useTrainingRuntimeLifecycle } from "./hooks/use-training-runtime-lifecycle";
 export { HfDatasetSubsetSplitSelectors } from "./components/hf-dataset-subset-split-selectors";
+export { useDatasetPreviewDialogStore } from "./stores/dataset-preview-dialog-store";
 export type { TrainingPhase } from "./types/runtime";
diff --git a/studio/frontend/src/features/training/stores/dataset-preview-dialog-store.ts b/studio/frontend/src/features/training/stores/dataset-preview-dialog-store.ts
new file mode 100644
index 0000000000..9129e31be7
--- /dev/null
+++ b/studio/frontend/src/features/training/stores/dataset-preview-dialog-store.ts
@@ -0,0 +1,33 @@
+import { create } from "zustand";
+import type { CheckFormatResponse } from "../types/datasets";
+
+export type DatasetPreviewDialogMode = "preview" | "mapping";
+
+type DatasetPreviewDialogState = {
+  open: boolean;
+  mode: DatasetPreviewDialogMode;
+  initialData: CheckFormatResponse | null;
+};
+
+type DatasetPreviewDialogActions = {
+  openPreview: () => void;
+  openMapping: (data: CheckFormatResponse) => void;
+  close: () => void;
+};
+
+const initialState: DatasetPreviewDialogState = {
+  open: false,
+  mode: "preview",
+  initialData: null,
+};
+
+export const useDatasetPreviewDialogStore = create<
+  DatasetPreviewDialogState & DatasetPreviewDialogActions
+>()((set) => ({
+  ...initialState,
+
+  openPreview: () => set({ open: true, mode: "preview", initialData: null }),
+  openMapping: (data) => set({ open: true, mode: "mapping", initialData: data }),
+  close: () => set({ open: false, initialData: null, mode: "preview" }),
+}));
+
diff --git a/studio/frontend/src/features/training/stores/training-config-store.ts b/studio/frontend/src/features/training/stores/training-config-store.ts
index 57d05db114..fdb0deb129 100644
--- a/studio/frontend/src/features/training/stores/training-config-store.ts
+++ b/studio/frontend/src/features/training/stores/training-config-store.ts
@@ -18,6 +18,7 @@ const initialState: TrainingConfigState = {
   dataset: null,
   datasetSubset: null,
   datasetSplit: null,
+  datasetManualMapping: { input: null, output: null },
   uploadedFile: null,
   ...DEFAULT_HYPERPARAMS,
 };
@@ -58,10 +59,22 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
       setDatasetSource: (datasetSource) => set({ datasetSource }),
       setDatasetFormat: (datasetFormat) => set({ datasetFormat }),
       setDataset: (dataset) =>
-        set({ dataset, datasetSubset: null, datasetSplit: null }),
+        set({
+          dataset,
+          datasetSubset: null,
+          datasetSplit: null,
+          datasetManualMapping: { input: null, output: null },
+        }),
       setDatasetSubset: (datasetSubset) =>
-        set({ datasetSubset, datasetSplit: null }),
-      setDatasetSplit: (datasetSplit) => set({ datasetSplit }),
+        set({
+          datasetSubset,
+          datasetSplit: null,
+          datasetManualMapping: { input: null, output: null },
+        }),
+      setDatasetSplit: (datasetSplit) =>
+        set({ datasetSplit, datasetManualMapping: { input: null, output: null } }),
+      setDatasetManualMapping: (datasetManualMapping) =>
+        set({ datasetManualMapping }),
       setUploadedFile: (uploadedFile) => set({ uploadedFile }),
       setEpochs: (epochs) => set({ epochs }),
       setContextLength: (contextLength) => set({ contextLength }),
diff --git a/studio/frontend/src/features/training/types/api.ts b/studio/frontend/src/features/training/types/api.ts
index 6011ff6b96..f6c43616d9 100644
--- a/studio/frontend/src/features/training/types/api.ts
+++ b/studio/frontend/src/features/training/types/api.ts
@@ -9,6 +9,7 @@ export interface TrainingStartRequest {
   hf_dataset_split: string | null;
   local_datasets: string[];
   format_type: string;
+  custom_format_mapping?: Record<string, string> | null;
   num_epochs: number;
   learning_rate: string;
   batch_size: number;
diff --git a/studio/frontend/src/features/training/types/config.ts b/studio/frontend/src/features/training/types/config.ts
index 77706d60b2..6a28500b7e 100644
--- a/studio/frontend/src/features/training/types/config.ts
+++ b/studio/frontend/src/features/training/types/config.ts
@@ -9,6 +9,11 @@ import type {
 
 export type LoraVariant = "lora" | "rslora" | "loftq";
 
+export type DatasetManualMapping = {
+  input: string | null;
+  output: string | null;
+};
+
 export interface TrainingConfigState {
   currentStep: StepNumber;
   modelType: ModelType | null;
@@ -20,6 +25,7 @@ export interface TrainingConfigState {
   dataset: string | null;
   datasetSubset: string | null;
   datasetSplit: string | null;
+  datasetManualMapping: DatasetManualMapping;
   uploadedFile: string | null;
   epochs: number;
   contextLength: number;
@@ -64,6 +70,7 @@ export interface TrainingConfigActions {
   setDataset: (dataset: string | null) => void;
   setDatasetSubset: (subset: string | null) => void;
   setDatasetSplit: (split: string | null) => void;
+  setDatasetManualMapping: (mapping: DatasetManualMapping) => void;
   setUploadedFile: (file: string | null) => void;
   setEpochs: (epochs: number) => void;
   setContextLength: (length: number) => void;
diff --git a/studio/frontend/src/features/training/types/datasets.ts b/studio/frontend/src/features/training/types/datasets.ts
new file mode 100644
index 0000000000..96cf699f89
--- /dev/null
+++ b/studio/frontend/src/features/training/types/datasets.ts
@@ -0,0 +1,13 @@
+export type CheckFormatResponse = {
+  requires_manual_mapping: boolean;
+  detected_format: string;
+  columns: string[];
+  suggested_mapping?: Record<string, string> | null;
+  detected_image_column?: string | null;
+  detected_text_column?: string | null;
+  preview_samples?: Record<string, unknown>[] | null;
+  total_rows?: number | null;
+  is_multimodal?: boolean;
+  multimodal_columns?: string[] | null;
+};
+
diff --git a/studio/frontend/src/main.tsx b/studio/frontend/src/main.tsx
index 0719812d19..9bfc4451bb 100644
--- a/studio/frontend/src/main.tsx
+++ b/studio/frontend/src/main.tsx
@@ -1,14 +1,24 @@
 import { StrictMode } from "react";
 import { createRoot } from "react-dom/client";
 
-if (!crypto.randomUUID) {
-  crypto.randomUUID = () =>
+const globalCrypto = globalThis.crypto as Crypto | undefined;
+const hasUuid =
+  globalCrypto && typeof (globalCrypto as Crypto).randomUUID === "function";
+
+if (globalCrypto && !hasUuid) {
+  // Some envs ship `crypto` but no `randomUUID()` (or a non-function stub).
+  // Provide a best-effort v4 UUID using `getRandomValues` when available.
+  const getRandomByte = () => {
+    if (typeof globalCrypto.getRandomValues === "function") {
+      return globalCrypto.getRandomValues(new Uint8Array(1))[0];
+    }
+    return Math.floor(Math.random() * 256);
+  };
+
+  (globalCrypto as Crypto).randomUUID = (() =>
     "10000000-1000-4000-8000-100000000000".replace(/[018]/g, (c) =>
-      (
-        +c ^
-        (crypto.getRandomValues(new Uint8Array(1))[0] & (15 >> (+c / 4)))
-      ).toString(16),
-    ) as `${string}-${string}-${string}-${string}-${string}`;
+      (+c ^ (getRandomByte() & (15 >> (+c / 4)))).toString(16),
+    )) as Crypto["randomUUID"];
 }
 
 import "./index.css";

From cc6cb00d562ee4232a1593513c72c74b9f8ae0d0 Mon Sep 17 00:00:00 2001
From: Shine1i <wasimysdev@gmail.com>
Date: Mon, 16 Feb 2026 01:47:28 +0100
Subject: [PATCH 2/4] fix: enforce unique input/output mapping and enhance UI
 feedback in dataset preview dialog

---
 .../dataset-preview-dialog-mapping.tsx        | 47 +++++++++++---
 .../sections/dataset-preview-dialog.tsx       | 62 ++++++++++++-------
 2 files changed, 75 insertions(+), 34 deletions(-)

diff --git a/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx b/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx
index 5f8833b742..d7ba18ff35 100644
--- a/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx
+++ b/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx
@@ -1,7 +1,7 @@
 import { Badge } from "@/components/ui/badge";
 import { Button } from "@/components/ui/button";
 import { Checkbox } from "@/components/ui/checkbox";
-import { AlertCircleIcon } from "@hugeicons/core-free-icons";
+import { AlertCircleIcon, CheckmarkCircle02Icon } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
 import type { CheckFormatResponse } from "@/features/training/types/datasets";
 
@@ -40,21 +40,46 @@ export function DatasetMappingCard({
   input: string | null;
   output: string | null;
 }) {
+  const tone = mappingOk ? "ok" : "warn";
   return (
-    <div className="rounded-xl corner-squircle ring-1 ring-amber-200/70 bg-amber-50/70 px-5 py-4 mb-4 text-amber-950 dark:ring-amber-900/50 dark:bg-amber-950/30 dark:text-amber-50">
+    <div
+      className={
+        tone === "ok"
+          ? "rounded-xl corner-squircle ring-1 ring-emerald-200/70 bg-emerald-50/70 px-5 py-4 mb-4 text-emerald-950 dark:ring-emerald-900/50 dark:bg-emerald-950/30 dark:text-emerald-50"
+          : "rounded-xl corner-squircle ring-1 ring-amber-200/70 bg-amber-50/70 px-5 py-4 mb-4 text-amber-950 dark:ring-amber-900/50 dark:bg-amber-950/30 dark:text-amber-50"
+      }
+    >
       <div className="flex items-start gap-3">
-        <div className="rounded-xl corner-squircle bg-amber-500/15 p-2 shrink-0">
+        <div
+          className={
+            tone === "ok"
+              ? "rounded-xl corner-squircle bg-emerald-500/15 p-2 shrink-0"
+              : "rounded-xl corner-squircle bg-amber-500/15 p-2 shrink-0"
+          }
+        >
           <HugeiconsIcon
-            icon={AlertCircleIcon}
-            className="size-4 text-amber-700 dark:text-amber-300"
+            icon={mappingOk ? CheckmarkCircle02Icon : AlertCircleIcon}
+            className={
+              tone === "ok"
+                ? "size-4 text-emerald-700 dark:text-emerald-300"
+                : "size-4 text-amber-700 dark:text-amber-300"
+            }
           />
         </div>
         <div className="min-w-0">
-          <p className="text-sm font-semibold tracking-tight">Map dataset columns</p>
-          <p className="text-xs text-amber-800/80 dark:text-amber-200/80 mt-0.5">
-            We couldn&apos;t auto-detect the format. Pick the{" "}
-            {leftLabel.toLowerCase()} column and the {rightLabel.toLowerCase()}{" "}
-            column. We&apos;ll convert it to a supported format automatically.
+          <p className="text-sm font-semibold tracking-tight">
+            {mappingOk ? "Mapping ready" : "Map dataset columns"}
+          </p>
+          <p
+            className={
+              tone === "ok"
+                ? "text-xs text-emerald-800/80 dark:text-emerald-200/80 mt-0.5"
+                : "text-xs text-amber-800/80 dark:text-amber-200/80 mt-0.5"
+            }
+          >
+            {mappingOk
+              ? "Looks good. We'll convert this dataset automatically."
+              : `We couldn't auto-detect the format. Pick the ${leftLabel.toLowerCase()} column and the ${rightLabel.toLowerCase()} column. We'll convert it to a supported format automatically.`}
           </p>
           <div className="mt-3 flex flex-wrap items-center gap-2">
             <Badge
@@ -143,10 +168,12 @@ export function deriveDefaultMapping(
       data.detected_image_column ?? pickRole(data.suggested_mapping, "image");
     const output =
       data.detected_text_column ?? pickRole(data.suggested_mapping, "text");
+    if (input && output && input === output) return { input, output: null };
     return { input: input ?? null, output: output ?? null };
   }
   const input = pickRole(data.suggested_mapping, "user");
   const output = pickRole(data.suggested_mapping, "assistant");
+  if (input && output && input === output) return { input, output: null };
   return { input: input ?? null, output: output ?? null };
 }
 
diff --git a/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx b/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx
index 11de320147..ce4528341a 100644
--- a/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx
+++ b/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx
@@ -67,6 +67,12 @@ export function DatasetPreviewDialog({
   const leftLabel = isVlm ? "Image" : "Input";
   const rightLabel = isVlm ? "Text" : "Output";
 
+  useEffect(() => {
+    if (!manualMapping.input || !manualMapping.output) return;
+    if (manualMapping.input !== manualMapping.output) return;
+    setManualMapping({ input: manualMapping.input, output: null });
+  }, [manualMapping.input, manualMapping.output, setManualMapping]);
+
   useEffect(() => {
     if (!open || !datasetName) {
       setData(null);
@@ -153,30 +159,38 @@ export function DatasetPreviewDialog({
           </span>
           {mappingEnabled && (
             <div className="flex items-center gap-3">
-              {(manualMapping.input == null || manualMapping.input === colName) && (
-                <HeaderPick
-                  label={leftLabel}
-                  checked={manualMapping.input === colName}
-                  onCheckedChange={(checked) => {
-                    setManualMapping({
-                      input: checked ? colName : null,
-                      output: manualMapping.output,
-                    });
-                  }}
-                />
-              )}
-              {(manualMapping.output == null || manualMapping.output === colName) && (
-                <HeaderPick
-                  label={rightLabel}
-                  checked={manualMapping.output === colName}
-                  onCheckedChange={(checked) => {
-                    setManualMapping({
-                      input: manualMapping.input,
-                      output: checked ? colName : null,
-                    });
-                  }}
-                />
-              )}
+              {(manualMapping.output == null || manualMapping.output !== colName) &&
+                (manualMapping.input == null || manualMapping.input === colName) && (
+                  <HeaderPick
+                    label={leftLabel}
+                    checked={manualMapping.input === colName}
+                    onCheckedChange={(checked) => {
+                      setManualMapping({
+                        input: checked ? colName : null,
+                        output:
+                          checked && manualMapping.output === colName
+                            ? null
+                            : manualMapping.output,
+                      });
+                    }}
+                  />
+                )}
+              {(manualMapping.input == null || manualMapping.input !== colName) &&
+                (manualMapping.output == null || manualMapping.output === colName) && (
+                  <HeaderPick
+                    label={rightLabel}
+                    checked={manualMapping.output === colName}
+                    onCheckedChange={(checked) => {
+                      setManualMapping({
+                        input:
+                          checked && manualMapping.input === colName
+                            ? null
+                            : manualMapping.input,
+                        output: checked ? colName : null,
+                      });
+                    }}
+                  />
+                )}
             </div>
           )}
         </div>

From 1a90741046c8f8e7e21b84f58f9043b3ac4ef952 Mon Sep 17 00:00:00 2001
From: Shine1i <wasimysdev@gmail.com>
Date: Mon, 16 Feb 2026 01:54:09 +0100
Subject: [PATCH 3/4] refactor: create utilities for dataset manual mapping and
 improve UI logic consistency in dataset preview dialog

---
 .../dataset-preview-dialog-mapping.tsx        |  66 +++++-----
 .../sections/dataset-preview-dialog.tsx       | 122 +++++++++---------
 .../studio/sections/dataset-section.tsx       |  29 +++--
 .../src/features/training/api/mappers.ts      |  23 ++--
 .../training/hooks/use-training-actions.ts    |  63 +++++----
 .../training/stores/training-config-store.ts  |  12 +-
 studio/frontend/src/main.tsx                  |  24 ++--
 7 files changed, 177 insertions(+), 162 deletions(-)

diff --git a/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx b/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx
index d7ba18ff35..e2dfb146f7 100644
--- a/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx
+++ b/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx
@@ -1,9 +1,10 @@
 import { Badge } from "@/components/ui/badge";
 import { Button } from "@/components/ui/button";
 import { Checkbox } from "@/components/ui/checkbox";
+import type { CheckFormatResponse } from "@/features/training/types/datasets";
+import { cn } from "@/lib/utils";
 import { AlertCircleIcon, CheckmarkCircle02Icon } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
-import type { CheckFormatResponse } from "@/features/training/types/datasets";
 
 export function HeaderPick({
   label,
@@ -40,30 +41,30 @@ export function DatasetMappingCard({
   input: string | null;
   output: string | null;
 }) {
-  const tone = mappingOk ? "ok" : "warn";
   return (
     <div
-      className={
-        tone === "ok"
-          ? "rounded-xl corner-squircle ring-1 ring-emerald-200/70 bg-emerald-50/70 px-5 py-4 mb-4 text-emerald-950 dark:ring-emerald-900/50 dark:bg-emerald-950/30 dark:text-emerald-50"
-          : "rounded-xl corner-squircle ring-1 ring-amber-200/70 bg-amber-50/70 px-5 py-4 mb-4 text-amber-950 dark:ring-amber-900/50 dark:bg-amber-950/30 dark:text-amber-50"
-      }
+      className={cn(
+        "rounded-xl corner-squircle ring-1 px-5 py-4 mb-4",
+        mappingOk
+          ? "ring-emerald-200/70 bg-emerald-50/70 text-emerald-950 dark:ring-emerald-900/50 dark:bg-emerald-950/30 dark:text-emerald-50"
+          : "ring-amber-200/70 bg-amber-50/70 text-amber-950 dark:ring-amber-900/50 dark:bg-amber-950/30 dark:text-amber-50",
+      )}
     >
       <div className="flex items-start gap-3">
         <div
-          className={
-            tone === "ok"
-              ? "rounded-xl corner-squircle bg-emerald-500/15 p-2 shrink-0"
-              : "rounded-xl corner-squircle bg-amber-500/15 p-2 shrink-0"
-          }
+          className={cn(
+            "rounded-xl corner-squircle p-2 shrink-0",
+            mappingOk ? "bg-emerald-500/15" : "bg-amber-500/15",
+          )}
         >
           <HugeiconsIcon
             icon={mappingOk ? CheckmarkCircle02Icon : AlertCircleIcon}
-            className={
-              tone === "ok"
-                ? "size-4 text-emerald-700 dark:text-emerald-300"
-                : "size-4 text-amber-700 dark:text-amber-300"
-            }
+            className={cn(
+              "size-4",
+              mappingOk
+                ? "text-emerald-700 dark:text-emerald-300"
+                : "text-amber-700 dark:text-amber-300",
+            )}
           />
         </div>
         <div className="min-w-0">
@@ -71,11 +72,12 @@ export function DatasetMappingCard({
             {mappingOk ? "Mapping ready" : "Map dataset columns"}
           </p>
           <p
-            className={
-              tone === "ok"
-                ? "text-xs text-emerald-800/80 dark:text-emerald-200/80 mt-0.5"
-                : "text-xs text-amber-800/80 dark:text-amber-200/80 mt-0.5"
-            }
+            className={cn(
+              "text-xs mt-0.5",
+              mappingOk
+                ? "text-emerald-800/80 dark:text-emerald-200/80"
+                : "text-amber-800/80 dark:text-amber-200/80",
+            )}
           >
             {mappingOk
               ? "Looks good. We'll convert this dataset automatically."
@@ -163,17 +165,17 @@ export function deriveDefaultMapping(
   data: CheckFormatResponse,
   isVlm: boolean,
 ): { input: string | null; output: string | null } {
-  if (isVlm) {
-    const input =
-      data.detected_image_column ?? pickRole(data.suggested_mapping, "image");
-    const output =
-      data.detected_text_column ?? pickRole(data.suggested_mapping, "text");
-    if (input && output && input === output) return { input, output: null };
-    return { input: input ?? null, output: output ?? null };
+  const input = isVlm
+    ? data.detected_image_column ?? pickRole(data.suggested_mapping, "image")
+    : pickRole(data.suggested_mapping, "user");
+  const output = isVlm
+    ? data.detected_text_column ?? pickRole(data.suggested_mapping, "text")
+    : pickRole(data.suggested_mapping, "assistant");
+
+  if (input && output && input === output) {
+    return { input, output: null };
   }
-  const input = pickRole(data.suggested_mapping, "user");
-  const output = pickRole(data.suggested_mapping, "assistant");
-  if (input && output && input === output) return { input, output: null };
+
   return { input: input ?? null, output: output ?? null };
 }
 
diff --git a/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx b/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx
index ce4528341a..9fc84f1013 100644
--- a/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx
+++ b/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx
@@ -1,3 +1,4 @@
+import { type ReactNode, useEffect, useMemo, useState } from "react";
 import type { ColumnDef } from "@tanstack/react-table";
 import {
   Dialog,
@@ -8,12 +9,12 @@ import {
 import { DataTable } from "@/components/ui/data-table";
 import { Badge } from "@/components/ui/badge";
 import { Spinner } from "@/components/ui/spinner";
-import { Database02Icon, AlertCircleIcon } from "@hugeicons/core-free-icons";
-import { HugeiconsIcon } from "@hugeicons/react";
-import { type ReactNode, useEffect, useMemo, useState } from "react";
 import { useTrainingActions, useTrainingConfigStore } from "@/features/training";
 import { checkDatasetFormat } from "@/features/training/api/datasets-api";
 import type { CheckFormatResponse } from "@/features/training/types/datasets";
+import { Database02Icon, AlertCircleIcon } from "@hugeicons/core-free-icons";
+import { HugeiconsIcon } from "@hugeicons/react";
+import { useShallow } from "zustand/react/shallow";
 import { collectPreviewImages, formatCell } from "./dataset-preview-dialog-utils";
 import {
   DatasetMappingCard,
@@ -22,10 +23,6 @@ import {
   deriveDefaultMapping,
 } from "./dataset-preview-dialog-mapping";
 
-// ---------------------------------------------------------------------------
-// Types
-// ---------------------------------------------------------------------------
-
 type DatasetPreviewDialogProps = {
   open: boolean;
   onOpenChange: (open: boolean) => void;
@@ -38,10 +35,6 @@ type DatasetPreviewDialogProps = {
   isVlm?: boolean;
 };
 
-// ---------------------------------------------------------------------------
-// Component
-// ---------------------------------------------------------------------------
-
 export function DatasetPreviewDialog({
   open,
   onOpenChange,
@@ -57,8 +50,12 @@ export function DatasetPreviewDialog({
   const [loading, setLoading] = useState(false);
   const [error, setError] = useState<string | null>(null);
 
-  const manualMapping = useTrainingConfigStore((s) => s.datasetManualMapping);
-  const setManualMapping = useTrainingConfigStore((s) => s.setDatasetManualMapping);
+  const { manualMapping, setManualMapping } = useTrainingConfigStore(
+    useShallow((s) => ({
+      manualMapping: s.datasetManualMapping,
+      setManualMapping: s.setDatasetManualMapping,
+    })),
+  );
   const { isStarting, startError, startTrainingRun } = useTrainingActions();
 
   const mappingEnabled = !!data?.requires_manual_mapping;
@@ -79,12 +76,14 @@ export function DatasetPreviewDialog({
       setError(null);
       return;
     }
+
     if (initialData) {
       setData(initialData);
       setError(null);
       setLoading(false);
       return;
     }
+
     let cancelled = false;
     setLoading(true);
     setError(null);
@@ -115,21 +114,13 @@ export function DatasetPreviewDialog({
   }, [open, datasetName, hfToken, datasetSubset, datasetSplit, isVlm, initialData]);
 
   useEffect(() => {
-    if (!open || !datasetName || !data?.requires_manual_mapping) return;
+    if (!open || !datasetName) return;
+    if (!data?.requires_manual_mapping) return;
     if (manualMapping.input || manualMapping.output) return;
     const derived = deriveDefaultMapping(data, isVlm);
     if (!derived.input && !derived.output) return;
     setManualMapping(derived);
-  }, [
-    open,
-    datasetName,
-    data?.requires_manual_mapping,
-    isVlm,
-    manualMapping.input,
-    manualMapping.output,
-    setManualMapping,
-    data,
-  ]);
+  }, [open, datasetName, data, isVlm, manualMapping.input, manualMapping.output, setManualMapping]);
 
   const rows = data?.preview_samples ?? [];
   const columns = data?.columns ?? [];
@@ -159,38 +150,30 @@ export function DatasetPreviewDialog({
           </span>
           {mappingEnabled && (
             <div className="flex items-center gap-3">
-              {(manualMapping.output == null || manualMapping.output !== colName) &&
-                (manualMapping.input == null || manualMapping.input === colName) && (
-                  <HeaderPick
-                    label={leftLabel}
-                    checked={manualMapping.input === colName}
-                    onCheckedChange={(checked) => {
-                      setManualMapping({
-                        input: checked ? colName : null,
-                        output:
-                          checked && manualMapping.output === colName
-                            ? null
-                            : manualMapping.output,
-                      });
-                    }}
-                  />
-                )}
-              {(manualMapping.input == null || manualMapping.input !== colName) &&
-                (manualMapping.output == null || manualMapping.output === colName) && (
-                  <HeaderPick
-                    label={rightLabel}
-                    checked={manualMapping.output === colName}
-                    onCheckedChange={(checked) => {
-                      setManualMapping({
-                        input:
-                          checked && manualMapping.input === colName
-                            ? null
-                            : manualMapping.input,
-                        output: checked ? colName : null,
-                      });
-                    }}
-                  />
-                )}
+              {canShowInputPicker(colName, manualMapping) && (
+                <HeaderPick
+                  label={leftLabel}
+                  checked={manualMapping.input === colName}
+                  onCheckedChange={(checked) => {
+                    setManualMapping({
+                      input: checked ? colName : null,
+                      output: manualMapping.output,
+                    });
+                  }}
+                />
+              )}
+              {canShowOutputPicker(colName, manualMapping) && (
+                <HeaderPick
+                  label={rightLabel}
+                  checked={manualMapping.output === colName}
+                  onCheckedChange={(checked) => {
+                    setManualMapping({
+                      input: manualMapping.input,
+                      output: checked ? colName : null,
+                    });
+                  }}
+                />
+              )}
             </div>
           )}
         </div>
@@ -235,8 +218,7 @@ export function DatasetPreviewDialog({
             </span>
           );
         }
-        const full =
-          typeof value === "string" ? value : JSON.stringify(value);
+        const full = typeof value === "string" ? value : JSON.stringify(value);
         return (
           <p
             className="text-[13px] leading-relaxed line-clamp-6"
@@ -325,10 +307,10 @@ export function DatasetPreviewDialog({
                       : "--"
                   }
                 />
-              <MetaRow
-                label="Columns"
-                value={
-                  <span className="flex items-center gap-1.5 flex-wrap">
+                <MetaRow
+                  label="Columns"
+                  value={
+                    <span className="flex items-center gap-1.5 flex-wrap">
                       {columns.map((col) => (
                         <Badge
                           key={col}
@@ -417,4 +399,18 @@ function MetaRow({
   );
 }
 
-// mapping UI extracted to ./dataset-preview-dialog-mapping.tsx
+function canShowInputPicker(
+  colName: string,
+  mapping: { input: string | null; output: string | null },
+): boolean {
+  if (mapping.output === colName) return false;
+  return mapping.input == null || mapping.input === colName;
+}
+
+function canShowOutputPicker(
+  colName: string,
+  mapping: { input: string | null; output: string | null },
+): boolean {
+  if (mapping.input === colName) return false;
+  return mapping.output == null || mapping.output === colName;
+}
diff --git a/studio/frontend/src/features/studio/sections/dataset-section.tsx b/studio/frontend/src/features/studio/sections/dataset-section.tsx
index 3d7e3f810f..055b4db9cf 100644
--- a/studio/frontend/src/features/studio/sections/dataset-section.tsx
+++ b/studio/frontend/src/features/studio/sections/dataset-section.tsx
@@ -190,20 +190,21 @@ export function DatasetSection() {
                   ref={scrollRef}
                   className="max-h-64 overflow-y-auto overscroll-contain [scrollbar-width:thin]"
                 >
-                  <ComboboxList className="p-1 !max-h-none !overflow-visible">
-                    {(id: string) => {
-                      const r = hfResults.find((ds) => ds.id === id);
-                      const detail = r?.totalExamples
-                        ? `${formatCompact(r.totalExamples)} rows`
-                        : r?.sizeCategory
-                          ? r.sizeCategory
-                          : r?.downloads != null
-                            ? `↓${formatCompact(r.downloads)}`
-                            : null;
-                      return (
-                        <ComboboxItem
-                          key={id}
-                          value={id}
+                    <ComboboxList className="p-1 !max-h-none !overflow-visible">
+                      {(id: string) => {
+                        const r = hfResults.find((ds) => ds.id === id);
+                        let detail: string | null = null;
+                        if (r?.totalExamples) {
+                          detail = `${formatCompact(r.totalExamples)} rows`;
+                        } else if (r?.sizeCategory) {
+                          detail = r.sizeCategory;
+                        } else if (r?.downloads != null) {
+                          detail = `↓${formatCompact(r.downloads)}`;
+                        }
+                        return (
+                          <ComboboxItem
+                            key={id}
+                            value={id}
                           className="justify-between"
                         >
                           <Tooltip>
diff --git a/studio/frontend/src/features/training/api/mappers.ts b/studio/frontend/src/features/training/api/mappers.ts
index ed2c3b752d..05fdd98959 100644
--- a/studio/frontend/src/features/training/api/mappers.ts
+++ b/studio/frontend/src/features/training/api/mappers.ts
@@ -14,15 +14,7 @@ export function buildTrainingStartPayload(
   const adapterMethod = config.trainingMethod !== "full";
   const isQlorMethod = config.trainingMethod === "qlora";
   const hfDataset = config.datasetSource === "huggingface" ? config.dataset : null;
-  const manual = config.datasetManualMapping;
-  const isVlm = config.modelType === "vision";
-  const customFormatMapping =
-    manual.input && manual.output
-      ? {
-          [manual.input]: isVlm ? "image" : "user",
-          [manual.output]: isVlm ? "text" : "assistant",
-        }
-      : undefined;
+  const customFormatMapping = buildCustomFormatMapping(config);
 
   return {
     model_name: config.selectedModel ?? "",
@@ -73,3 +65,16 @@ export function buildTrainingStartPayload(
       : null,
   };
 }
+
+function buildCustomFormatMapping(
+  config: TrainingConfigState,
+): Record<string, string> | undefined {
+  const { input, output } = config.datasetManualMapping;
+  if (!input || !output) return undefined;
+
+  if (config.modelType === "vision") {
+    return { [input]: "image", [output]: "text" };
+  }
+
+  return { [input]: "user", [output]: "assistant" };
+}
diff --git a/studio/frontend/src/features/training/hooks/use-training-actions.ts b/studio/frontend/src/features/training/hooks/use-training-actions.ts
index 27e36f9765..e3600432e9 100644
--- a/studio/frontend/src/features/training/hooks/use-training-actions.ts
+++ b/studio/frontend/src/features/training/hooks/use-training-actions.ts
@@ -1,12 +1,13 @@
 import { useCallback } from "react";
-import { useTrainingConfigStore } from "../stores/training-config-store";
-import { useTrainingRuntimeStore } from "../stores/training-runtime-store";
-import { useDatasetPreviewDialogStore } from "../stores/dataset-preview-dialog-store";
-import { startTraining, stopTraining, resetTraining } from "../api/train-api";
-import { buildTrainingStartPayload } from "../api/mappers";
 import { checkDatasetFormat } from "../api/datasets-api";
+import { buildTrainingStartPayload } from "../api/mappers";
+import { startTraining, stopTraining, resetTraining } from "../api/train-api";
 import { syncTrainingRuntimeFromBackend } from "../lib/sync-runtime";
 import { validateTrainingConfig } from "../lib/validation";
+import { useDatasetPreviewDialogStore } from "../stores/dataset-preview-dialog-store";
+import { useTrainingConfigStore } from "../stores/training-config-store";
+import { useTrainingRuntimeStore } from "../stores/training-runtime-store";
+import type { TrainingConfigState } from "../types/config";
 
 export function useTrainingActions() {
   const isStarting = useTrainingRuntimeStore((state) => state.isStarting);
@@ -27,8 +28,7 @@ export function useTrainingActions() {
     runtimeStore.setStarting(true);
 
     try {
-      const datasetName =
-        config.datasetSource === "huggingface" ? config.dataset : config.uploadedFile;
+      const datasetName = getDatasetName(config);
       const isVlm = config.modelType === "vision";
 
       if (datasetName) {
@@ -40,29 +40,24 @@ export function useTrainingActions() {
           isVlm,
         });
 
-        if (check.requires_manual_mapping) {
-          const existing = useTrainingConfigStore.getState().datasetManualMapping;
-          const hasMapping = !!existing.input && !!existing.output;
+        if (check.requires_manual_mapping && !hasManualMapping(config)) {
+          const hintInput = isVlm
+            ? check.detected_image_column
+            : pickRoleColumn(check.suggested_mapping, "user");
+          const hintOutput = isVlm
+            ? check.detected_text_column
+            : pickRoleColumn(check.suggested_mapping, "assistant");
 
-          if (!hasMapping) {
-            const hintInput = isVlm
-              ? check.detected_image_column
-              : pickRoleColumn(check.suggested_mapping, "user");
-            const hintOutput = isVlm
-              ? check.detected_text_column
-              : pickRoleColumn(check.suggested_mapping, "assistant");
-
-            if (hintInput || hintOutput) {
-              useTrainingConfigStore.getState().setDatasetManualMapping({
-                input: hintInput ?? null,
-                output: hintOutput ?? null,
-              });
-            }
-
-            runtimeStore.setStarting(false);
-            dialogStore.openMapping(check);
-            return false;
+          if (hintInput || hintOutput) {
+            useTrainingConfigStore.getState().setDatasetManualMapping({
+              input: hintInput ?? null,
+              output: hintOutput ?? null,
+            });
           }
+
+          runtimeStore.setStarting(false);
+          dialogStore.openMapping(check);
+          return false;
         }
       }
 
@@ -121,6 +116,18 @@ export function useTrainingActions() {
   };
 }
 
+function getDatasetName(config: TrainingConfigState): string | null {
+  return config.datasetSource === "huggingface"
+    ? config.dataset
+    : config.uploadedFile;
+}
+
+function hasManualMapping(config: TrainingConfigState): boolean {
+  return (
+    !!config.datasetManualMapping.input && !!config.datasetManualMapping.output
+  );
+}
+
 function pickRoleColumn(
   mapping: Record<string, string> | null | undefined,
   role: string,
diff --git a/studio/frontend/src/features/training/stores/training-config-store.ts b/studio/frontend/src/features/training/stores/training-config-store.ts
index fdb0deb129..731ded9d80 100644
--- a/studio/frontend/src/features/training/stores/training-config-store.ts
+++ b/studio/frontend/src/features/training/stores/training-config-store.ts
@@ -7,6 +7,10 @@ import type { TrainingConfigState, TrainingConfigStore } from "../types/config";
 const MIN_STEP: StepNumber = 1;
 const MAX_STEP: StepNumber = STEPS.length as StepNumber;
 
+function emptyManualMapping(): TrainingConfigState["datasetManualMapping"] {
+  return { input: null, output: null };
+}
+
 const initialState: TrainingConfigState = {
   currentStep: MIN_STEP,
   modelType: null,
@@ -18,7 +22,7 @@ const initialState: TrainingConfigState = {
   dataset: null,
   datasetSubset: null,
   datasetSplit: null,
-  datasetManualMapping: { input: null, output: null },
+  datasetManualMapping: emptyManualMapping(),
   uploadedFile: null,
   ...DEFAULT_HYPERPARAMS,
 };
@@ -63,16 +67,16 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
           dataset,
           datasetSubset: null,
           datasetSplit: null,
-          datasetManualMapping: { input: null, output: null },
+          datasetManualMapping: emptyManualMapping(),
         }),
       setDatasetSubset: (datasetSubset) =>
         set({
           datasetSubset,
           datasetSplit: null,
-          datasetManualMapping: { input: null, output: null },
+          datasetManualMapping: emptyManualMapping(),
         }),
       setDatasetSplit: (datasetSplit) =>
-        set({ datasetSplit, datasetManualMapping: { input: null, output: null } }),
+        set({ datasetSplit, datasetManualMapping: emptyManualMapping() }),
       setDatasetManualMapping: (datasetManualMapping) =>
         set({ datasetManualMapping }),
       setUploadedFile: (uploadedFile) => set({ uploadedFile }),
diff --git a/studio/frontend/src/main.tsx b/studio/frontend/src/main.tsx
index 9bfc4451bb..4358241c94 100644
--- a/studio/frontend/src/main.tsx
+++ b/studio/frontend/src/main.tsx
@@ -1,29 +1,29 @@
 import { StrictMode } from "react";
 import { createRoot } from "react-dom/client";
 
-const globalCrypto = globalThis.crypto as Crypto | undefined;
-const hasUuid =
-  globalCrypto && typeof (globalCrypto as Crypto).randomUUID === "function";
+import "./index.css";
+import { App } from "./app/app";
 
-if (globalCrypto && !hasUuid) {
+const globalCrypto = globalThis.crypto as Crypto | undefined;
+
+if (globalCrypto && typeof globalCrypto.randomUUID !== "function") {
   // Some envs ship `crypto` but no `randomUUID()` (or a non-function stub).
   // Provide a best-effort v4 UUID using `getRandomValues` when available.
-  const getRandomByte = () => {
-    if (typeof globalCrypto.getRandomValues === "function") {
-      return globalCrypto.getRandomValues(new Uint8Array(1))[0];
+  const cryptoRef = globalCrypto;
+
+  function getRandomByte(): number {
+    if (typeof cryptoRef.getRandomValues === "function") {
+      return cryptoRef.getRandomValues(new Uint8Array(1))[0];
     }
     return Math.floor(Math.random() * 256);
-  };
+  }
 
-  (globalCrypto as Crypto).randomUUID = (() =>
+  cryptoRef.randomUUID = (() =>
     "10000000-1000-4000-8000-100000000000".replace(/[018]/g, (c) =>
       (+c ^ (getRandomByte() & (15 >> (+c / 4)))).toString(16),
     )) as Crypto["randomUUID"];
 }
 
-import "./index.css";
-import { App } from "./app/app";
-
 const rootElement = document.getElementById("root");
 if (!rootElement) {
   throw new Error("Root element not found");

From d0bcf02784dd8e174cfdc809ece849de6fc51038 Mon Sep 17 00:00:00 2001
From: Shine1i <wasimysdev@gmail.com>
Date: Mon, 16 Feb 2026 02:05:53 +0100
Subject: [PATCH 4/4] feat: lock model selector during guided tour and refine
 open state logic

---
 .../frontend/src/features/chat/chat-page.tsx  | 29 +++++++++++++++++--
 .../src/features/studio/studio-page.tsx       |  3 +-
 2 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx
index 0f57841249..c79f4b387c 100644
--- a/studio/frontend/src/features/chat/chat-page.tsx
+++ b/studio/frontend/src/features/chat/chat-page.tsx
@@ -205,6 +205,7 @@ export function ChatPage(): ReactElement {
   });
   const [settingsOpen, setSettingsOpen] = useState(false);
   const [modelSelectorOpen, setModelSelectorOpen] = useState(false);
+  const [modelSelectorLocked, setModelSelectorLocked] = useState(false);
   const [sidebarOpen, setSidebarOpen] = useState(true);
   const viewBeforeCompareRef = useRef<ChatView | null>(null);
   const inferenceParams = useChatRuntimeStore((state) => state.params);
@@ -239,8 +240,23 @@ export function ChatPage(): ReactElement {
     [],
   );
 
-  const openModelSelector = useCallback(() => setModelSelectorOpen(true), []);
-  const closeModelSelector = useCallback(() => setModelSelectorOpen(false), []);
+  const openModelSelector = useCallback(() => {
+    setModelSelectorLocked(true);
+    setModelSelectorOpen(true);
+  }, []);
+
+  const closeModelSelector = useCallback(() => {
+    setModelSelectorLocked(false);
+    setModelSelectorOpen(false);
+  }, []);
+
+  const handleModelSelectorOpenChange = useCallback(
+    (open: boolean) => {
+      if (!open && modelSelectorLocked) return;
+      setModelSelectorOpen(open);
+    },
+    [modelSelectorLocked],
+  );
   const openSettings = useCallback(() => setSettingsOpen(true), []);
   const closeSettings = useCallback(() => setSettingsOpen(false), []);
   const openSidebar = useCallback(() => setSidebarOpen(true), []);
@@ -313,6 +329,13 @@ export function ChatPage(): ReactElement {
     steps: tourSteps,
   });
 
+  useEffect(() => {
+    if (tour.open) return;
+    if (!modelSelectorLocked) return;
+    setModelSelectorLocked(false);
+    setModelSelectorOpen(false);
+  }, [modelSelectorLocked, tour.open]);
+
   return (
     <div className="h-[calc(100vh-4rem)] bg-background overflow-hidden">
     <GuidedTour {...tour.tourProps} />
@@ -355,7 +378,7 @@ export function ChatPage(): ReactElement {
               onEject={handleEject}
               variant="ghost"
               open={modelSelectorOpen}
-              onOpenChange={setModelSelectorOpen}
+              onOpenChange={handleModelSelectorOpenChange}
               triggerDataTour="chat-model-selector"
               contentDataTour="chat-model-selector-popover"
             />
diff --git a/studio/frontend/src/features/studio/studio-page.tsx b/studio/frontend/src/features/studio/studio-page.tsx
index dd2c075a67..8b92636fde 100644
--- a/studio/frontend/src/features/studio/studio-page.tsx
+++ b/studio/frontend/src/features/studio/studio-page.tsx
@@ -55,7 +55,7 @@ export function StudioPage(): ReactElement {
   return (
     <div className="min-h-screen bg-background">
       <main className="mx-auto max-w-7xl px-6 py-4">
-        <GuidedTour {...tour.tourProps} celebrate={true} />
+        <GuidedTour {...tour.tourProps} celebrate={isConfigTour} />
 
         <DatasetPreviewDialog
           open={dialogOpen}
@@ -114,4 +114,3 @@ export function StudioPage(): ReactElement {
     </div>
   );
 }
-