From 8a433c6e1fdd178a8ee92df9d4e8a4e2f580af75 Mon Sep 17 00:00:00 2001 From: Hannes Rudolph Date: Thu, 29 Jan 2026 12:28:47 -0700 Subject: [PATCH 1/3] feat: add initial checkpoint as state-driven UI element Show the initial checkpoint (workspace state at task start) as a fixed UI element at the top of the chat, with visual states for pending/ready/failed. This replaces the previous approach of a header button and provides better UX: - Pending state: greyed out with spinner while initializing - Ready state: full color with restore/diff menu options - Failed state: warning styling with error indication Changes: - Add initialCheckpointState and initialCheckpointHash to extension state - New InitialCheckpoint component with three visual states - CheckpointMenu now supports isInitial prop to hide irrelevant options - Backend emits state transitions instead of messages --- packages/types/src/vscode-extension-host.ts | 8 + src/core/checkpoints/index.ts | 27 ++- webview-ui/src/components/chat/ChatView.tsx | 11 ++ .../chat/checkpoints/CheckpointMenu.tsx | 37 ++-- .../chat/checkpoints/CheckpointSaved.tsx | 5 +- .../chat/checkpoints/InitialCheckpoint.tsx | 131 ++++++++++++++ .../__tests__/CheckpointSaved.spec.tsx | 144 ++++++++++++++++ .../__tests__/InitialCheckpoint.spec.tsx | 161 ++++++++++++++++++ .../src/components/chat/checkpoints/schema.ts | 1 + .../src/context/ExtensionStateContext.tsx | 11 ++ webview-ui/src/i18n/locales/en/chat.json | 4 + 11 files changed, 523 insertions(+), 17 deletions(-) create mode 100644 webview-ui/src/components/chat/checkpoints/InitialCheckpoint.tsx create mode 100644 webview-ui/src/components/chat/checkpoints/__tests__/InitialCheckpoint.spec.tsx diff --git a/packages/types/src/vscode-extension-host.ts b/packages/types/src/vscode-extension-host.ts index d8c421f4b8d..77c8597a285 100644 --- a/packages/types/src/vscode-extension-host.ts +++ b/packages/types/src/vscode-extension-host.ts @@ -60,6 +60,7 @@ export interface ExtensionMessage { | "deleteCustomModeCheck" | "currentCheckpointUpdated" | "checkpointInitWarning" + | "initialCheckpointState" | "browserToolEnabled" | "browserConnectionResult" | "remoteBrowserEnabled" @@ -116,6 +117,9 @@ export interface ExtensionMessage { type: "WAIT_TIMEOUT" | "INIT_TIMEOUT" timeout: number } + // Initial checkpoint state for state-driven UI + initialCheckpointState?: "pending" | "ready" | "failed" | null + initialCheckpointHash?: string | null action?: | "chatButtonClicked" | "settingsButtonClicked" @@ -408,6 +412,10 @@ export type ExtensionState = Pick< featureRoomoteControlEnabled: boolean openAiCodexIsAuthenticated?: boolean debug?: boolean + + // Initial checkpoint state for state-driven UI + initialCheckpointState?: "pending" | "ready" | "failed" | null + initialCheckpointHash?: string | null } export interface Command { diff --git a/src/core/checkpoints/index.ts b/src/core/checkpoints/index.ts index 26a137b939c..d29e4ee3011 100644 --- a/src/core/checkpoints/index.ts +++ b/src/core/checkpoints/index.ts @@ -25,6 +25,14 @@ function sendCheckpointInitWarn(task: Task, type?: "WAIT_TIMEOUT" | "INIT_TIMEOU }) } +function sendInitialCheckpointState(task: Task, state: "pending" | "ready" | "failed" | null, hash?: string) { + task.providerRef.deref()?.postMessageToWebview({ + type: "initialCheckpointState", + initialCheckpointState: state, + initialCheckpointHash: hash ?? null, + }) +} + export async function getCheckpointService(task: Task, { interval = 250 }: { interval?: number } = {}) { if (!task.enableCheckpoints) { return undefined @@ -98,6 +106,7 @@ export async function getCheckpointService(task: Task, { interval = 250 }: { int ) if (!task?.checkpointService) { sendCheckpointInitWarn(task, "INIT_TIMEOUT", task.checkpointTimeout) + sendInitialCheckpointState(task, "failed") task.enableCheckpoints = false return undefined } else { @@ -121,6 +130,7 @@ export async function getCheckpointService(task: Task, { interval = 250 }: { int } catch (err) { if (err.name === "TimeoutError" && task.enableCheckpoints) { sendCheckpointInitWarn(task, "INIT_TIMEOUT", task.checkpointTimeout) + sendInitialCheckpointState(task, "failed") } log(`[Task#getCheckpointService] ${err.message}`) task.enableCheckpoints = false @@ -140,6 +150,7 @@ async function checkGitInstallation( if (!gitInstalled) { log("[Task#getCheckpointService] Git is not installed, disabling checkpoints") + sendInitialCheckpointState(task, "failed") task.enableCheckpoints = false task.checkpointServiceInitializing = false @@ -157,9 +168,18 @@ async function checkGitInstallation( } // Git is installed, proceed with initialization - service.on("initialize", () => { + service.on("initialize", ({ baseHash }) => { log("[Task#getCheckpointService] service initialized") task.checkpointServiceInitializing = false + + // Send initial checkpoint state as ready with the hash + sendInitialCheckpointState(task, "ready", baseHash) + + // Update webview with initial checkpoint hash (for currentCheckpoint tracking) + provider?.postMessageToWebview({ + type: "currentCheckpointUpdated", + text: baseHash, + }) }) service.on("checkpoint", ({ fromHash: from, toHash: to, suppressMessage }) => { @@ -195,15 +215,20 @@ async function checkGitInstallation( log("[Task#getCheckpointService] initializing shadow git") + // Set pending state before starting initialization + sendInitialCheckpointState(task, "pending") + try { await service.initShadowGit() } catch (err) { log(`[Task#getCheckpointService] initShadowGit -> ${err.message}`) + sendInitialCheckpointState(task, "failed") task.enableCheckpoints = false } } catch (err) { log(`[Task#getCheckpointService] Unexpected error during Git check: ${err.message}`) console.error("Git check error:", err) + sendInitialCheckpointState(task, "failed") task.enableCheckpoints = false task.checkpointServiceInitializing = false } diff --git a/webview-ui/src/components/chat/ChatView.tsx b/webview-ui/src/components/chat/ChatView.tsx index 5dcdf1998e7..32cb53cc428 100644 --- a/webview-ui/src/components/chat/ChatView.tsx +++ b/webview-ui/src/components/chat/ChatView.tsx @@ -44,6 +44,7 @@ import TaskHeader from "./TaskHeader" import SystemPromptWarning from "./SystemPromptWarning" import ProfileViolationWarning from "./ProfileViolationWarning" import { CheckpointWarning } from "./CheckpointWarning" +import { InitialCheckpoint } from "./checkpoints/InitialCheckpoint" import { QueuedMessages } from "./QueuedMessages" import { WorktreeSelector } from "./WorktreeSelector" import DismissibleUpsell from "../common/DismissibleUpsell" @@ -97,6 +98,9 @@ const ChatViewComponent: React.ForwardRefRenderFunction )} + + {/* Initial Checkpoint UI - state-driven element */} + {initialCheckpointState && enableCheckpoints && ( +
+ +
+ )} ) : (
diff --git a/webview-ui/src/components/chat/checkpoints/CheckpointMenu.tsx b/webview-ui/src/components/chat/checkpoints/CheckpointMenu.tsx index 4895d05b3a0..d3f0a7f3437 100644 --- a/webview-ui/src/components/chat/checkpoints/CheckpointMenu.tsx +++ b/webview-ui/src/components/chat/checkpoints/CheckpointMenu.tsx @@ -12,6 +12,7 @@ type CheckpointMenuBaseProps = { ts: number commitHash: string checkpoint: Checkpoint + isInitial?: boolean } type CheckpointMenuControlledProps = { onOpenChange: (open: boolean) => void @@ -21,7 +22,7 @@ type CheckpointMenuUncontrolledProps = { } type CheckpointMenuProps = CheckpointMenuBaseProps & (CheckpointMenuControlledProps | CheckpointMenuUncontrolledProps) -export const CheckpointMenu = ({ ts, commitHash, checkpoint, onOpenChange }: CheckpointMenuProps) => { +export const CheckpointMenu = ({ ts, commitHash, checkpoint, isInitial, onOpenChange }: CheckpointMenuProps) => { const { t } = useTranslation() const [internalRestoreOpen, setInternalRestoreOpen] = useState(false) const [restoreConfirming, setRestoreConfirming] = useState(false) @@ -95,11 +96,14 @@ export const CheckpointMenu = ({ ts, commitHash, checkpoint, onOpenChange }: Che return (
- - - + {/* Hide "View Diff" for initial checkpoint - no previous checkpoint to diff against */} + {!isInitial && ( + + + + )} { @@ -175,15 +179,18 @@ export const CheckpointMenu = ({ ts, commitHash, checkpoint, onOpenChange }: Che
- + {/* Hide "View All Changes" for initial checkpoint - already at init */} + {!isInitial && ( + + )}
diff --git a/webview-ui/src/components/chat/checkpoints/InitialCheckpoint.tsx b/webview-ui/src/components/chat/checkpoints/InitialCheckpoint.tsx new file mode 100644 index 00000000000..63328e26f65 --- /dev/null +++ b/webview-ui/src/components/chat/checkpoints/InitialCheckpoint.tsx @@ -0,0 +1,131 @@ +import { useMemo, useRef, useState, useEffect, useCallback } from "react" +import { useTranslation } from "react-i18next" +import { cn } from "@/lib/utils" + +import { CheckpointMenu } from "./CheckpointMenu" +import { GitCommitVertical, Loader2, AlertTriangle } from "lucide-react" +import { StandardTooltip } from "@/components/ui" + +type InitialCheckpointProps = { + state: "pending" | "ready" | "failed" + hash?: string | null +} + +export const InitialCheckpoint = ({ state, hash }: InitialCheckpointProps) => { + const { t } = useTranslation() + const [isPopoverOpen, setIsPopoverOpen] = useState(false) + const [isClosing, setIsClosing] = useState(false) + const [isHovering, setIsHovering] = useState(false) + const closeTimer = useRef(null) + + useEffect(() => { + return () => { + if (closeTimer.current) { + window.clearTimeout(closeTimer.current) + closeTimer.current = null + } + } + }, []) + + const handlePopoverOpenChange = useCallback((open: boolean) => { + setIsPopoverOpen(open) + if (open) { + setIsClosing(false) + if (closeTimer.current) { + window.clearTimeout(closeTimer.current) + closeTimer.current = null + } + } else { + setIsClosing(true) + closeTimer.current = window.setTimeout(() => { + setIsClosing(false) + closeTimer.current = null + }, 200) // keep menu visible briefly to avoid popover jump + } + }, []) + + const handleMouseEnter = useCallback(() => { + setIsHovering(true) + }, []) + + const handleMouseLeave = useCallback(() => { + setIsHovering(false) + }, []) + + // Menu is visible when hovering, popover is open, or briefly after popover closes + // But only when the state is 'ready' + const menuVisible = state === "ready" && (isHovering || isPopoverOpen || isClosing) + + // Create checkpoint metadata for the menu + const checkpointMetadata = useMemo(() => { + if (!hash) { + return undefined + } + return { + from: hash, + to: hash, + isInitial: true, + } + }, [hash]) + + const isPending = state === "pending" + const isReady = state === "ready" + const isFailed = state === "failed" + + return ( +
+
+ {isPending && } + {isReady && } + {isFailed && } + + {isPending && t("chat:checkpoint.initializing")} + {isReady && t("chat:checkpoint.initial")} + {isFailed && ( + + {t("chat:checkpoint.failed")} + + )} + +
+ + + {/* Only show menu when ready and hash is available */} + {isReady && hash && checkpointMetadata && ( +
+ +
+ )} +
+ ) +} diff --git a/webview-ui/src/components/chat/checkpoints/__tests__/CheckpointSaved.spec.tsx b/webview-ui/src/components/chat/checkpoints/__tests__/CheckpointSaved.spec.tsx index d2b6d48a3ff..8dd77920570 100644 --- a/webview-ui/src/components/chat/checkpoints/__tests__/CheckpointSaved.spec.tsx +++ b/webview-ui/src/components/chat/checkpoints/__tests__/CheckpointSaved.spec.tsx @@ -186,3 +186,147 @@ describe("CheckpointSaved popover visibility", () => { }) }) }) + +describe("CheckpointSaved label rendering", () => { + const baseProps = { + ts: 123, + commitHash: "abc123", + currentHash: "zzz999", + } + + it("renders initial checkpoint label when isInitial is true", () => { + const { getByText } = render( + } + />, + ) + + // Test uses i18n key since translations may not be loaded in test environment + expect(getByText("chat:checkpoint.initial")).toBeTruthy() + }) + + it("renders regular checkpoint label when isInitial is false", () => { + const { getByText } = render( + } + />, + ) + + expect(getByText("chat:checkpoint.regular")).toBeTruthy() + }) + + it("renders regular checkpoint label when isInitial is undefined", () => { + const { getByText } = render( + } + />, + ) + + expect(getByText("chat:checkpoint.regular")).toBeTruthy() + }) +}) + +describe("CheckpointMenu isInitial behavior", () => { + const baseProps = { + ts: 123, + commitHash: "abc123", + currentHash: "zzz999", + } + + it("hides View Diff button when isInitial is true", () => { + const { container } = render( + } + />, + ) + + // The View Diff button should not be rendered + const diffButton = container.querySelector('[aria-label="View Diff"]') + expect(diffButton).toBeNull() + }) + + it("shows View Diff button when isInitial is false", async () => { + const { container } = render( + } + />, + ) + + // Hover to make menu visible + const parentDiv = container.querySelector("[class*='flex items-center justify-between']") as HTMLElement + fireEvent.mouseEnter(parentDiv) + + // The View Diff button should be rendered (using codicon class as identifier) + await waitFor(() => { + const diffIcon = container.querySelector(".codicon-diff-single") + expect(diffIcon).toBeTruthy() + }) + }) + + it("hides View All Changes button when isInitial is true", async () => { + const { container } = render( + } + />, + ) + + // Hover to make menu visible + const parentDiv = container.querySelector("[class*='flex items-center justify-between']") as HTMLElement + fireEvent.mouseEnter(parentDiv) + + // Open the "more" popover + await waitForOpenHandler() + + // The "View All Changes" button with codicon-versions should not be rendered + const versionsIcon = container.querySelector(".codicon-versions") + expect(versionsIcon).toBeNull() + }) + + it("shows View Changes Since This Checkpoint regardless of isInitial", async () => { + const { container } = render( + } + />, + ) + + // Hover to make menu visible + const parentDiv = container.querySelector("[class*='flex items-center justify-between']") as HTMLElement + fireEvent.mouseEnter(parentDiv) + + // The "View Changes Since This Checkpoint" button with codicon-diff should be present + await waitFor(() => { + const diffIcon = container.querySelector(".codicon-diff") + expect(diffIcon).toBeTruthy() + }) + }) + + it("shows restore options regardless of isInitial", async () => { + const { getByTestId, container } = render( + } + />, + ) + + // Hover to make menu visible + const parentDiv = container.querySelector("[class*='flex items-center justify-between']") as HTMLElement + fireEvent.mouseEnter(parentDiv) + + // Open the restore popover + await waitForOpenHandler() + lastOnOpenChange?.(true) + + // Restore buttons should be available + await waitFor(() => { + expect(getByTestId("restore-files-btn")).toBeTruthy() + expect(getByTestId("restore-files-and-task-btn")).toBeTruthy() + }) + }) +}) diff --git a/webview-ui/src/components/chat/checkpoints/__tests__/InitialCheckpoint.spec.tsx b/webview-ui/src/components/chat/checkpoints/__tests__/InitialCheckpoint.spec.tsx new file mode 100644 index 00000000000..5cda181a293 --- /dev/null +++ b/webview-ui/src/components/chat/checkpoints/__tests__/InitialCheckpoint.spec.tsx @@ -0,0 +1,161 @@ +// npx vitest run src/components/chat/checkpoints/__tests__/InitialCheckpoint.spec.tsx + +// Capture onOpenChange from Popover to control open/close in tests +let _lastOnOpenChange: ((open: boolean) => void) | undefined + +vi.mock("@/components/ui", () => { + // Minimal UI primitives to ensure deterministic behavior in tests + return { + Button: ({ children, ...rest }: any) => , + StandardTooltip: ({ children }: any) => <>{children}, + Popover: (props: any) => { + const { children, onOpenChange, open, ...rest } = props + if (rest["data-testid"] === "restore-popover") { + _lastOnOpenChange = onOpenChange + } + return ( +
+ {children} +
+ ) + }, + PopoverTrigger: ({ children }: any) =>
{children}
, + PopoverContent: ({ children, className, ...rest }: any) => ( +
+ {children} +
+ ), + } +}) + +import { render, screen, fireEvent, waitFor } from "@/utils/test-utils" +import React from "react" +import { InitialCheckpoint } from "../InitialCheckpoint" + +describe("InitialCheckpoint", () => { + beforeEach(() => { + _lastOnOpenChange = undefined + }) + + describe("visual states", () => { + it("renders pending state correctly", () => { + const { getByTestId, getByText } = render() + + const container = getByTestId("initial-checkpoint") + expect(container).toBeTruthy() + + // Should show spinner + expect(getByTestId("initial-checkpoint-spinner")).toBeTruthy() + + // Should show initializing text + expect(getByText("chat:checkpoint.initializing")).toBeTruthy() + + // Menu should not be visible + expect(screen.queryByTestId("initial-checkpoint-menu-container")).toBeNull() + }) + + it("renders ready state correctly", () => { + const { getByTestId, getByText } = render() + + const container = getByTestId("initial-checkpoint") + expect(container).toBeTruthy() + + // Should not show spinner + expect(screen.queryByTestId("initial-checkpoint-spinner")).toBeNull() + + // Should show "Initial State" text + expect(getByText("chat:checkpoint.initial")).toBeTruthy() + + // Menu container should exist + expect(getByTestId("initial-checkpoint-menu-container")).toBeTruthy() + }) + + it("renders failed state correctly", () => { + const { getByTestId, getByText } = render() + + const container = getByTestId("initial-checkpoint") + expect(container).toBeTruthy() + + // Should not show spinner + expect(screen.queryByTestId("initial-checkpoint-spinner")).toBeNull() + + // Should show failed text + expect(getByText("chat:checkpoint.failed")).toBeTruthy() + + // Menu should not be visible + expect(screen.queryByTestId("initial-checkpoint-menu-container")).toBeNull() + }) + }) + + describe("menu visibility in ready state", () => { + it("hides menu by default when not hovering", () => { + const { getByTestId } = render() + + const menuContainer = getByTestId("initial-checkpoint-menu-container") + expect(menuContainer.className).toContain("hidden") + }) + + it("shows menu when hovering", async () => { + const { getByTestId } = render() + + const container = getByTestId("initial-checkpoint") + const menuContainer = getByTestId("initial-checkpoint-menu-container") + + // Initially hidden + expect(menuContainer.className).toContain("hidden") + + // Hover to show menu + fireEvent.mouseEnter(container) + + await waitFor(() => { + expect(menuContainer.className).toContain("block") + expect(menuContainer.className).not.toContain("hidden") + }) + + // Mouse leave to hide menu + fireEvent.mouseLeave(container) + + await waitFor(() => { + expect(menuContainer.className).toContain("hidden") + }) + }) + }) + + describe("styling", () => { + it("applies opacity styling for pending state", () => { + const { getByTestId } = render() + + const container = getByTestId("initial-checkpoint") + expect(container.className).toContain("opacity-50") + }) + + it("applies opacity styling for failed state", () => { + const { getByTestId } = render() + + const container = getByTestId("initial-checkpoint") + expect(container.className).toContain("opacity-75") + }) + + it("does not apply opacity for ready state", () => { + const { getByTestId } = render() + + const container = getByTestId("initial-checkpoint") + expect(container.className).not.toContain("opacity-50") + expect(container.className).not.toContain("opacity-75") + }) + }) + + describe("menu not rendered without hash in ready state", () => { + it("does not render menu when hash is null in ready state", () => { + render() + + expect(screen.queryByTestId("initial-checkpoint-menu-container")).toBeNull() + }) + + it("does not render menu when hash is undefined in ready state", () => { + render() + + expect(screen.queryByTestId("initial-checkpoint-menu-container")).toBeNull() + }) + }) +}) diff --git a/webview-ui/src/components/chat/checkpoints/schema.ts b/webview-ui/src/components/chat/checkpoints/schema.ts index 3c72a755608..e24b1789e41 100644 --- a/webview-ui/src/components/chat/checkpoints/schema.ts +++ b/webview-ui/src/components/chat/checkpoints/schema.ts @@ -3,6 +3,7 @@ import { z } from "zod" export const checkpointSchema = z.object({ from: z.string(), to: z.string(), + isInitial: z.boolean().optional(), }) export type Checkpoint = z.infer diff --git a/webview-ui/src/context/ExtensionStateContext.tsx b/webview-ui/src/context/ExtensionStateContext.tsx index 1d2c43ff008..6b3270d3e66 100644 --- a/webview-ui/src/context/ExtensionStateContext.tsx +++ b/webview-ui/src/context/ExtensionStateContext.tsx @@ -39,6 +39,8 @@ export interface ExtensionStateContextType extends ExtensionState { mcpServers: McpServer[] hasSystemPromptOverride?: boolean currentCheckpoint?: string + initialCheckpointState?: "pending" | "ready" | "failed" | null + initialCheckpointHash?: string | null currentTaskTodos?: TodoItem[] // Initial todos for the current task filePaths: string[] openedTabs: Array<{ label: string; isActive: boolean; path?: string }> @@ -283,6 +285,8 @@ export const ExtensionStateContextProvider: React.FC<{ children: React.ReactNode const [skills, setSkills] = useState([]) const [mcpServers, setMcpServers] = useState([]) const [currentCheckpoint, setCurrentCheckpoint] = useState() + const [initialCheckpointState, setInitialCheckpointState] = useState<"pending" | "ready" | "failed" | null>(null) + const [initialCheckpointHash, setInitialCheckpointHash] = useState(null) const [extensionRouterModels, setExtensionRouterModels] = useState(undefined) const [marketplaceItems, setMarketplaceItems] = useState([]) const [alwaysAllowFollowupQuestions, setAlwaysAllowFollowupQuestions] = useState(false) // Add state for follow-up questions auto-approve @@ -405,6 +409,11 @@ export const ExtensionStateContextProvider: React.FC<{ children: React.ReactNode setCurrentCheckpoint(message.text) break } + case "initialCheckpointState": { + setInitialCheckpointState(message.initialCheckpointState ?? null) + setInitialCheckpointHash(message.initialCheckpointHash ?? null) + break + } case "listApiConfig": { setListApiConfigMeta(message.listApiConfig ?? []) break @@ -492,6 +501,8 @@ export const ExtensionStateContextProvider: React.FC<{ children: React.ReactNode theme, mcpServers, currentCheckpoint, + initialCheckpointState, + initialCheckpointHash, filePaths, openedTabs, commands, diff --git a/webview-ui/src/i18n/locales/en/chat.json b/webview-ui/src/i18n/locales/en/chat.json index d167a19ff3e..653ca52e2c0 100644 --- a/webview-ui/src/i18n/locales/en/chat.json +++ b/webview-ui/src/i18n/locales/en/chat.json @@ -171,6 +171,10 @@ }, "checkpoint": { "regular": "Checkpoint", + "initial": "Initial State", + "initializing": "Initializing checkpoint...", + "failed": "Checkpoint initialization failed", + "failedDescription": "Unable to create initial checkpoint. Restore features may be unavailable for this task.", "initializingWarning": "Still initializing checkpoint... If this takes too long, you can disable checkpoints in settings and restart your task.", "menu": { "viewDiff": "View Diff", From a89acc58ed47ec80768d72b3d53691fd3650b9ee Mon Sep 17 00:00:00 2001 From: Hannes Rudolph Date: Thu, 29 Jan 2026 13:49:31 -0700 Subject: [PATCH 2/3] fix: initial checkpoint restore now uses checkpointRestoreToBase The initial checkpoint has no message in clineMessages (ts=0 doesn't exist), so checkpointRestore was returning early without doing anything. Now when isInitial=true is passed in the restore payload, the handler calls checkpointRestoreToBase() instead which properly restores to the baseHash without needing to find a message. --- packages/types/src/vscode-extension-host.ts | 1 + src/core/webview/webviewMessageHandler.ts | 13 ++++++++++++- .../components/chat/checkpoints/CheckpointMenu.tsx | 14 ++++++++++---- 3 files changed, 23 insertions(+), 5 deletions(-) diff --git a/packages/types/src/vscode-extension-host.ts b/packages/types/src/vscode-extension-host.ts index 77c8597a285..f88cf9fcc41 100644 --- a/packages/types/src/vscode-extension-host.ts +++ b/packages/types/src/vscode-extension-host.ts @@ -737,6 +737,7 @@ export const checkoutRestorePayloadSchema = z.object({ ts: z.number(), commitHash: z.string(), mode: z.enum(["preview", "restore"]), + isInitial: z.boolean().optional(), }) export type CheckpointRestorePayload = z.infer diff --git a/src/core/webview/webviewMessageHandler.ts b/src/core/webview/webviewMessageHandler.ts index 8246dda472e..de272773c94 100644 --- a/src/core/webview/webviewMessageHandler.ts +++ b/src/core/webview/webviewMessageHandler.ts @@ -1200,7 +1200,18 @@ export const webviewMessageHandler = async ( } try { - await provider.getCurrentTask()?.checkpointRestore(result.data) + // For initial checkpoint, use checkpointRestoreToBase instead of checkpointRestore + // because there's no message with ts to find in clineMessages + if (result.data.isInitial) { + const { checkpointRestoreToBase } = await import("../checkpoints") + const success = await checkpointRestoreToBase(provider.getCurrentTask()!) + + if (!success) { + vscode.window.showErrorMessage(t("common:errors.checkpoint_restore_base_failed")) + } + } else { + await provider.getCurrentTask()?.checkpointRestore(result.data) + } } catch (error) { vscode.window.showErrorMessage(t("common:errors.checkpoint_failed")) } diff --git a/webview-ui/src/components/chat/checkpoints/CheckpointMenu.tsx b/webview-ui/src/components/chat/checkpoints/CheckpointMenu.tsx index d3f0a7f3437..d3a77301142 100644 --- a/webview-ui/src/components/chat/checkpoints/CheckpointMenu.tsx +++ b/webview-ui/src/components/chat/checkpoints/CheckpointMenu.tsx @@ -75,14 +75,20 @@ export const CheckpointMenu = ({ ts, commitHash, checkpoint, isInitial, onOpenCh }, [ts, commitHash]) const onPreview = useCallback(() => { - vscode.postMessage({ type: "checkpointRestore", payload: { ts, commitHash, mode: "preview" } }) + vscode.postMessage({ + type: "checkpointRestore", + payload: { ts, commitHash, mode: "preview", ...(isInitial && { isInitial: true }) }, + }) setRestoreOpen(false) - }, [ts, commitHash, setRestoreOpen]) + }, [ts, commitHash, isInitial, setRestoreOpen]) const onRestore = useCallback(() => { - vscode.postMessage({ type: "checkpointRestore", payload: { ts, commitHash, mode: "restore" } }) + vscode.postMessage({ + type: "checkpointRestore", + payload: { ts, commitHash, mode: "restore", ...(isInitial && { isInitial: true }) }, + }) setRestoreOpen(false) - }, [ts, commitHash, setRestoreOpen]) + }, [ts, commitHash, isInitial, setRestoreOpen]) const handleOpenChange = useCallback( (open: boolean) => { From 7ee88cc638962aeb36cf3d15990e38a02d8ee708 Mon Sep 17 00:00:00 2001 From: Hannes Rudolph Date: Thu, 29 Jan 2026 13:59:33 -0700 Subject: [PATCH 3/3] fix: add missing checkpointRestoreToBase function This function was in the old PR #10908 but was not included when we created the new branch from origin/main. It's needed for restoring to the initial checkpoint state. --- src/core/checkpoints/index.ts | 38 +++++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/src/core/checkpoints/index.ts b/src/core/checkpoints/index.ts index d29e4ee3011..1859b45ca9f 100644 --- a/src/core/checkpoints/index.ts +++ b/src/core/checkpoints/index.ts @@ -326,6 +326,44 @@ export async function checkpointRestore( } } +/** + * Restore to the initial checkpoint (base state when task started). + * This is used by the InitialCheckpoint component since it doesn't have a message timestamp. + */ +export async function checkpointRestoreToBase(task: Task): Promise { + const service = await getCheckpointService(task) + + if (!service) { + return false + } + + const baseHash = service.baseHash + + if (!baseHash) { + const provider = task.providerRef.deref() + provider?.log("[checkpointRestoreToBase] no baseHash available") + return false + } + + const provider = task.providerRef.deref() + + try { + await service.restoreCheckpoint(baseHash) + TelemetryService.instance.captureCheckpointRestored(task.taskId) + await provider?.postMessageToWebview({ type: "currentCheckpointUpdated", text: baseHash }) + + // Cancel the task to reinitialize with the restored state + // This follows the same pattern as checkpointRestore + provider?.cancelTask() + + return true + } catch (err) { + provider?.log("[checkpointRestoreToBase] disabling checkpoints for this task") + task.enableCheckpoints = false + return false + } +} + export type CheckpointDiffOptions = { ts?: number previousCommitHash?: string