From 305ad95a20e7ac8a111fd8e10e57e5ec11fc6507 Mon Sep 17 00:00:00 2001 From: can1357 Date: Fri, 31 Jul 2026 19:10:14 +0200 Subject: [PATCH] fix(task): guard deferred launch work --- packages/coding-agent/src/task/executor.ts | 18 +-- .../test/task/executor-launch-startup.test.ts | 103 ++++++++++++++++++ .../test/task/executor-prewalk.test.ts | 29 +++++ 3 files changed, 143 insertions(+), 7 deletions(-) create mode 100644 packages/coding-agent/test/task/executor-launch-startup.test.ts diff --git a/packages/coding-agent/src/task/executor.ts b/packages/coding-agent/src/task/executor.ts index 098f43e09..d9b72882f 100644 --- a/packages/coding-agent/src/task/executor.ts +++ b/packages/coding-agent/src/task/executor.ts @@ -2685,6 +2685,16 @@ export async function runSubprocess(options: ExecutorOptions): Promise {}); // Per-agent prewalk: the agent definition's `prewalk` frontmatter or the // `task.agentPrewalk` settings override hands the subagent off to a // fast/cheap target at its first edit/write — the same mechanism as the @@ -2697,6 +2707,7 @@ export async function runSubprocess(options: ExecutorOptions): Promise { + vi.restoreAllMocks(); + for (const authStorage of authStorages.splice(0)) await authStorage.close(); + for (const tempDir of tempDirs.splice(0)) tempDir[Symbol.dispose](); +}); + +it("overlaps registry refresh with session-file opening and session setup", async () => { + const tempDir = TempDir.createSync("@pi-task-launch-"); + tempDirs.push(tempDir); + const authStorage = await AuthStorage.create(tempDir.join("auth.db")); + authStorages.push(authStorage); + + const refreshGate = Promise.withResolvers(); + vi.spyOn(ModelRegistry.prototype, "refresh").mockImplementation(() => refreshGate.promise); + + const sessionManager = SessionManager.inMemory(tempDir.path()); + const openGate = Promise.withResolvers(); + const openStarted = Promise.withResolvers(); + const openSpy = vi.spyOn(SessionManager, "open").mockImplementation(() => { + openStarted.resolve(); + return openGate.promise; + }); + + const sessionCreationStarted = Promise.withResolvers(); + let sessionCreated = false; + const listeners: Array<(event: AgentSessionEvent) => void> = []; + const session = { + state: { messages: [] }, + agent: { state: { systemPrompt: ["test"] } }, + model: undefined, + extensionRunner: undefined, + sessionManager: { appendSessionInit: () => {} }, + getActiveToolNames: () => ["yield"], + getEnabledToolNames: () => ["yield"], + setActiveToolsByName: async () => {}, + subscribe: (listener: (event: AgentSessionEvent) => void) => { + listeners.push(listener); + return () => {}; + }, + prompt: async () => { + for (const listener of listeners) { + listener({ + type: "tool_execution_end", + toolCallId: "yield", + toolName: "yield", + result: { content: [], details: { status: "success", data: { ok: true } } }, + isError: false, + } as AgentSessionEvent); + } + }, + waitForIdle: async () => {}, + getLastAssistantMessage: () => undefined, + abort: async () => {}, + dispose: async () => {}, + } as unknown as AgentSession; + vi.spyOn(sdkModule, "createAgentSession").mockImplementation(async () => { + sessionCreationStarted.resolve(); + sessionCreated = true; + return { + session, + extensionsResult: {}, + setToolUIContext: () => {}, + eventBus: new EventBus(), + } as CreateAgentSessionResult; + }); + + const run = runSubprocess({ + cwd: tempDir.path(), + artifactsDir: tempDir.path(), + agent: { name: "task", description: "test", systemPrompt: "test", source: "bundled" }, + task: "test", + index: 0, + id: "task-launch-overlap", + authStorage, + enableLsp: false, + enableIrc: false, + }); + await openStarted.promise; + + expect(openSpy).toHaveBeenCalledTimes(1); + expect(sessionCreated).toBe(false); + + openGate.resolve(sessionManager); + await sessionCreationStarted.promise; + expect(sessionCreated).toBe(true); + + refreshGate.resolve(); + expect((await run).exitCode).toBe(0); +}); diff --git a/packages/coding-agent/test/task/executor-prewalk.test.ts b/packages/coding-agent/test/task/executor-prewalk.test.ts index dfdf8fb83..9096c2124 100644 --- a/packages/coding-agent/test/task/executor-prewalk.test.ts +++ b/packages/coding-agent/test/task/executor-prewalk.test.ts @@ -96,6 +96,7 @@ function createModelRegistry(models: Model[]): ModelRegistry { return { authStorage: {}, refresh: async () => {}, + awaitBackgroundRefresh: async () => {}, getAvailable: () => models, getApiKey: async () => "test-key", hasConfiguredAuth: () => true, @@ -149,6 +150,34 @@ describe("runSubprocess per-agent prewalk", () => { expect(forwarded?.prewalk?.target.provider).toBe(target.provider); }); + it("waits for background discovery before resolving a configured prewalk target", async () => { + const models = [primary]; + const registry = createModelRegistry(models); + const refreshGate = Promise.withResolvers(); + vi.spyOn(registry, "awaitBackgroundRefresh").mockImplementation(async () => { + await refreshGate.promise; + models.push(target); + }); + const spy = vi + .spyOn(sdkModule, "createAgentSession") + .mockResolvedValue(createSessionResult(yieldEmittingSession())); + + const run = runSubprocess({ + ...baseOptions("subagent-prewalk-discovery", Settings.isolated()), + modelRegistry: registry, + agent: { + ...baseAgent, + model: [`${primary.provider}/${primary.id}`], + prewalk: `${target.provider}/${target.id}`, + }, + }); + expect(spy).not.toHaveBeenCalled(); + + refreshGate.resolve(); + expect((await run).exitCode).toBe(0); + expect(spy.mock.calls[0]?.[0]?.prewalk?.target.id).toBe(target.id); + }); + it("reports the prewalk target as the active model after handoff", async () => { const progressModels: string[] = []; vi.spyOn(sdkModule, "createAgentSession").mockResolvedValue(