diff --git a/studio/frontend/src/features/studio/sections/progress-section.tsx b/studio/frontend/src/features/studio/sections/progress-section.tsx index a04273f568..10a547b82d 100644 --- a/studio/frontend/src/features/studio/sections/progress-section.tsx +++ b/studio/frontend/src/features/studio/sections/progress-section.tsx @@ -237,6 +237,7 @@ export function ProgressSection(): ReactElement { onClick={() => { setStopRequested(true); setStopDialogOpen(false); + useTrainingRuntimeStore.getState().setStopRequested(true); void stopTrainingRun(false).then((ok) => { if (!ok) setStopRequested(false); }); @@ -248,6 +249,7 @@ export function ProgressSection(): ReactElement { onClick={() => { setStopRequested(true); setStopDialogOpen(false); + useTrainingRuntimeStore.getState().setStopRequested(true); void stopTrainingRun(true).then((ok) => { if (!ok) setStopRequested(false); }); diff --git a/studio/frontend/src/features/studio/studio-page.tsx b/studio/frontend/src/features/studio/studio-page.tsx index 26353b9b7e..69abc1969e 100644 --- a/studio/frontend/src/features/studio/studio-page.tsx +++ b/studio/frontend/src/features/studio/studio-page.tsx @@ -45,11 +45,16 @@ export function StudioPage(): ReactElement { const dialogInitial = useDatasetPreviewDialogStore((s) => s.initialData); const closeDialog = useDatasetPreviewDialogStore((s) => s.close); + const stopRequested = useTrainingRuntimeStore((state) => state.stopRequested); const canGoBack = showTrainingView && - !isTrainingRunning && !isHydratingRuntime && - (runtimePhase === "stopped" || runtimePhase === "error" || runtimePhase === "completed" || runtimePhase === "idle"); + (stopRequested || + (!isTrainingRunning && + (runtimePhase === "stopped" || + runtimePhase === "error" || + runtimePhase === "completed" || + runtimePhase === "idle"))); const tourEnabled = hasHydratedRuntime && !isHydratingRuntime; const isConfigTour = !showTrainingView; const tourSteps = showTrainingView ? studioTrainingTourSteps : studioTourSteps; diff --git a/studio/frontend/src/features/studio/training-start-overlay.tsx b/studio/frontend/src/features/studio/training-start-overlay.tsx index b6c413b400..03078302f6 100644 --- a/studio/frontend/src/features/studio/training-start-overlay.tsx +++ b/studio/frontend/src/features/studio/training-start-overlay.tsx @@ -1,9 +1,23 @@ +import { + AlertDialog, + AlertDialogAction, + AlertDialogCancel, + AlertDialogContent, + AlertDialogDescription, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogTitle, +} from "@/components/ui/alert-dialog"; +import { Button } from "@/components/ui/button"; import { AnimatedSpan, Terminal, TypingAnimation, -} from "@/components/ui/terminal" -import type { ReactElement } from "react" +} from "@/components/ui/terminal"; +import { useTrainingActions, useTrainingRuntimeStore } from "@/features/training"; +import { Cancel01Icon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; +import { useEffect, useState, type ReactElement } from "react"; type TrainingStartOverlayProps = { message: string @@ -14,18 +28,65 @@ export function TrainingStartOverlay({ message, currentStep, }: TrainingStartOverlayProps): ReactElement { + const { stopTrainingRun } = useTrainingActions(); + const isStarting = useTrainingRuntimeStore((s) => s.isStarting); + const [cancelDialogOpen, setCancelDialogOpen] = useState(false); + const [cancelRequested, setCancelRequested] = useState(false); + + useEffect(() => { + if (!isStarting) { + setCancelRequested(false); + } + }, [isStarting]); + return (
-