refactor(agent): snapshot harness turn state

This commit is contained in:
Mario Zechner
2026-05-09 23:34:46 +02:00
parent 76131673d3
commit 322759a3f0
8 changed files with 587 additions and 305 deletions

View File

@@ -0,0 +1,136 @@
# AgentHarness lifecycle
`AgentHarness` is the orchestration layer above the low-level `Agent`. It owns session persistence, runtime configuration, resource resolution, operation locking, and extension-facing mutation semantics.
## State model
The harness separates state into four categories.
### Harness config
Harness config is the latest runtime configuration set by the application or extensions:
- model
- thinking level
- tools
- active tool names
- resources or resource provider
- system prompt or system prompt provider
Getters return harness config. They do not return the snapshot used by an in-flight provider request.
Setters update harness config immediately, including while a turn is in flight. Changes affect the next turn snapshot, not the currently running provider request.
### Turn snapshot
A turn snapshot is the concrete state used for one LLM turn. It is created by `createTurnState()` and contains:
- persisted session messages
- resolved resources
- resolved system prompt
- model
- thinking level
- active tools
Static option values are used directly. Provider callbacks are invoked once per `createTurnState()` call. All logic for that turn uses the same snapshot.
### Session
The session contains persisted entries only. Session reads return persisted state and do not include queued writes.
### Pending session writes
Session writes requested while an operation is active are queued as pending session writes. Pending writes are visible through an explicit pending-writes API, not through normal session reads.
Pending session writes are always persisted. They are flushed at save points, at operation settlement, and in failure cleanup.
## Operation phases
The harness has an explicit phase:
```ts
type AgentHarnessPhase = "idle" | "turn" | "compaction" | "branch_summary" | "retry";
```
Structural operations require `phase === "idle"` and synchronously set the phase before the first `await`:
- `prompt`
- `skill`
- `promptFromTemplate`
- `compact`
- `navigateTree`
Starting another structural operation while the harness is not idle throws.
The following operations are allowed during a turn where appropriate:
- `steer`
- `followUp`
- `nextTurn`
- `abort`
- runtime config setters
- session facade writes
## Turn execution
`prompt`, `skill`, and `promptFromTemplate` follow the same flow:
1. Assert idle and set phase to `"turn"`.
2. Create a turn snapshot with `createTurnState()`.
3. Derive invocation text from that snapshot.
4. Execute the turn with `executeTurn()`.
`skill` and `promptFromTemplate` resolve their resource from the same snapshot that is passed to the turn. They do not resolve resources separately.
`nextTurn` queues arbitrary `AgentMessage`s for the next user-initiated turn. Queued messages are inserted before the new user message.
## Save points
A save point occurs after an assistant turn and its tool-result messages have completed.
At a save point the harness:
1. flushes pending session writes after the agent-emitted messages for that turn
2. creates a fresh turn snapshot if the low-level loop may continue
3. applies the fresh context/model/reasoning state before the next provider request
This lets model, thinking level, tool, resource, and system prompt changes made during a turn affect the next turn in the same run, while never mutating an in-flight provider request. The loop callbacks are not recreated at save points.
No state refresh is needed on `agent_end` except flushing leftover pending session writes and clearing the operation phase.
## Session facade
Extensions and callbacks interact with a harness-scoped session facade rather than the raw session.
Reads delegate to persisted session state. Writes behave as follows:
- idle: persist immediately
- busy: enqueue as pending session writes
The facade exposes pending writes explicitly for diagnostics and UI:
```ts
getPendingWrites(): readonly PendingSessionWrite[]
```
Agent-emitted messages are persisted on `message_end` to preserve transcript ordering. Pending extension/session writes flush after those messages at save points.
## Extension context
Event payloads describe what is happening. Harness getters describe latest config for future snapshots.
Event contexts expose the harness and session facade. Events that belong to a turn also expose the immutable turn snapshot used for that turn. Extensions may update harness config at any time; updates affect the next snapshot.
## Abort
Abort is allowed during a turn. It aborts the low-level run and clears low-level steering/follow-up queues.
Abort does not discard pending session writes. Pending writes flush at the next save point if reached, at `agent_end`, or in operation failure cleanup.
## Compaction and tree navigation
Compaction and tree navigation are structural session mutations.
They are allowed only while idle and are not queued. They operate on persisted session state. The next prompt creates a fresh turn snapshot.
Branch summary generation is part of the tree navigation operation.

View File

