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

Merge pull request #7343 from earendil-works/feat/agent-harness-shutdown

Browse source 
feat(agent): add harness shutdown lifecycle
This commit is contained in:
Christian Klotz 2026-07-30 21:35:01 +03:00 • committed by GitHub
commit 7d86e6233a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 263 additions and 38 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -177,8 +177,10 @@ export class AgentHarness<
private session: Session;
readonly models: Models;
private phase: AgentHarnessPhase = "idle";
private runAbortController?: AbortController;
private runPromise?: Promise<void>;
private activeAbortController?: AbortController;
private activeOperationPromise?: Promise<void>;
private shutdownPromise?: Promise<void>;
private isShutdown = false;
private pendingSessionWrites: PendingSessionWrite[] = [];
private model: Model<any>;
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<AgentHarnessHandler> | 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<void>((resolve) => {
this.activeAbortController = abortController;
this.activeOperationPromise = new Promise<void>((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<AgentHarnessTurnState<TContext, TSkill, TPromptTemplate, TTool>> {
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<TContext, TSkill, TPromptTemplate, TTool>,
text: string,
signal: AbortSignal,
options?: { images?: ImageContent[] },
): Promise<AssistantMessage> {
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<TContext, TSkill, TPromptTemplate, TTool>) => {
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<AssistantMessage> {
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<AssistantMessage> {
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<AssistantMessage> {
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<void> {
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<void> {
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<void> {
this.assertNotShutDown();
this.nextTurnQueue.push(createUserMessage(text, options?.images));
await this.emitQueueUpdate();
}
async appendMessage(message: AgentMessage): Promise<void> {
this.assertNotShutDown();
try {
if (this.phase === "idle") {
await this.session.appendMessage(message);
@ -734,8 +750,10 @@ export class AgentHarness<
}
async compact(customInstructions?: string): Promise<CompactResult> {
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<NavigateTreeResult> {
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<any>): Promise<void> {
this.assertNotShutDown();
try {
const previousModel = this.model;
if (this.phase === "idle") {
@ -905,6 +933,7 @@ export class AgentHarness<
}
async setThinkingLevel(level: ThinkingLevel): Promise<void> {
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<void> {
this.assertNotShutDown();
try {
this.validateUniqueNames(
tools.map((tool) => tool.name),
@ -959,6 +989,7 @@ export class AgentHarness<
}
async setActiveTools(toolNames: string[]): Promise<void> {
this.assertNotShutDown();
try {
this.validateToolNames(toolNames);
const previousToolNames = [...this.tools.keys()];
@ -987,6 +1018,7 @@ export class AgentHarness<
}
async setSteeringMode(mode: QueueMode): Promise<void> {
this.assertNotShutDown();
this.steeringQueueMode = mode;
}
@ -995,6 +1027,7 @@ export class AgentHarness<
}
async setFollowUpMode(mode: QueueMode): Promise<void> {
this.assertNotShutDown();
this.followUpQueueMode = mode;
}
@ -1006,6 +1039,7 @@ export class AgentHarness<
}
async setResources(resources: AgentHarnessResources<TSkill, TPromptTemplate>): Promise<void> {
this.assertNotShutDown();
const previousResources = this.getResources();
this.resources = {
skills: resources.skills?.slice(),
@ -1019,15 +1053,33 @@ export class AgentHarness<
}
async setStreamOptions(streamOptions: AgentHarnessStreamOptions): Promise<void> {
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<void> {
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<AbortResult> {
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<void> {
await this.runPromise;
await this.activeOperationPromise;
}
subscribe(
listener: (event: AgentHarnessEvent<TSkill, TPromptTemplate>, signal?: AbortSignal) => Promise<void> | 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<AgentHarnessOwnEvent, { type: TType }>,
) => Promise<AgentHarnessEventResultMap[TType]> | AgentHarnessEventResultMap[TType],
): () => void {
this.assertNotShutDown();
let handlers = this.handlers.get(type);
if (!handlers) {
handlers = new Set();

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

@ -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[] = [];