diff --git a/packages/agent/src/harness/agent-harness.ts b/packages/agent/src/harness/agent-harness.ts index cbc2d0a4a..f97d0210e 100644 --- a/packages/agent/src/harness/agent-harness.ts +++ b/packages/agent/src/harness/agent-harness.ts @@ -177,8 +177,10 @@ export class AgentHarness< private session: Session; readonly models: Models; private phase: AgentHarnessPhase = "idle"; - private runAbortController?: AbortController; - private runPromise?: Promise; + private activeAbortController?: AbortController; + private activeOperationPromise?: Promise; + private shutdownPromise?: Promise; + private isShutdown = false; private pendingSessionWrites: PendingSessionWrite[] = []; private model: Model; private thinkingLevel: ThinkingLevel; @@ -222,6 +224,10 @@ export class AgentHarness< this.followUpQueueMode = options.followUpMode ?? "one-at-a-time"; } + private assertNotShutDown(): void { + if (this.isShutdown) throw new AgentHarnessError("invalid_state", "AgentHarness has been shut down"); + } + private getHandlers(type: string): Set | undefined { return this.handlers.get(type); } @@ -326,14 +332,20 @@ export class AgentHarness< }); } - private startRunPromise(): () => void { + private startOperation(): { signal: AbortSignal; finish: () => void } { + const abortController = new AbortController(); let finish = () => {}; - this.runPromise = new Promise((resolve) => { + this.activeAbortController = abortController; + this.activeOperationPromise = new Promise((resolve) => { finish = resolve; }); - return () => { - this.runPromise = undefined; - finish(); + return { + signal: abortController.signal, + finish: () => { + this.activeAbortController = undefined; + this.activeOperationPromise = undefined; + finish(); + }, }; } @@ -352,6 +364,7 @@ export class AgentHarness< } private async createTurnState(): Promise> { + this.assertNotShutDown(); const context = await this.session.buildContext(); const resources = this.getResources(); const sessionMetadata = await this.session.getMetadata(); @@ -581,8 +594,10 @@ export class AgentHarness< private async executeTurn( turnState: AgentHarnessTurnState, text: string, + signal: AbortSignal, options?: { images?: ImageContent[] }, ): Promise { + this.assertNotShutDown(); let activeTurnState = turnState; let messages: AgentMessage[] = [createUserMessage(text, options?.images)]; if (this.nextTurnQueue.length > 0) { @@ -602,32 +617,26 @@ export class AgentHarness< systemPrompt: turnState.systemPrompt, resources: turnState.resources, }); + this.assertNotShutDown(); if (beforeResult?.messages) messages = [...messages, ...beforeResult.messages]; - const abortController = new AbortController(); const getTurnState = () => activeTurnState; const setTurnState = (nextTurnState: AgentHarnessTurnState) => { activeTurnState = nextTurnState; }; - this.runAbortController = abortController; const runResultPromise = (async () => { try { return await runAgentLoop( messages, this.createContext(turnState, beforeResult?.systemPrompt), this.createLoopConfig(getTurnState, setTurnState), - (event) => this.handleAgentEvent(event, abortController.signal), - abortController.signal, + (event) => this.handleAgentEvent(event, signal), + signal, this.createStreamFn(getTurnState), ); } catch (error) { try { - return await this.emitRunFailure( - activeTurnState.model, - error, - abortController.signal.aborted, - abortController.signal, - ); + return await this.emitRunFailure(activeTurnState.model, error, signal.aborted, signal); } catch (failureError) { const cause = new AggregateError( [toError(error), toError(failureError)], @@ -647,81 +656,88 @@ export class AgentHarness< } throw new AgentHarnessError("invalid_state", "AgentHarness prompt completed without an assistant message"); } finally { - try { - await this.flushPendingSessionWrites(); - } finally { - this.runAbortController = undefined; - } + await this.flushPendingSessionWrites(); } } async prompt(text: string, options?: { images?: ImageContent[] }): Promise { + this.assertNotShutDown(); if (this.phase !== "idle") throw new AgentHarnessError("busy", "AgentHarness is busy"); this.phase = "turn"; - const finishRunPromise = this.startRunPromise(); + const operation = this.startOperation(); try { const turnState = await this.createTurnState(); - return await this.executeTurn(turnState, text, options); + return await this.executeTurn(turnState, text, operation.signal, options); } catch (error) { this.phase = "idle"; throw normalizeHarnessError(error, "unknown"); } finally { - finishRunPromise(); + operation.finish(); } } async skill(name: string, additionalInstructions?: string): Promise { + this.assertNotShutDown(); if (this.phase !== "idle") throw new AgentHarnessError("busy", "AgentHarness is busy"); this.phase = "turn"; - const finishRunPromise = this.startRunPromise(); + const operation = this.startOperation(); try { const turnState = await this.createTurnState(); const skill = (turnState.resources.skills ?? []).find((candidate) => candidate.name === name); if (!skill) throw new AgentHarnessError("invalid_argument", `Unknown skill: ${name}`); - return await this.executeTurn(turnState, formatSkillInvocation(skill, additionalInstructions)); + return await this.executeTurn( + turnState, + formatSkillInvocation(skill, additionalInstructions), + operation.signal, + ); } catch (error) { this.phase = "idle"; throw normalizeHarnessError(error, "unknown"); } finally { - finishRunPromise(); + operation.finish(); } } async promptFromTemplate(name: string, args: string[] = []): Promise { + this.assertNotShutDown(); if (this.phase !== "idle") throw new AgentHarnessError("busy", "AgentHarness is busy"); this.phase = "turn"; - const finishRunPromise = this.startRunPromise(); + const operation = this.startOperation(); try { const turnState = await this.createTurnState(); const template = (turnState.resources.promptTemplates ?? []).find((candidate) => candidate.name === name); if (!template) throw new AgentHarnessError("invalid_argument", `Unknown prompt template: ${name}`); - return await this.executeTurn(turnState, formatPromptTemplateInvocation(template, args)); + return await this.executeTurn(turnState, formatPromptTemplateInvocation(template, args), operation.signal); } catch (error) { this.phase = "idle"; throw normalizeHarnessError(error, "unknown"); } finally { - finishRunPromise(); + operation.finish(); } } async steer(text: string, options?: { images?: ImageContent[] }): Promise { + this.assertNotShutDown(); if (this.phase === "idle") throw new AgentHarnessError("invalid_state", "Cannot steer while idle"); this.steerQueue.push(createUserMessage(text, options?.images)); await this.emitQueueUpdate(); } async followUp(text: string, options?: { images?: ImageContent[] }): Promise { + this.assertNotShutDown(); if (this.phase === "idle") throw new AgentHarnessError("invalid_state", "Cannot follow up while idle"); this.followUpQueue.push(createUserMessage(text, options?.images)); await this.emitQueueUpdate(); } async nextTurn(text: string, options?: { images?: ImageContent[] }): Promise { + this.assertNotShutDown(); this.nextTurnQueue.push(createUserMessage(text, options?.images)); await this.emitQueueUpdate(); } async appendMessage(message: AgentMessage): Promise { + this.assertNotShutDown(); try { if (this.phase === "idle") { await this.session.appendMessage(message); @@ -734,8 +750,10 @@ export class AgentHarness< } async compact(customInstructions?: string): Promise { + this.assertNotShutDown(); if (this.phase !== "idle") throw new AgentHarnessError("busy", "compact() requires idle harness"); this.phase = "compaction"; + const operation = this.startOperation(); try { const model = this.model; if (!model) throw new AgentHarnessError("invalid_state", "No model set for compaction"); @@ -749,7 +767,7 @@ export class AgentHarness< preparation, branchEntries, customInstructions, - signal: new AbortController().signal, + signal: operation.signal, }); if (hookResult?.cancel) throw new AgentHarnessError("compaction", "Compaction cancelled"); const provided = hookResult?.compaction; @@ -760,13 +778,14 @@ export class AgentHarness< this.models, model, customInstructions, - undefined, + operation.signal, this.thinkingLevel, this.retry, this.retryCallbacks("compaction"), ); if (!compactResult.ok) throw compactResult.error; const result = compactResult.value; + this.assertNotShutDown(); const entryId = await this.session.appendCompaction( result.summary, result.firstKeptEntryId, @@ -785,6 +804,7 @@ export class AgentHarness< throw normalizeHarnessError(error, "compaction"); } finally { this.phase = "idle"; + operation.finish(); } } @@ -792,8 +812,10 @@ export class AgentHarness< targetId: string, options?: { summarize?: boolean; customInstructions?: string; replaceInstructions?: boolean; label?: string }, ): Promise { + this.assertNotShutDown(); if (this.phase !== "idle") throw new AgentHarnessError("busy", "navigateTree() requires idle harness"); this.phase = "branch_summary"; + const operation = this.startOperation(); try { const oldLeafId = await this.session.getLeafId(); if (oldLeafId === targetId) return { cancelled: false }; @@ -810,8 +832,11 @@ export class AgentHarness< replaceInstructions: options?.replaceInstructions, label: options?.label, }; - const signal = new AbortController().signal; - const hookResult = await this.emitHook({ type: "session_before_tree", preparation, signal }); + const hookResult = await this.emitHook({ + type: "session_before_tree", + preparation, + signal: operation.signal, + }); if (hookResult?.cancel) return { cancelled: true }; let summaryEntry: NavigateTreeResult["summaryEntry"]; let summaryText: string | undefined = hookResult?.summary?.summary; @@ -823,7 +848,7 @@ export class AgentHarness< const branchSummary = await generateBranchSummary(entries, { models: this.models, model, - signal: new AbortController().signal, + signal: operation.signal, customInstructions: hookResult?.customInstructions ?? options?.customInstructions, replaceInstructions: hookResult?.replaceInstructions ?? options?.replaceInstructions, retry: this.retry, @@ -851,6 +876,7 @@ export class AgentHarness< } else { newLeafId = targetId; } + this.assertNotShutDown(); const summaryId = await this.session.moveTo( newLeafId, summaryText @@ -878,6 +904,7 @@ export class AgentHarness< throw normalizeHarnessError(error, "branch_summary"); } finally { this.phase = "idle"; + operation.finish(); } } @@ -886,6 +913,7 @@ export class AgentHarness< } async setModel(model: Model): Promise { + this.assertNotShutDown(); try { const previousModel = this.model; if (this.phase === "idle") { @@ -905,6 +933,7 @@ export class AgentHarness< } async setThinkingLevel(level: ThinkingLevel): Promise { + this.assertNotShutDown(); try { const previousLevel = this.thinkingLevel; if (this.phase === "idle") { @@ -924,6 +953,7 @@ export class AgentHarness< } async setTools(tools: TTool[], activeToolNames?: string[]): Promise { + this.assertNotShutDown(); try { this.validateUniqueNames( tools.map((tool) => tool.name), @@ -959,6 +989,7 @@ export class AgentHarness< } async setActiveTools(toolNames: string[]): Promise { + this.assertNotShutDown(); try { this.validateToolNames(toolNames); const previousToolNames = [...this.tools.keys()]; @@ -987,6 +1018,7 @@ export class AgentHarness< } async setSteeringMode(mode: QueueMode): Promise { + this.assertNotShutDown(); this.steeringQueueMode = mode; } @@ -995,6 +1027,7 @@ export class AgentHarness< } async setFollowUpMode(mode: QueueMode): Promise { + this.assertNotShutDown(); this.followUpQueueMode = mode; } @@ -1006,6 +1039,7 @@ export class AgentHarness< } async setResources(resources: AgentHarnessResources): Promise { + this.assertNotShutDown(); const previousResources = this.getResources(); this.resources = { skills: resources.skills?.slice(), @@ -1019,15 +1053,33 @@ export class AgentHarness< } async setStreamOptions(streamOptions: AgentHarnessStreamOptions): Promise { + this.assertNotShutDown(); this.streamOptions = cloneStreamOptions(streamOptions); } + /** + * Permanently stop this harness instance without deleting its durable session. + * Clears queued work, aborts the active operation, and waits for it to settle. + */ + async shutdown(): Promise { + if (this.shutdownPromise) return this.shutdownPromise; + this.isShutdown = true; + this.pendingSessionWrites = []; + this.steerQueue = []; + this.followUpQueue = []; + this.nextTurnQueue = []; + this.activeAbortController?.abort(); + this.shutdownPromise = this.waitForIdle(); + return this.shutdownPromise; + } + async abort(): Promise { + this.assertNotShutDown(); const clearedSteer = [...this.steerQueue]; const clearedFollowUp = [...this.followUpQueue]; this.steerQueue = []; this.followUpQueue = []; - this.runAbortController?.abort(); + this.activeAbortController?.abort(); const errors: Error[] = []; try { await this.emitQueueUpdate(); @@ -1052,12 +1104,13 @@ export class AgentHarness< } async waitForIdle(): Promise { - await this.runPromise; + await this.activeOperationPromise; } subscribe( listener: (event: AgentHarnessEvent, signal?: AbortSignal) => Promise | void, ): () => void { + this.assertNotShutDown(); let handlers = this.handlers.get(SUBSCRIBER_EVENT_TYPE); if (!handlers) { handlers = new Set(); @@ -1073,6 +1126,7 @@ export class AgentHarness< event: Extract, ) => Promise | AgentHarnessEventResultMap[TType], ): () => void { + this.assertNotShutDown(); let handlers = this.handlers.get(type); if (!handlers) { handlers = new Set(); diff --git a/packages/agent/test/harness/agent-harness.test.ts b/packages/agent/test/harness/agent-harness.test.ts index 53e8958f7..d678427b0 100644 --- a/packages/agent/test/harness/agent-harness.test.ts +++ b/packages/agent/test/harness/agent-harness.test.ts @@ -112,6 +112,177 @@ describe("AgentHarness", () => { expect(harness.getFollowUpMode()).toBe("one-at-a-time"); }); + it("shuts down active work permanently and idempotently", async () => { + const registration = newFaux(); + const entered = deferred(); + const release = deferred(); + let signal: AbortSignal | undefined; + registration.setResponses([ + async (_context, options) => { + signal = options?.signal; + entered.resolve(); + await release.promise; + return fauxAssistantMessage("finished after shutdown"); + }, + ]); + const harness = new AgentHarness({ + models, + session: new Session(new InMemorySessionStorage()), + model: registration.getModel(), + }); + const prompt = harness.prompt("hello"); + await entered.promise; + await harness.steer("queued steer"); + await harness.followUp("queued follow-up"); + await harness.nextTurn("queued next turn"); + + let firstShutdownSettled = false; + const firstShutdown = harness.shutdown().then(() => { + firstShutdownSettled = true; + }); + const secondShutdown = harness.shutdown(); + await Promise.resolve(); + + expect(signal?.aborted).toBe(true); + expect(firstShutdownSettled).toBe(false); + release.resolve(); + await expect(prompt).resolves.toMatchObject({ role: "assistant" }); + await expect(Promise.all([firstShutdown, secondShutdown])).resolves.toEqual([undefined, undefined]); + await expect(harness.prompt("again")).rejects.toMatchObject({ + code: "invalid_state", + message: "AgentHarness has been shut down", + }); + await expect(harness.nextTurn("again")).rejects.toMatchObject({ code: "invalid_state" }); + await expect(harness.appendMessage(createUserMessage("again"))).rejects.toMatchObject({ + code: "invalid_state", + }); + }); + + it("does not start a provider request when shutdown occurs during before_agent_start", async () => { + const registration = newFaux(); + const entered = deferred(); + const release = deferred(); + let providerCalls = 0; + registration.setResponses([ + () => { + providerCalls++; + return fauxAssistantMessage("must not run"); + }, + ]); + const harness = new AgentHarness({ + models, + session: new Session(new InMemorySessionStorage()), + model: registration.getModel(), + }); + harness.on("before_agent_start", async () => { + entered.resolve(); + await release.promise; + return undefined; + }); + const prompt = harness.prompt("hello"); + await entered.promise; + + let shutdownSettled = false; + const shutdown = harness.shutdown().then(() => { + shutdownSettled = true; + }); + await Promise.resolve(); + + expect(shutdownSettled).toBe(false); + release.resolve(); + await expect(prompt).rejects.toMatchObject({ code: "invalid_state" }); + await shutdown; + expect(providerCalls).toBe(0); + }); + + it("aborts and awaits active compaction without persisting its result", async () => { + const registration = newFaux(); + const entered = deferred(); + const release = deferred(); + let signal: AbortSignal | undefined; + registration.setResponses([ + async (_context, options) => { + signal = options?.signal; + entered.resolve(); + await release.promise; + return fauxAssistantMessage("summary produced after shutdown"); + }, + ]); + const session = new Session(new InMemorySessionStorage()); + await session.appendMessage(createUserMessage("one")); + await session.appendMessage(createAssistantMessage("two")); + const harness = new AgentHarness({ models, session, model: registration.getModel() }); + const compaction = harness.compact(); + await entered.promise; + + let shutdownSettled = false; + const shutdown = harness.shutdown().then(() => { + shutdownSettled = true; + }); + await Promise.resolve(); + + expect(signal?.aborted).toBe(true); + expect(shutdownSettled).toBe(false); + release.resolve(); + await expect(compaction).rejects.toMatchObject({ code: "compaction" }); + await shutdown; + expect((await session.getEntries()).some((entry) => entry.type === "compaction")).toBe(false); + }); + + it("aborts and awaits active tree navigation without moving the session leaf", async () => { + const registration = newFaux(); + const entered = deferred(); + const release = deferred(); + let signal: AbortSignal | undefined; + registration.setResponses([ + async (_context, options) => { + signal = options?.signal; + entered.resolve(); + await release.promise; + return fauxAssistantMessage("summary produced after shutdown"); + }, + ]); + const session = new Session(new InMemorySessionStorage()); + const targetId = await session.appendMessage(createUserMessage("first branch")); + await session.appendMessage(createAssistantMessage("first reply")); + await session.appendMessage(createUserMessage("abandoned work")); + const originalLeafId = await session.appendMessage(createAssistantMessage("abandoned reply")); + const harness = new AgentHarness({ models, session, model: registration.getModel() }); + const navigation = harness.navigateTree(targetId, { summarize: true }); + await entered.promise; + + let shutdownSettled = false; + const shutdown = harness.shutdown().then(() => { + shutdownSettled = true; + }); + await Promise.resolve(); + + expect(signal?.aborted).toBe(true); + expect(shutdownSettled).toBe(false); + release.resolve(); + await expect(navigation).resolves.toEqual({ cancelled: true }); + await shutdown; + expect(await session.getLeafId()).toBe(originalLeafId); + }); + + it("shuts down an idle harness without modifying its durable session", async () => { + const session = new Session(new InMemorySessionStorage()); + await session.appendMessage(createUserMessage("existing")); + const harness = new AgentHarness({ + models, + session, + model: getModel("anthropic", "claude-sonnet-4-5"), + }); + + await harness.shutdown(); + + const messages = (await session.getEntries()).flatMap((entry) => + entry.type === "message" ? [entry.message] : [], + ); + expect(messages).toEqual([expect.objectContaining({ role: "user" })]); + await expect(harness.compact()).rejects.toMatchObject({ code: "invalid_state" }); + }); + it("drains one queued steering message at a time and emits queue updates", async () => { const registration = newFaux(); const userCounts: number[] = [];