@@ -153,13 +153,15 @@ function createAgentStream(): EventStream<AgentEvent, AgentMessage[]> {
* Main loop logic shared by agentLoop and agentLoopContinue.
*/
async function runLoop(
currentContext: AgentContext,
initialContext: AgentContext,
newMessages: AgentMessage[],
config: AgentLoopConfig,
initialConfig: AgentLoopConfig,
signal: AbortSignal | undefined,
emit: AgentEventSink,
streamFn?: StreamFn,
): Promise<void> {
let currentContext = initialContext;
let config = initialConfig;
let firstTurn = true;
// Check for steering messages at start (user may have typed while waiting)
let pendingMessages: AgentMessage[] = (await config.getSteeringMessages?.()) || [];
@@ -215,6 +217,27 @@ async function runLoop(
await emit({ type: "turn_end", message, toolResults });
const nextTurnContext = {
message,
toolResults,
context: currentContext,
newMessages,
};
const nextTurnSnapshot = await config.prepareNextTurn?.(nextTurnContext);
if (nextTurnSnapshot) {
currentContext = nextTurnSnapshot.context ?? currentContext;
config = {
...config,
model: nextTurnSnapshot.model ?? config.model,
reasoning:
nextTurnSnapshot.thinkingLevel === undefined
? config.reasoning
: nextTurnSnapshot.thinkingLevel === "off"
? undefined
: nextTurnSnapshot.thinkingLevel,
};
}
if (
await config.shouldStopAfterTurn?.({
message,

View File

@@ -15,6 +15,7 @@ import type {
AgentContext,
AgentEvent,
AgentLoopConfig,
AgentLoopTurnUpdate,
AgentMessage,
AgentState,
AgentTool,
@@ -52,7 +53,7 @@ const DEFAULT_MODEL = {
maxTokens: 0,
} satisfies Model<any>;
type QueueMode = "all" | "one-at-a-time";
export type QueueMode = "all" | "one-at-a-time";
type MutableAgentState = Omit<AgentState, "isStreaming" | "streamingMessage" | "pendingToolCalls" | "errorMessage"> & {
isStreaming: boolean;
@@ -101,6 +102,9 @@ export interface AgentOptions {
onResponse?: SimpleStreamOptions["onResponse"];
beforeToolCall?: (context: BeforeToolCallContext, signal?: AbortSignal) => Promise<BeforeToolCallResult | undefined>;
afterToolCall?: (context: AfterToolCallContext, signal?: AbortSignal) => Promise<AfterToolCallResult | undefined>;
prepareNextTurn?: (
signal?: AbortSignal,
) => Promise<AgentLoopTurnUpdate | undefined> | AgentLoopTurnUpdate | undefined;
steeringMode?: QueueMode;
followUpMode?: QueueMode;
sessionId?: string;
@@ -175,6 +179,9 @@ export class Agent {
context: AfterToolCallContext,
signal?: AbortSignal,
) => Promise<AfterToolCallResult | undefined>;
public prepareNextTurn?: (
signal?: AbortSignal,
) => Promise<AgentLoopTurnUpdate | undefined> | AgentLoopTurnUpdate | undefined;
private activeRun?: ActiveRun;
/** Session identifier forwarded to providers for cache-aware backends. */
public sessionId?: string;
@@ -197,6 +204,7 @@ export class Agent {
this.onResponse = options.onResponse;
this.beforeToolCall = options.beforeToolCall;
this.afterToolCall = options.afterToolCall;
this.prepareNextTurn = options.prepareNextTurn;
this.steeringQueue = new PendingMessageQueue(options.steeringMode ?? "one-at-a-time");
this.followUpQueue = new PendingMessageQueue(options.followUpMode ?? "one-at-a-time");
this.sessionId = options.sessionId;
@@ -421,6 +429,7 @@ export class Agent {
toolExecution: this.toolExecution,
beforeToolCall: this.beforeToolCall,
afterToolCall: this.afterToolCall,
prepareNextTurn: this.prepareNextTurn ? async () => await this.prepareNextTurn?.(this.signal) : undefined,
convertToLlm: this.convertToLlm,
transformContext: this.transformContext,
getApiKey: this.getApiKey,

View File

@@ -1,6 +1,5 @@
import { randomUUID } from "node:crypto";
import type { AssistantMessage, ImageContent, Model } from "@earendil-works/pi-ai";
import { Agent } from "../agent.js";
import type { AssistantMessage, ImageContent, Model, UserMessage } from "@earendil-works/pi-ai";
import { Agent, type QueueMode } from "../agent.js";
import type { AgentEvent, AgentMessage, AgentTool, ThinkingLevel } from "../types.js";
import { collectEntriesForBranchSummary, generateBranchSummary } from "./compaction/branch-summarization.js";
import { compact, DEFAULT_COMPACTION_SETTINGS, prepareCompaction } from "./compaction/compaction.js";
@@ -8,20 +7,20 @@ import { formatPromptTemplateInvocation } from "./prompt-templates.js";
import { formatSkillInvocation } from "./skills.js";
import type {
AbortResult,
AgentHarnessContext,
AgentHarnessConversationState,
AgentHarnessEvent,
AgentHarnessEventResultMap,
AgentHarnessOperationState,
AgentHarnessOptions,
AgentHarnessOwnEvent,
AgentHarnessPhase,
AgentHarnessResources,
AgentHarnessTurnState,
ExecutionEnv,
NavigateTreeResult,
PendingSessionWrite,
Session,
} from "./types.js";
function createUserMessage(text: string, images?: ImageContent[]): AgentMessage {
function createUserMessage(text: string, images?: ImageContent[]): UserMessage {
const content: Array<{ type: "text"; text: string } | ImageContent> = [{ type: "text", text }];
if (images) content.push(...images);
return { role: "user", content, timestamp: Date.now() };
@@ -30,19 +29,21 @@ function createUserMessage(text: string, images?: ImageContent[]): AgentMessage
export class AgentHarness {
readonly agent: Agent;
readonly env: ExecutionEnv;
readonly conversation: AgentHarnessConversationState;
readonly operation: AgentHarnessOperationState;
private session: Session;
private resourcesInput?: AgentHarnessOptions["resources"];
private model: Model<any>;
private thinkingLevel: ThinkingLevel;
private activeToolNames: string[];
private nextTurnQueue: AgentMessage[] = [];
private phase: AgentHarnessPhase = "idle";
private steerQueue: UserMessage[] = [];
private followUpQueue: UserMessage[] = [];
private pendingSessionWrites: PendingSessionWrite[] = [];
private resources?: AgentHarnessOptions["resources"];
private systemPrompt: AgentHarnessOptions["systemPrompt"];
private requestAuth?: AgentHarnessOptions["requestAuth"];
private toolRegistry = new Map<string, AgentTool>();
private getApiKeyAndHeaders?: AgentHarnessOptions["getApiKeyAndHeaders"];
private tools = new Map<string, AgentTool>();
private listeners = new Set<(event: AgentHarnessEvent, signal?: AbortSignal) => Promise<void> | void>();
private hooks = new Map<
keyof AgentHarnessEventResultMap,
Set<(event: any, ctx: AgentHarnessContext) => Promise<any> | any>
>();
private hooks = new Map<keyof AgentHarnessEventResultMap, Set<(event: any) => Promise<any> | any>>();
constructor(options: AgentHarnessOptions) {
this.agent = new Agent({
@@ -51,119 +52,79 @@ export class AgentHarness {
thinkingLevel: options.thinkingLevel,
tools: options.tools ?? [],
},
steeringMode: options.steeringMode,
followUpMode: options.followUpMode,
});
this.env = options.env;
this.session = options.session;
this.resourcesInput = options.resources;
this.resources = options.resources;
this.systemPrompt = options.systemPrompt;
this.requestAuth = options.requestAuth;
this.getApiKeyAndHeaders = options.getApiKeyAndHeaders;
for (const tool of this.agent.state.tools) {
this.toolRegistry.set(tool.name, tool);
this.tools.set(tool.name, tool);
}
this.conversation = {
session: options.session,
model: options.model,
thinkingLevel: options.thinkingLevel ?? this.agent.state.thinkingLevel,
activeToolNames: options.activeToolNames ?? this.agent.state.tools.map((tool) => tool.name),
nextTurnQueue: [],
};
this.agent.state.model = this.conversation.model;
this.agent.state.thinkingLevel = this.conversation.thinkingLevel;
this.model = options.model;
this.thinkingLevel = options.thinkingLevel ?? this.agent.state.thinkingLevel;
this.activeToolNames = options.activeToolNames ?? this.agent.state.tools.map((tool) => tool.name);
this.agent.state.model = this.model;
this.agent.state.thinkingLevel = this.thinkingLevel;
this.agent.getApiKey = async (provider) => {
const model = this.conversation.model;
if (!this.requestAuth || model.provider !== provider) return undefined;
return (await this.requestAuth(model))?.apiKey;
const model = this.model;
if (!this.getApiKeyAndHeaders || model.provider !== provider) return undefined;
return (await this.getApiKeyAndHeaders(model))?.apiKey;
};
this.operation = {
idle: true,
abortRequested: false,
steerQueue: [],
followUpQueue: [],
pendingMutations: {
appendMessages: [],
},
};
this.agent.transformContext = async (messages, signal) => {
const result = await this.emitHook("context", { type: "context", messages: [...messages] }, signal);
this.agent.transformContext = async (messages) => {
const result = await this.emitHook({ type: "context", messages: [...messages] });
return result?.messages ?? messages;
};
this.agent.beforeToolCall = async ({ toolCall, args }, signal) => {
const result = await this.emitHook(
"tool_call",
{
type: "tool_call",
toolCallId: toolCall.id,
toolName: toolCall.name,
input: args as Record<string, unknown>,
},
signal,
);
this.agent.beforeToolCall = async ({ toolCall, args }) => {
const result = await this.emitHook({
type: "tool_call",
toolCallId: toolCall.id,
toolName: toolCall.name,
input: args as Record<string, unknown>,
});
return result ? { block: result.block, reason: result.reason } : undefined;
};
this.agent.afterToolCall = async ({ toolCall, args, result, isError }, signal) => {
const patch = await this.emitHook(
"tool_result",
{
type: "tool_result",
toolCallId: toolCall.id,
toolName: toolCall.name,
input: args as Record<string, unknown>,
content: result.content,
details: result.details,
isError,
},
signal,
);
this.agent.afterToolCall = async ({ toolCall, args, result, isError }) => {
const patch = await this.emitHook({
type: "tool_result",
toolCallId: toolCall.id,
toolName: toolCall.name,
input: args as Record<string, unknown>,
content: result.content,
details: result.details,
isError,
});
return patch
? { content: patch.content, details: patch.details, isError: patch.isError, terminate: patch.terminate }
: undefined;
};
this.agent.onPayload = async (payload) => {
const result = await this.emitHook("before_provider_request", { type: "before_provider_request", payload });
const result = await this.emitHook({ type: "before_provider_request", payload });
return result?.payload ?? payload;
};
this.agent.onResponse = async (response) => {
const headers = { ...(response.headers as Record<string, string>) };
await this.emitOwn({ type: "after_provider_response", status: response.status, headers }, this.agent.signal);
};
this.agent.prepareNextTurn = async () => {
await this.flushPendingSessionWrites();
const turnState = await this.createTurnState();
this.applyTurnState(turnState);
return {
context: {
systemPrompt: turnState.systemPrompt,
messages: turnState.messages.slice(),
tools: turnState.activeTools.slice(),
},
model: turnState.model,
thinkingLevel: turnState.thinkingLevel,
};
};
this.agent.subscribe(async (event, signal) => {
await this.handleAgentEvent(event, signal);
});
void this.syncFromTree();
}
private createContext(signal?: AbortSignal): AgentHarnessContext {
return {
env: this.env,
conversation: this.conversation,
operation: this.operation,
abortSignal: signal,
};
}
private async resolveResources(signal?: AbortSignal): Promise<AgentHarnessResources> {
return typeof this.resourcesInput === "function"
? await this.resourcesInput(this.createContext(signal))
: (this.resourcesInput ?? {});
}
private getActiveTools(): AgentTool[] {
return this.conversation.activeToolNames
.map((name) => this.toolRegistry.get(name))
.filter((tool): tool is AgentTool => tool !== undefined);
}
private async resolveSystemPrompt(resources: AgentHarnessResources): Promise<string> {
if (!this.systemPrompt) return "You are a helpful assistant.";
if (typeof this.systemPrompt === "string") return this.systemPrompt;
return await this.systemPrompt({
env: this.env,
session: this.session,
model: this.conversation.model,
thinkingLevel: this.conversation.thinkingLevel,
activeTools: this.getActiveTools(),
resources,
});
}
private async emitOwn(event: AgentHarnessOwnEvent, signal?: AbortSignal): Promise<void> {
@@ -179,15 +140,13 @@ export class AgentHarness {
}
private async emitHook<TType extends keyof AgentHarnessEventResultMap>(
type: TType,
event: Extract<AgentHarnessOwnEvent, { type: TType }>,
signal?: AbortSignal,
): Promise<AgentHarnessEventResultMap[TType] | undefined> {
const handlers = this.hooks.get(type);
const handlers = this.hooks.get(event.type as TType);
if (!handlers || handlers.size === 0) return undefined;
let lastResult: AgentHarnessEventResultMap[TType] | undefined;
for (const handler of handlers) {
const result = await handler(event, this.createContext(signal));
const result = await handler(event);
if (result !== undefined) {
lastResult = result;
}
@@ -198,69 +157,89 @@ export class AgentHarness {
private async emitQueueUpdate(): Promise<void> {
await this.emitOwn({
type: "queue_update",
steer: [...this.operation.steerQueue],
followUp: [...this.operation.followUpQueue],
nextTurn: [...this.conversation.nextTurnQueue],
steer: [...this.steerQueue],
followUp: [...this.followUpQueue],
nextTurn: [...this.nextTurnQueue],
});
}
private async syncFromTree(): Promise<void> {
private async createTurnState(): Promise<AgentHarnessTurnState> {
const context = await this.session.buildContext();
this.agent.state.messages = context.messages;
if (context.model && this.conversation.model) {
// leave active model untouched; harness-level model is source of truth
const resources = typeof this.resources === "function" ? await this.resources() : (this.resources ?? {});
const tools = [...this.tools.values()];
const activeTools = this.activeToolNames
.map((name) => this.tools.get(name))
.filter((tool): tool is AgentTool => tool !== undefined);
let systemPrompt = "You are a helpful assistant.";
if (typeof this.systemPrompt === "string") {
systemPrompt = this.systemPrompt;
} else if (this.systemPrompt) {
systemPrompt = await this.systemPrompt({
env: this.env,
session: this.session,
model: this.model,
thinkingLevel: this.thinkingLevel,
activeTools,
resources,
});
}
this.agent.state.systemPrompt = await this.resolveSystemPrompt(await this.resolveResources());
return {
messages: context.messages,
resources,
systemPrompt,
model: this.model,
thinkingLevel: this.thinkingLevel,
tools,
activeTools,
};
}
private async applyPendingMutations(): Promise<void> {
for (const message of this.operation.pendingMutations.appendMessages) {
await this.session.appendMessage(message);
}
this.operation.pendingMutations.appendMessages = [];
private applyTurnState(turnState: AgentHarnessTurnState): void {
this.agent.state.messages = turnState.messages;
this.agent.state.systemPrompt = turnState.systemPrompt;
this.agent.state.model = turnState.model;
this.agent.state.thinkingLevel = turnState.thinkingLevel;
this.agent.state.tools = turnState.activeTools;
}
if (this.operation.pendingMutations.model) {
const model = this.operation.pendingMutations.model;
const previousModel = this.conversation.model;
this.conversation.model = model;
this.agent.state.model = model;
await this.session.appendModelChange(model.provider, model.id);
await this.emitOwn({ type: "model_select", model, previousModel, source: "set" });
this.operation.pendingMutations.model = undefined;
}
private validateToolNames(toolNames: string[]): void {
const missing = toolNames.filter((name) => !this.tools.has(name));
if (missing.length > 0) throw new Error(`Unknown tool(s): ${missing.join(", ")}`);
}
if (this.operation.pendingMutations.thinkingLevel !== undefined) {
const level = this.operation.pendingMutations.thinkingLevel;
const previousLevel = this.conversation.thinkingLevel;
this.conversation.thinkingLevel = level;
this.agent.state.thinkingLevel = level;
await this.session.appendThinkingLevelChange(level);
await this.emitOwn({ type: "thinking_level_select", level, previousLevel });
this.operation.pendingMutations.thinkingLevel = undefined;
private async flushPendingSessionWrites(): Promise<void> {
const writes = this.pendingSessionWrites;
this.pendingSessionWrites = [];
for (const write of writes) {
if (write.type === "message") {
await this.session.appendMessage(write.message);
} else if (write.type === "model_change") {
await this.session.appendModelChange(write.provider, write.modelId);
} else if (write.type === "thinking_level_change") {
await this.session.appendThinkingLevelChange(write.thinkingLevel);
} else if (write.type === "custom") {
await this.session.appendCustomEntry(write.customType, write.data);
} else if (write.type === "custom_message") {
await this.session.appendCustomMessageEntry(write.customType, write.content, write.display, write.details);
} else if (write.type === "label") {
await this.session.appendLabel(write.targetId, write.label);
} else if (write.type === "session_info") {
await this.session.appendSessionName(write.name ?? "");
}
}
if (this.operation.pendingMutations.activeToolNames) {
this.conversation.activeToolNames = [...this.operation.pendingMutations.activeToolNames];
this.agent.state.tools = this.conversation.activeToolNames
.map((name) => this.toolRegistry.get(name))
.filter((tool): tool is (typeof this.agent.state.tools)[number] => tool !== undefined);
this.operation.pendingMutations.activeToolNames = undefined;
}
await this.syncFromTree();
}
private async handleAgentEvent(event: AgentEvent, signal?: AbortSignal): Promise<void> {
await this.emitAny(event, signal);
if (event.type === "message_start") {
const steerIndex = this.operation.steerQueue.indexOf(event.message);
if (event.type === "message_start" && event.message.role === "user") {
const steerIndex = this.steerQueue.indexOf(event.message);
if (steerIndex !== -1) {
this.operation.steerQueue.splice(steerIndex, 1);
this.steerQueue.splice(steerIndex, 1);
await this.emitQueueUpdate();
} else {
const followUpIndex = this.operation.followUpQueue.indexOf(event.message);
const followUpIndex = this.followUpQueue.indexOf(event.message);
if (followUpIndex !== -1) {
this.operation.followUpQueue.splice(followUpIndex, 1);
this.followUpQueue.splice(followUpIndex, 1);
await this.emitQueueUpdate();
}
}
@@ -269,55 +248,47 @@ export class AgentHarness {
await this.session.appendMessage(event.message);
}
if (event.type === "turn_end") {
const hadPendingMutations =
this.operation.pendingMutations.appendMessages.length > 0 ||
this.operation.pendingMutations.model !== undefined ||
this.operation.pendingMutations.thinkingLevel !== undefined ||
this.operation.pendingMutations.activeToolNames !== undefined;
await this.emitOwn(
{ type: "save_point", liveOperationId: this.operation.liveOperationId ?? "unknown", hadPendingMutations },
signal,
);
if (hadPendingMutations) {
await this.applyPendingMutations();
}
const hadPendingMutations = this.pendingSessionWrites.length > 0;
await this.flushPendingSessionWrites();
await this.emitOwn({
type: "save_point",
hadPendingMutations,
});
}
if (event.type === "agent_end") {
this.operation.idle = true;
this.operation.liveOperationId = undefined;
this.operation.abortRequested = false;
await this.syncFromTree();
await this.emitOwn({ type: "settled", nextTurnCount: this.conversation.nextTurnQueue.length }, signal);
await this.flushPendingSessionWrites();
this.phase = "idle";
await this.emitOwn({ type: "settled", nextTurnCount: this.nextTurnQueue.length }, signal);
}
}
async prompt(text: string, options?: { images?: ImageContent[] }): Promise<AssistantMessage> {
if (!this.operation.idle) throw new Error("AgentHarness is busy");
private async executeTurn(
turnState: AgentHarnessTurnState,
text: string,
options?: { images?: ImageContent[] },
): Promise<AssistantMessage> {
this.applyTurnState(turnState);
const beforeLength = this.agent.state.messages.length;
this.operation.idle = false;
this.operation.liveOperationId = randomUUID();
const resources = await this.resolveResources(this.agent.signal);
let messages: AgentMessage[] = [createUserMessage(text, options?.images)];
if (this.conversation.nextTurnQueue.length > 0) {
messages = [messages[0]!, ...this.conversation.nextTurnQueue];
this.conversation.nextTurnQueue = [];
if (this.nextTurnQueue.length > 0) {
messages = [...this.nextTurnQueue, messages[0]!];
this.nextTurnQueue = [];
await this.emitQueueUpdate();
}
this.agent.state.systemPrompt = await this.resolveSystemPrompt(resources);
const beforeResult = await this.emitHook(
"before_agent_start",
{
type: "before_agent_start",
prompt: text,
images: options?.images,
systemPrompt: this.agent.state.systemPrompt,
resources,
},
this.agent.signal,
);
const beforeResult = await this.emitHook({
type: "before_agent_start",
prompt: text,
images: options?.images,
systemPrompt: turnState.systemPrompt,
resources: turnState.resources,
});
if (beforeResult?.messages) messages = [...beforeResult.messages, ...messages];
if (beforeResult?.systemPrompt) this.agent.state.systemPrompt = beforeResult.systemPrompt;
await this.agent.prompt(messages);
try {
await this.agent.prompt(messages);
} finally {
await this.flushPendingSessionWrites();
}
let response: AssistantMessage | undefined;
const newMessages = this.agent.state.messages.slice(beforeLength);
for (let i = newMessages.length - 1; i >= 0; i--) {
@@ -331,81 +302,98 @@ export class AgentHarness {
return response;
}
async prompt(text: string, options?: { images?: ImageContent[] }): Promise<AssistantMessage> {
if (this.phase !== "idle") throw new Error("AgentHarness is busy");
this.phase = "turn";
try {
const turnState = await this.createTurnState();
return await this.executeTurn(turnState, text, options);
} catch (error) {
this.phase = "idle";
throw error;
}
}
async skill(name: string, additionalInstructions?: string): Promise<AssistantMessage> {
const resources = await this.resolveResources();
const skill = (resources.skills ?? []).find((candidate) => candidate.name === name);
if (!skill) throw new Error(`Unknown skill: ${name}`);
return await this.prompt(formatSkillInvocation(skill, additionalInstructions));
if (this.phase !== "idle") throw new Error("AgentHarness is busy");
this.phase = "turn";
try {
const turnState = await this.createTurnState();
const skill = (turnState.resources.skills ?? []).find((candidate) => candidate.name === name);
if (!skill) throw new Error(`Unknown skill: ${name}`);
return await this.executeTurn(turnState, formatSkillInvocation(skill, additionalInstructions));
} catch (error) {
this.phase = "idle";
throw error;
}
}
async promptFromTemplate(name: string, args: string[] = []): Promise<AssistantMessage> {
const resources = await this.resolveResources();
const template = (resources.promptTemplates ?? []).find((candidate) => candidate.name === name);
if (!template) throw new Error(`Unknown prompt template: ${name}`);
return await this.prompt(formatPromptTemplateInvocation(template, args));
if (this.phase !== "idle") throw new Error("AgentHarness is busy");
this.phase = "turn";
try {
const turnState = await this.createTurnState();
const template = (turnState.resources.promptTemplates ?? []).find((candidate) => candidate.name === name);
if (!template) throw new Error(`Unknown prompt template: ${name}`);
return await this.executeTurn(turnState, formatPromptTemplateInvocation(template, args));
} catch (error) {
this.phase = "idle";
throw error;
}
}
steer(message: AgentMessage): void {
if (this.operation.idle) throw new Error("Cannot steer while idle");
this.operation.steerQueue.push(message);
steer(text: string, options?: { images?: ImageContent[] }): void {
if (this.phase === "idle") throw new Error("Cannot steer while idle");
const message = createUserMessage(text, options?.images);
this.steerQueue.push(message);
this.agent.steer(message);
void this.emitQueueUpdate();
}
followUp(message: AgentMessage): void {
if (this.operation.idle) throw new Error("Cannot follow up while idle");
this.operation.followUpQueue.push(message);
followUp(text: string, options?: { images?: ImageContent[] }): void {
if (this.phase === "idle") throw new Error("Cannot follow up while idle");
const message = createUserMessage(text, options?.images);
this.followUpQueue.push(message);
this.agent.followUp(message);
void this.emitQueueUpdate();
}
nextTurn(message: AgentMessage): void {
this.conversation.nextTurnQueue.push(message);
nextTurn(text: string, options?: { images?: ImageContent[] }): void {
this.nextTurnQueue.push(createUserMessage(text, options?.images));
void this.emitQueueUpdate();
}
async appendMessage(message: AgentMessage): Promise<void> {
if (this.operation.idle) {
if (this.phase === "idle") {
await this.session.appendMessage(message);
await this.syncFromTree();
} else {
this.operation.pendingMutations.appendMessages.push(message);
this.pendingSessionWrites.push({ type: "message", message });
}
}
async shell(
command: string,
options?: {
cwd?: string;
env?: Record<string, string>;
timeout?: number;
signal?: AbortSignal;
onStdout?: (chunk: string) => void;
onStderr?: (chunk: string) => void;
},
): Promise<{ stdout: string; stderr: string; exitCode: number }> {
return await this.env.exec(command, options);
}
async compact(
customInstructions?: string,
): Promise<{ summary: string; firstKeptEntryId: string; tokensBefore: number; details?: unknown }> {
if (!this.operation.idle) throw new Error("compact() requires idle harness");
const model = this.conversation.model;
if (this.phase !== "idle") throw new Error("compact() requires idle harness");
this.phase = "compaction";
const model = this.model;
if (!model) throw new Error("No model set for compaction");
const auth = await this.requestAuth?.(model);
const auth = await this.getApiKeyAndHeaders?.(model);
if (!auth) throw new Error("No auth available for compaction");
const branchEntries = await this.session.getBranch();
const preparation = prepareCompaction(branchEntries, DEFAULT_COMPACTION_SETTINGS);
if (!preparation) throw new Error("Nothing to compact");
const hookResult = await this.emitHook("session_before_compact", {
const hookResult = await this.emitHook({
type: "session_before_compact",
preparation,
branchEntries,
customInstructions,
signal: new AbortController().signal,
});
if (hookResult?.cancel) throw new Error("Compaction cancelled");
if (hookResult?.cancel) {
this.phase = "idle";
throw new Error("Compaction cancelled");
}
const provided = hookResult?.compaction;
const result =
provided ??
@@ -416,7 +404,7 @@ export class AgentHarness {
auth.headers,
customInstructions,
undefined,
this.conversation.thinkingLevel,
this.thinkingLevel,
));
const entryId = await this.session.appendCompaction(
result.summary,
@@ -426,10 +414,10 @@ export class AgentHarness {
provided !== undefined,
);
const entry = await this.session.getEntry(entryId);
await this.syncFromTree();
if (entry?.type === "compaction") {
await this.emitOwn({ type: "session_compact", compactionEntry: entry, fromHook: provided !== undefined });
}
this.phase = "idle";
return result;
}
@@ -437,9 +425,13 @@ export class AgentHarness {
targetId: string,
options?: { summarize?: boolean; customInstructions?: string; replaceInstructions?: boolean; label?: string },
): Promise<NavigateTreeResult> {
if (!this.operation.idle) throw new Error("navigateTree() requires idle harness");
if (this.phase !== "idle") throw new Error("navigateTree() requires idle harness");
this.phase = "branch_summary";
const oldLeafId = await this.session.getLeafId();
if (oldLeafId === targetId) return { cancelled: false };
if (oldLeafId === targetId) {
this.phase = "idle";
return { cancelled: false };
}
const targetEntry = await this.session.getEntry(targetId);
if (!targetEntry) throw new Error(`Entry ${targetId} not found`);
const { entries, commonAncestorId } = await collectEntriesForBranchSummary(this.session, oldLeafId, targetId);
@@ -454,19 +446,22 @@ export class AgentHarness {
label: options?.label,
};
const signal = new AbortController().signal;
const hookResult = await this.emitHook("session_before_tree", {
const hookResult = await this.emitHook({
type: "session_before_tree",
preparation,
signal,
});
if (hookResult?.cancel) return { cancelled: true };
if (hookResult?.cancel) {
this.phase = "idle";
return { cancelled: true };
}
let summaryEntry: any | undefined;
let summaryText: string | undefined = hookResult?.summary?.summary;
let summaryDetails: unknown = hookResult?.summary?.details;
if (!summaryText && options?.summarize && entries.length > 0) {
const model = this.conversation.model;
const model = this.model;
if (!model) throw new Error("No model set for branch summary");
const auth = await this.requestAuth?.(model);
const auth = await this.getApiKeyAndHeaders?.(model);
if (!auth) throw new Error("No auth available for branch summary");
const branchSummary = await generateBranchSummary(entries, {
model,
@@ -476,7 +471,10 @@ export class AgentHarness {
customInstructions: hookResult?.customInstructions ?? options?.customInstructions,
replaceInstructions: hookResult?.replaceInstructions ?? options?.replaceInstructions,
});
if (branchSummary.aborted) return { cancelled: true };
if (branchSummary.aborted) {
this.phase = "idle";
return { cancelled: true };
}
if (branchSummary.error) throw new Error(branchSummary.error);
summaryText = branchSummary.summary;
summaryDetails = {
@@ -521,7 +519,6 @@ export class AgentHarness {
if (summaryId) {
summaryEntry = await this.session.getEntry(summaryId);
}
await this.syncFromTree();
await this.emitOwn({
type: "session_tree",
newLeafId: await this.session.getLeafId(),
@@ -529,48 +526,80 @@ export class AgentHarness {
summaryEntry,
fromHook: hookResult?.summary !== undefined,
});
this.phase = "idle";
return { cancelled: false, editorText, summaryEntry };
}
async setModel(model: Model<any>): Promise<void> {
if (this.operation.idle) {
const previousModel = this.conversation.model;
this.conversation.model = model;
const previousModel = this.model;
this.model = model;
if (this.phase === "idle") {
this.agent.state.model = model;
await this.session.appendModelChange(model.provider, model.id);
await this.emitOwn({ type: "model_select", model, previousModel, source: "set" });
} else {
this.operation.pendingMutations.model = model;
this.pendingSessionWrites.push({ type: "model_change", provider: model.provider, modelId: model.id });
}
await this.emitOwn({ type: "model_select", model, previousModel, source: "set" });
}
async setThinkingLevel(level: ThinkingLevel): Promise<void> {
if (this.operation.idle) {
const previousLevel = this.conversation.thinkingLevel;
this.conversation.thinkingLevel = level;
const previousLevel = this.thinkingLevel;
this.thinkingLevel = level;
if (this.phase === "idle") {
this.agent.state.thinkingLevel = level;
await this.session.appendThinkingLevelChange(level);
await this.emitOwn({ type: "thinking_level_select", level, previousLevel });
} else {
this.operation.pendingMutations.thinkingLevel = level;
this.pendingSessionWrites.push({ type: "thinking_level_change", thinkingLevel: level });
}
await this.emitOwn({ type: "thinking_level_select", level, previousLevel });
}
async setActiveTools(toolNames: string[]): Promise<void> {
if (this.operation.idle) {
this.conversation.activeToolNames = [...toolNames];
this.agent.state.tools = this.getActiveTools();
this.validateToolNames(toolNames);
this.activeToolNames = [...toolNames];
if (this.phase === "idle") {
this.agent.state.tools = this.activeToolNames.map((name) => this.tools.get(name)!);
}
}
get steeringMode(): QueueMode {
return this.agent.steeringMode;
}
set steeringMode(mode: QueueMode) {
this.agent.steeringMode = mode;
}
get followUpMode(): QueueMode {
return this.agent.followUpMode;
}
set followUpMode(mode: QueueMode) {
this.agent.followUpMode = mode;
}
async setResources(resources: AgentHarnessResources): Promise<void> {
this.resources = resources;
}
async setTools(tools: AgentTool[], activeToolNames?: string[]): Promise<void> {
this.tools = new Map(tools.map((tool) => [tool.name, tool]));
if (activeToolNames) {
this.validateToolNames(activeToolNames);
this.activeToolNames = [...activeToolNames];
} else {
this.operation.pendingMutations.activeToolNames = [...toolNames];
this.validateToolNames(this.activeToolNames);
}
if (this.phase === "idle") {
this.agent.state.tools = this.activeToolNames.map((name) => this.tools.get(name)!);
}
}
async abort(): Promise<AbortResult> {
this.operation.abortRequested = true;
const clearedSteer = [...this.operation.steerQueue];
const clearedFollowUp = [...this.operation.followUpQueue];
this.operation.steerQueue = [];
this.operation.followUpQueue = [];
const clearedSteer = [...this.steerQueue];
const clearedFollowUp = [...this.followUpQueue];
this.steerQueue = [];
this.followUpQueue = [];
this.agent.clearAllQueues();
await this.emitQueueUpdate();
this.agent.abort();
@@ -591,8 +620,7 @@ export class AgentHarness {
on<TType extends keyof AgentHarnessEventResultMap>(
type: TType,
handler: (
event: Extract<import("./types.js").AgentHarnessOwnEvent, { type: TType }>,
ctx: AgentHarnessContext,
event: Extract<AgentHarnessOwnEvent, { type: TType }>,
) => Promise<AgentHarnessEventResultMap[TType]> | AgentHarnessEventResultMap[TType],
): () => void {
let handlers = this.hooks.get(type);

View File

@@ -1,4 +1,5 @@
import type { ImageContent, Model, TextContent } from "@earendil-works/pi-ai";
import type { QueueMode } from "../agent.js";
import type { AgentEvent, AgentMessage, AgentTool, ThinkingLevel } from "../index.js";
import type { Session } from "./session/session.js";
@@ -286,43 +287,22 @@ export interface JsonlSessionListOptions {
export interface JsonlSessionRepoApi
extends SessionRepo<JsonlSessionMetadata, JsonlSessionCreateOptions, JsonlSessionListOptions> {}
export interface AgentHarnessPendingMutations {
appendMessages: AgentMessage[];
model?: Model<any>;
thinkingLevel?: ThinkingLevel;
activeToolNames?: string[];
}
export type AgentHarnessPhase = "idle" | "turn" | "compaction" | "branch_summary" | "retry";
export interface AgentHarnessConversationState {
session: Session;
export type PendingSessionWrite = SessionTreeEntry extends infer TEntry
? TEntry extends SessionTreeEntry
? Omit<TEntry, "id" | "parentId" | "timestamp">
: never
: never;
export interface AgentHarnessTurnState {
messages: AgentMessage[];
resources: AgentHarnessResources;
systemPrompt: string;
model: Model<any>;
thinkingLevel: ThinkingLevel;
activeToolNames: string[];
nextTurnQueue: AgentMessage[];
}
export interface AgentHarnessOperationState {
idle: boolean;
liveOperationId?: string;
abortRequested: boolean;
steerQueue: AgentMessage[];
followUpQueue: AgentMessage[];
pendingMutations: AgentHarnessPendingMutations;
}
export interface SavePointSnapshot {
messages: AgentMessage[];
model: Model<any> | undefined;
thinkingLevel: ThinkingLevel;
activeToolNames: string[];
systemPrompt: string;
}
export interface AgentHarnessContext {
env: ExecutionEnv;
conversation: AgentHarnessConversationState;
operation: AgentHarnessOperationState;
abortSignal?: AbortSignal;
tools: AgentTool[];
activeTools: AgentTool[];
}
export interface QueueUpdateEvent {
@@ -334,7 +314,6 @@ export interface QueueUpdateEvent {
export interface SavePointEvent {
type: "save_point";
liveOperationId: string;
hadPendingMutations: boolean;
}
@@ -586,9 +565,7 @@ export interface AgentHarnessOptions {
env: ExecutionEnv;
session: Session;
tools?: AgentTool[];
resources?:
| AgentHarnessResources
| ((context: AgentHarnessContext) => AgentHarnessResources | Promise<AgentHarnessResources>);
resources?: AgentHarnessResources | (() => AgentHarnessResources | Promise<AgentHarnessResources>);
systemPrompt?:
| string
| ((context: {
@@ -599,10 +576,14 @@ export interface AgentHarnessOptions {
activeTools: AgentTool[];
resources: AgentHarnessResources;
}) => string | Promise<string>);
requestAuth?: (model: Model<any>) => Promise<{ apiKey: string; headers?: Record<string, string> } | undefined>;
getApiKeyAndHeaders?: (
model: Model<any>,
) => Promise<{ apiKey: string; headers?: Record<string, string> } | undefined>;
model: Model<any>;
thinkingLevel?: ThinkingLevel;
activeToolNames?: string[];
steeringMode?: QueueMode;
followUpMode?: QueueMode;
}
export type { AgentHarness } from "./agent-harness.js";

View File

@@ -112,6 +112,18 @@ export interface ShouldStopAfterTurnContext {
newMessages: AgentMessage[];
}
/** Replacement runtime state used by the agent loop before starting another provider request. */
export interface AgentLoopTurnUpdate {
/** Context for the next provider request. */
context?: AgentContext;
/** Model for the next provider request. */
model?: Model<any>;
/** Thinking level for the next provider request. */
thinkingLevel?: ThinkingLevel;
}
export interface PrepareNextTurnContext extends ShouldStopAfterTurnContext {}
export interface AgentLoopConfig extends SimpleStreamOptions {
model: Model<any>;
@@ -187,6 +199,15 @@ export interface AgentLoopConfig extends SimpleStreamOptions {
*/
shouldStopAfterTurn?: (context: ShouldStopAfterTurnContext) => boolean | Promise<boolean>;
/**
* Called after `turn_end` and before the loop decides whether another provider request should start.
* Return replacement context/model/thinking state to affect the next turn in this run.
* Return undefined to keep using the current context/config.
*/
prepareNextTurn?: (
context: PrepareNextTurnContext,
) => AgentLoopTurnUpdate | undefined | Promise<AgentLoopTurnUpdate | undefined>;
/**
* Returns steering messages to inject into the conversation mid-run.
*

View File

@@ -894,6 +894,79 @@ describe("agentLoop with AgentMessage", () => {
expect(parallelObserved).toBe(true);
});
it("should use prepareNextTurn snapshot before continuing", async () => {
const toolSchema = Type.Object({ value: Type.String() });
const tool: AgentTool<typeof toolSchema, { value: string }> = {
name: "echo",
label: "Echo",
description: "Echo tool",
parameters: toolSchema,
async execute(_toolCallId, params) {
return {
content: [{ type: "text", text: `echoed: ${params.value}` }],
details: { value: params.value },
};
},
};
const context: AgentContext = {
systemPrompt: "first prompt",
messages: [],
tools: [tool],
};
let convertedSecondTurnSystemPrompt = "";
let prepared = false;
const config: AgentLoopConfig = {
model: createModel(),
convertToLlm: identityConverter,
prepareNextTurn: async ({ context: currentContext }) => {
if (prepared) return undefined;
prepared = true;
return {
context: {
systemPrompt: "second prompt",
messages: currentContext.messages.slice(),
tools: currentContext.tools,
},
};
},
};
let llmCalls = 0;
const stream = agentLoop([createUserMessage("echo something")], context, config, undefined, (_model, ctx) => {
llmCalls++;
if (llmCalls === 2) {
convertedSecondTurnSystemPrompt = ctx.systemPrompt ?? "";
}
const mockStream = new MockAssistantStream();
queueMicrotask(() => {
if (llmCalls === 1) {
mockStream.push({
type: "done",
reason: "toolUse",
message: createAssistantMessage(
[{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "hello" } }],
"toolUse",
),
});
} else {
mockStream.push({
type: "done",
reason: "stop",
message: createAssistantMessage([{ type: "text", text: "done" }]),
});
}
});
return mockStream;
});
for await (const _event of stream) {
// consume
}
expect(llmCalls).toBe(2);
expect(convertedSecondTurnSystemPrompt).toBe("second prompt");
});
it("should stop after the current turn when shouldStopAfterTurn returns true", async () => {
const toolSchema = Type.Object({ value: Type.String() });
const executed: string[] = [];

View File

@@ -18,10 +18,21 @@ describe("harness factories", () => {
const session = createSession(new InMemorySessionStorage());
const env = new NodeExecutionEnv({ cwd: process.cwd() });
const initialModel = getModel("anthropic", "claude-sonnet-4-5");
const harness = createAgentHarness({ env, session, model: initialModel, systemPrompt: "You are helpful." });
const harness = createAgentHarness({
env,
session,
model: initialModel,
systemPrompt: "You are helpful.",
steeringMode: "all",
followUpMode: "all",
});
expect(harness.env).toBe(env);
expect(harness.conversation.session).toBe(session);
expect(harness.conversation.model).toBe(initialModel);
expect(harness.agent.state.model).toBe(initialModel);
expect(harness.steeringMode).toBe("all");
expect(harness.followUpMode).toBe("all");
harness.steeringMode = "one-at-a-time";
harness.followUpMode = "one-at-a-time";
expect(harness.agent.steeringMode).toBe("one-at-a-time");
expect(harness.agent.followUpMode).toBe("one-at-a-time");
});
});