diff --git a/src/core/checkpoints/__tests__/checkpoint.test.ts b/src/core/checkpoints/__tests__/checkpoint.test.ts index e073c0cb92..d8e87c9710 100644 --- a/src/core/checkpoints/__tests__/checkpoint.test.ts +++ b/src/core/checkpoints/__tests__/checkpoint.test.ts @@ -1,4 +1,4 @@ -import { describe, it, expect, vi, beforeEach, afterEach } from "vitest" +import { describe, it, expect, vi, beforeEach, afterEach, beforeAll, afterAll } from "vitest" import { Task } from "../../task/Task" import { ClineProvider } from "../../webview/ClineProvider" import { checkpointSave, checkpointRestore, checkpointDiff, getCheckpointService } from "../index" @@ -394,11 +394,40 @@ describe("Checkpoint functionality", () => { expect(service).toBeUndefined() }) - it("should return undefined if service is still initializing", async () => { + it("should handle initialization with isInitialCall flag", async () => { + // Test that isInitialCall=true doesn't use timeout mockTask.checkpointService = undefined - mockTask.checkpointServiceInitializing = true + mockTask.checkpointServiceInitializing = false + mockTask.enableCheckpoints = true + + // Mock the RepoPerTaskCheckpointService.create to return our mock service + const checkpointsModule = await import("../../../services/checkpoints") + vi.mocked(checkpointsModule.RepoPerTaskCheckpointService.create).mockReturnValue(mockCheckpointService) + + // Call with isInitialCall=true + const service = await getCheckpointService(mockTask, { isInitialCall: true }) + + // Should create and return the service + expect(service).toBe(mockCheckpointService) + expect(mockTask.checkpointService).toBe(mockCheckpointService) + }) + + it("should handle initialization without isInitialCall flag", async () => { + // Test that without isInitialCall, it uses the default timeout + mockTask.checkpointService = undefined + mockTask.checkpointServiceInitializing = false + mockTask.enableCheckpoints = true + + // Mock the RepoPerTaskCheckpointService.create to return our mock service + const checkpointsModule = await import("../../../services/checkpoints") + vi.mocked(checkpointsModule.RepoPerTaskCheckpointService.create).mockReturnValue(mockCheckpointService) + + // Call without isInitialCall (defaults to false) const service = await getCheckpointService(mockTask) - expect(service).toBeUndefined() + + // Should create and return the service + expect(service).toBe(mockCheckpointService) + expect(mockTask.checkpointService).toBe(mockCheckpointService) }) it("should create new service if none exists", async () => { diff --git a/src/core/checkpoints/index.ts b/src/core/checkpoints/index.ts index bc842c9f18..fb45b17c12 100644 --- a/src/core/checkpoints/index.ts +++ b/src/core/checkpoints/index.ts @@ -18,7 +18,11 @@ import { CheckpointServiceOptions, RepoPerTaskCheckpointService } from "../../se export async function getCheckpointService( task: Task, - { interval = 250, timeout = 15_000 }: { interval?: number; timeout?: number } = {}, + { + interval = 250, + timeout = 15_000, + isInitialCall = false, + }: { interval?: number; timeout?: number; isInitialCall?: boolean } = {}, ) { if (!task.enableCheckpoints) { return undefined @@ -67,13 +71,25 @@ export async function getCheckpointService( } if (task.checkpointServiceInitializing) { - await pWaitFor( - () => { + // For initial call during task startup, don't use timeout to allow large repos to initialize + const waitOptions = isInitialCall ? { interval } : { interval, timeout } + + try { + await pWaitFor(() => { console.log("[Task#getCheckpointService] waiting for service to initialize") return !!task.checkpointService && !!task?.checkpointService?.isInitialized - }, - { interval, timeout }, - ) + }, waitOptions) + } catch (err) { + // Only disable checkpoints if this is not the initial call and we hit a timeout + if (!isInitialCall) { + log("[Task#getCheckpointService] timeout waiting for service initialization, disabling checkpoints") + task.enableCheckpoints = false + return undefined + } + // For initial call, continue waiting or handle the error differently + throw err + } + if (!task?.checkpointService) { task.enableCheckpoints = false return undefined diff --git a/src/core/task/Task.ts b/src/core/task/Task.ts index c5be865731..30956aa658 100644 --- a/src/core/task/Task.ts +++ b/src/core/task/Task.ts @@ -1661,7 +1661,8 @@ export class Task extends EventEmitter implements TaskLike { private async initiateTaskLoop(userContent: Anthropic.Messages.ContentBlockParam[]): Promise { // Kicks off the checkpoints initialization process in the background. - getCheckpointService(this) + // For the initial call, don't use a timeout to allow large repositories to complete initialization + getCheckpointService(this, { isInitialCall: true }) let nextUserContent = userContent let includeFileDetails = true