From bc031ae4571081d2edae18787dd45486d0a17ce1 Mon Sep 17 00:00:00 2001 From: Christian Klotz Date: Thu, 30 Jul 2026 22:07:37 +0300 Subject: [PATCH] fix(agent): await mutations during harness shutdown --- packages/agent/src/harness/agent-harness.ts | 104 ++++++++++++------ .../agent/test/harness/agent-harness.test.ts | 46 ++++++++ 2 files changed, 119 insertions(+), 31 deletions(-) diff --git a/packages/agent/src/harness/agent-harness.ts b/packages/agent/src/harness/agent-harness.ts index f97d0210e..493c392d0 100644 --- a/packages/agent/src/harness/agent-harness.ts +++ b/packages/agent/src/harness/agent-harness.ts @@ -179,6 +179,7 @@ export class AgentHarness< private phase: AgentHarnessPhase = "idle"; private activeAbortController?: AbortController; private activeOperationPromise?: Promise; + private readonly activeMutationPromises = new Set>(); private shutdownPromise?: Promise; private isShutdown = false; private pendingSessionWrites: PendingSessionWrite[] = []; @@ -349,6 +350,30 @@ export class AgentHarness< }; } + private trackMutation(operation: () => Promise): Promise { + const promise = operation(); + this.activeMutationPromises.add(promise); + void promise.then( + () => this.activeMutationPromises.delete(promise), + () => this.activeMutationPromises.delete(promise), + ); + return promise; + } + + private async waitForMutations(): Promise { + while (this.activeMutationPromises.size > 0) { + const mutations = [...this.activeMutationPromises]; + await Promise.all( + mutations.map((mutation) => + mutation.then( + () => undefined, + () => undefined, + ), + ), + ); + } + } + private async resolveToolContext(): Promise { if (typeof this.toolContext === "function") { return await (this.toolContext as () => TContext | Promise)(); @@ -738,15 +763,17 @@ export class AgentHarness< async appendMessage(message: AgentMessage): Promise { this.assertNotShutDown(); - try { - if (this.phase === "idle") { - await this.session.appendMessage(message); - } else { - this.pendingSessionWrites.push({ type: "message", message }); + return this.trackMutation(async () => { + try { + if (this.phase === "idle") { + await this.session.appendMessage(message); + } else { + this.pendingSessionWrites.push({ type: "message", message }); + } + } catch (error) { + throw normalizeHarnessError(error, "session"); } - } catch (error) { - throw normalizeHarnessError(error, "session"); - } + }); } async compact(customInstructions?: string): Promise { @@ -914,18 +941,20 @@ export class AgentHarness< async setModel(model: Model): Promise { this.assertNotShutDown(); - try { - const previousModel = this.model; - if (this.phase === "idle") { - await this.session.appendModelChange(model.provider, model.id); - } else { - this.pendingSessionWrites.push({ type: "model_change", provider: model.provider, modelId: model.id }); + return this.trackMutation(async () => { + try { + const previousModel = this.model; + if (this.phase === "idle") { + await this.session.appendModelChange(model.provider, model.id); + } else { + this.pendingSessionWrites.push({ type: "model_change", provider: model.provider, modelId: model.id }); + } + this.model = model; + await this.emitOwn({ type: "model_update", model, previousModel, source: "set" }); + } catch (error) { + throw normalizeHarnessError(error, "session"); } - this.model = model; - await this.emitOwn({ type: "model_update", model, previousModel, source: "set" }); - } catch (error) { - throw normalizeHarnessError(error, "session"); - } + }); } getThinkingLevel(): ThinkingLevel { @@ -934,18 +963,20 @@ export class AgentHarness< async setThinkingLevel(level: ThinkingLevel): Promise { this.assertNotShutDown(); - try { - const previousLevel = this.thinkingLevel; - if (this.phase === "idle") { - await this.session.appendThinkingLevelChange(level); - } else { - this.pendingSessionWrites.push({ type: "thinking_level_change", thinkingLevel: level }); + return this.trackMutation(async () => { + try { + const previousLevel = this.thinkingLevel; + if (this.phase === "idle") { + await this.session.appendThinkingLevelChange(level); + } else { + this.pendingSessionWrites.push({ type: "thinking_level_change", thinkingLevel: level }); + } + this.thinkingLevel = level; + await this.emitOwn({ type: "thinking_level_update", level, previousLevel }); + } catch (error) { + throw normalizeHarnessError(error, "session"); } - this.thinkingLevel = level; - await this.emitOwn({ type: "thinking_level_update", level, previousLevel }); - } catch (error) { - throw normalizeHarnessError(error, "session"); - } + }); } getTools(): TTool[] { @@ -954,6 +985,10 @@ export class AgentHarness< async setTools(tools: TTool[], activeToolNames?: string[]): Promise { this.assertNotShutDown(); + return this.trackMutation(() => this.applyTools(tools, activeToolNames)); + } + + private async applyTools(tools: TTool[], activeToolNames?: string[]): Promise { try { this.validateUniqueNames( tools.map((tool) => tool.name), @@ -990,6 +1025,10 @@ export class AgentHarness< async setActiveTools(toolNames: string[]): Promise { this.assertNotShutDown(); + return this.trackMutation(() => this.applyActiveTools(toolNames)); + } + + private async applyActiveTools(toolNames: string[]): Promise { try { this.validateToolNames(toolNames); const previousToolNames = [...this.tools.keys()]; @@ -1069,7 +1108,10 @@ export class AgentHarness< this.followUpQueue = []; this.nextTurnQueue = []; this.activeAbortController?.abort(); - this.shutdownPromise = this.waitForIdle(); + this.shutdownPromise = (async () => { + await this.waitForIdle(); + await this.waitForMutations(); + })(); return this.shutdownPromise; } diff --git a/packages/agent/test/harness/agent-harness.test.ts b/packages/agent/test/harness/agent-harness.test.ts index d678427b0..954e5f2b8 100644 --- a/packages/agent/test/harness/agent-harness.test.ts +++ b/packages/agent/test/harness/agent-harness.test.ts @@ -89,6 +89,25 @@ function createAssistantMessage(text: string): AgentMessage { }; } +class BlockingSessionStorage extends InMemorySessionStorage { + readonly allWritesStarted = deferred(); + readonly releaseWrites = deferred(); + private readonly expectedWrites: number; + private writesStarted = 0; + + constructor(expectedWrites: number) { + super(); + this.expectedWrites = expectedWrites; + } + + override async appendEntry(entry: Parameters[0]): Promise { + this.writesStarted++; + if (this.writesStarted === this.expectedWrites) this.allWritesStarted.resolve(); + await this.releaseWrites.promise; + await super.appendEntry(entry); + } +} + describe("AgentHarness", () => { it("constructs directly and exposes queue modes", () => { const session = new Session(new InMemorySessionStorage()); @@ -265,6 +284,33 @@ describe("AgentHarness", () => { expect(await session.getLeafId()).toBe(originalLeafId); }); + it("awaits concurrent idle session mutations before shutdown resolves", async () => { + const storage = new BlockingSessionStorage(3); + const harness = new AgentHarness({ + models, + session: new Session(storage), + model: getModel("anthropic", "claude-sonnet-4-5"), + }); + const nextModel = getModel("anthropic", "claude-haiku-4-5"); + const mutations = [ + harness.appendMessage(createUserMessage("concurrent")), + harness.setModel(nextModel), + harness.setThinkingLevel("high"), + ]; + await storage.allWritesStarted.promise; + + const shutdown = harness.shutdown(); + const firstSettlement = await Promise.race([ + shutdown.then(() => "shutdown" as const), + new Promise<"writes-pending">((resolve) => setImmediate(() => resolve("writes-pending"))), + ]); + + expect(firstSettlement).toBe("writes-pending"); + storage.releaseWrites.resolve(); + await Promise.all([...mutations, shutdown]); + expect(await storage.getEntries()).toHaveLength(3); + }); + it("shuts down an idle harness without modifying its durable session", async () => { const session = new Session(new InMemorySessionStorage()); await session.appendMessage(createUserMessage("existing"));