sscoble/pi-mono
1
0
Fork 0
mirror of https://github.com/badlogic/pi-mono.git synced 2026-11-05 23:37:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

fix(agent): await mutations during harness shutdown

Browse source
This commit is contained in:
Christian Klotz 2026-07-30 22:07:37 +03:00
parent 56eb685b69
commit bc031ae457
No known key found for this signature in database
2 changed files with 119 additions and 31 deletions
Download patch file Download diff file Expand all files Collapse all files

104
packages/agent/src/harness/agent-harness.ts
View file

@ -179,6 +179,7 @@ export class AgentHarness<
private phase: AgentHarnessPhase = "idle";
private activeAbortController?: AbortController;
private activeOperationPromise?: Promise<void>;
private readonly activeMutationPromises = new Set<Promise<void>>();
private shutdownPromise?: Promise<void>;
private isShutdown = false;
private pendingSessionWrites: PendingSessionWrite[] = [];
@ -349,6 +350,30 @@ export class AgentHarness<
};
}
private trackMutation(operation: () => Promise<void>): Promise<void> {
const promise = operation();
this.activeMutationPromises.add(promise);
void promise.then(
() => this.activeMutationPromises.delete(promise),
() => this.activeMutationPromises.delete(promise),
);
return promise;
}
private async waitForMutations(): Promise<void> {
while (this.activeMutationPromises.size > 0) {
const mutations = [...this.activeMutationPromises];
await Promise.all(
mutations.map((mutation) =>
mutation.then(
() => undefined,
() => undefined,
),
),
);
}
}
private async resolveToolContext(): Promise<TContext> {
if (typeof this.toolContext === "function") {
return await (this.toolContext as () => TContext | Promise<TContext>)();
@ -738,15 +763,17 @@ export class AgentHarness<
async appendMessage(message: AgentMessage): Promise<void> {
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<CompactResult> {
@ -914,18 +941,20 @@ export class AgentHarness<
async setModel(model: Model<any>): Promise<void> {
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<void> {
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<void> {
this.assertNotShutDown();
return this.trackMutation(() => this.applyTools(tools, activeToolNames));
}
private async applyTools(tools: TTool[], activeToolNames?: string[]): Promise<void> {
try {
this.validateUniqueNames(
tools.map((tool) => tool.name),
@ -990,6 +1025,10 @@ export class AgentHarness<
async setActiveTools(toolNames: string[]): Promise<void> {
this.assertNotShutDown();
return this.trackMutation(() => this.applyActiveTools(toolNames));
}
private async applyActiveTools(toolNames: string[]): Promise<void> {
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;
}

46
packages/agent/test/harness/agent-harness.test.ts
View file

@ -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<InMemorySessionStorage["appendEntry"]>[0]): Promise<void> {
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"));