shumengya/sproutclaw
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(agent): snapshot harness turn state

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

27
packages/agent/src/agent-loop.ts
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,

11
packages/agent/src/agent.ts
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,

548
packages/agent/src/harness/agent-harness.ts
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);

59
packages/agent/src/harness/types.ts
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";

21
packages/agent/src/types.ts
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.
*