feat(agent): add initial harness foundation
This commit is contained in:
Mario Zechner
613
packages/agent/src/harness/agent-harness.ts
Normal file
613
packages/agent/src/harness/agent-harness.ts
Normal file
@@ -0,0 +1,613 @@
|
||||
import { randomUUID } from "node:crypto";
|
||||
import { readFileSync } from "node:fs";
|
||||
import type { ImageContent, Model } from "@mariozechner/pi-ai";
|
||||
import type { Agent } from "../agent.js";
|
||||
import type { AgentEvent, AgentMessage, AgentTool, ThinkingLevel } from "../types.js";
|
||||
import {
|
||||
collectEntriesForBranchSummary,
|
||||
compact,
|
||||
DEFAULT_COMPACTION_SETTINGS,
|
||||
generateBranchSummary,
|
||||
prepareCompaction,
|
||||
} from "./compaction.js";
|
||||
import { expandPromptTemplate } from "./prompt-templates.js";
|
||||
import type {
|
||||
AbortResult,
|
||||
AgentHarness,
|
||||
AgentHarnessContext,
|
||||
AgentHarnessConversationState,
|
||||
AgentHarnessEvent,
|
||||
AgentHarnessEventResultMap,
|
||||
AgentHarnessOperationState,
|
||||
AgentHarnessOptions,
|
||||
AgentHarnessOwnEvent,
|
||||
ExecutionEnv,
|
||||
NavigateTreeResult,
|
||||
PromptTemplate,
|
||||
SessionBeforeCompactResult,
|
||||
SessionBeforeTreeResult,
|
||||
SessionTree,
|
||||
Skill,
|
||||
SystemPromptInputs,
|
||||
} from "./types.js";
|
||||
|
||||
function buildSystemPrompt(inputs: SystemPromptInputs, cwd: string): string {
|
||||
const parts: string[] = [];
|
||||
if (inputs.basePrompt) parts.push(inputs.basePrompt);
|
||||
if (inputs.appendPrompt) parts.push(inputs.appendPrompt);
|
||||
if (inputs.contextFiles && inputs.contextFiles.length > 0) {
|
||||
parts.push("# Project Context\n");
|
||||
for (const file of inputs.contextFiles) {
|
||||
parts.push(`## ${file.path}\n\n${file.content}`);
|
||||
}
|
||||
}
|
||||
if (inputs.skills && inputs.skills.length > 0) {
|
||||
parts.push("# Available Skills\n");
|
||||
for (const skill of inputs.skills.filter((s) => !s.disableModelInvocation)) {
|
||||
parts.push(`- ${skill.name}: ${skill.description}`);
|
||||
}
|
||||
}
|
||||
parts.push(`Current working directory: ${cwd}`);
|
||||
return parts.filter(Boolean).join("\n\n");
|
||||
}
|
||||
|
||||
function createUserMessage(text: string, images?: ImageContent[]): AgentMessage {
|
||||
const content: Array<{ type: "text"; text: string } | ImageContent> = [{ type: "text", text }];
|
||||
if (images) content.push(...images);
|
||||
return { role: "user", content, timestamp: Date.now() };
|
||||
}
|
||||
|
||||
export class DefaultAgentHarness implements AgentHarness {
|
||||
readonly agent: Agent;
|
||||
readonly env: ExecutionEnv;
|
||||
readonly conversation: AgentHarnessConversationState;
|
||||
readonly operation: AgentHarnessOperationState;
|
||||
|
||||
private sessionTree: SessionTree;
|
||||
private promptTemplates: PromptTemplate[];
|
||||
private skills: Skill[];
|
||||
private requestAuth?: AgentHarnessOptions["requestAuth"];
|
||||
private toolRegistry = 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>
|
||||
>();
|
||||
|
||||
constructor(options: AgentHarnessOptions) {
|
||||
this.agent = options.agent;
|
||||
this.env = options.env;
|
||||
this.sessionTree = options.sessionTree;
|
||||
this.promptTemplates = options.promptTemplates ?? [];
|
||||
this.skills = options.skills ?? [];
|
||||
this.requestAuth = options.requestAuth;
|
||||
for (const tool of this.agent.state.tools) {
|
||||
this.toolRegistry.set(tool.name, tool);
|
||||
}
|
||||
this.conversation = {
|
||||
sessionTree: options.sessionTree,
|
||||
model: options.initialModel ?? this.agent.state.model,
|
||||
thinkingLevel: options.initialThinkingLevel ?? this.agent.state.thinkingLevel,
|
||||
activeToolNames: options.initialActiveToolNames ?? this.agent.state.tools.map((tool) => tool.name),
|
||||
systemPromptInputs: options.initialSystemPromptInputs ?? { skills: this.skills },
|
||||
nextTurnQueue: [],
|
||||
};
|
||||
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);
|
||||
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,
|
||||
);
|
||||
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,
|
||||
);
|
||||
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 });
|
||||
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.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 emitOwn(event: AgentHarnessOwnEvent, signal?: AbortSignal): Promise<void> {
|
||||
for (const listener of this.listeners) {
|
||||
await listener(event, signal);
|
||||
}
|
||||
}
|
||||
|
||||
private async emitAny(event: AgentHarnessEvent, signal?: AbortSignal): Promise<void> {
|
||||
for (const listener of this.listeners) {
|
||||
await listener(event, signal);
|
||||
}
|
||||
}
|
||||
|
||||
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);
|
||||
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));
|
||||
if (result !== undefined) {
|
||||
lastResult = result;
|
||||
}
|
||||
}
|
||||
return lastResult;
|
||||
}
|
||||
|
||||
private async emitQueueUpdate(): Promise<void> {
|
||||
await this.emitOwn({
|
||||
type: "queue_update",
|
||||
steer: [...this.operation.steerQueue],
|
||||
followUp: [...this.operation.followUpQueue],
|
||||
nextTurn: [...this.conversation.nextTurnQueue],
|
||||
});
|
||||
}
|
||||
|
||||
private async syncFromTree(): Promise<void> {
|
||||
const context = await this.sessionTree.buildContext();
|
||||
this.agent.state.messages = context.messages;
|
||||
if (context.model && this.conversation.model) {
|
||||
// leave active model untouched; harness-level model is source of truth
|
||||
}
|
||||
this.agent.state.systemPrompt = buildSystemPrompt(this.conversation.systemPromptInputs, this.env.cwd);
|
||||
}
|
||||
|
||||
private async applyPendingMutations(): Promise<void> {
|
||||
for (const message of this.operation.pendingMutations.appendMessages) {
|
||||
await this.sessionTree.appendMessage(message);
|
||||
}
|
||||
this.operation.pendingMutations.appendMessages = [];
|
||||
|
||||
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.sessionTree.appendModelChange(model.provider, model.id);
|
||||
await this.emitOwn({ type: "model_select", model, previousModel, source: "set" });
|
||||
this.operation.pendingMutations.model = undefined;
|
||||
}
|
||||
|
||||
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.sessionTree.appendThinkingLevelChange(level);
|
||||
await this.emitOwn({ type: "thinking_level_select", level, previousLevel });
|
||||
this.operation.pendingMutations.thinkingLevel = undefined;
|
||||
}
|
||||
|
||||
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;
|
||||
}
|
||||
|
||||
if (this.operation.pendingMutations.systemPromptInputs) {
|
||||
this.conversation.systemPromptInputs = this.operation.pendingMutations.systemPromptInputs;
|
||||
this.agent.state.systemPrompt = buildSystemPrompt(this.conversation.systemPromptInputs, this.env.cwd);
|
||||
this.operation.pendingMutations.systemPromptInputs = 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 (steerIndex !== -1) {
|
||||
this.operation.steerQueue.splice(steerIndex, 1);
|
||||
await this.emitQueueUpdate();
|
||||
} else {
|
||||
const followUpIndex = this.operation.followUpQueue.indexOf(event.message);
|
||||
if (followUpIndex !== -1) {
|
||||
this.operation.followUpQueue.splice(followUpIndex, 1);
|
||||
await this.emitQueueUpdate();
|
||||
}
|
||||
}
|
||||
}
|
||||
if (event.type === "message_end") {
|
||||
await this.sessionTree.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 ||
|
||||
this.operation.pendingMutations.systemPromptInputs !== undefined;
|
||||
await this.emitOwn(
|
||||
{ type: "save_point", liveOperationId: this.operation.liveOperationId ?? "unknown", hadPendingMutations },
|
||||
signal,
|
||||
);
|
||||
if (hadPendingMutations) {
|
||||
await this.applyPendingMutations();
|
||||
}
|
||||
}
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
async prompt(text: string, options?: { images?: ImageContent[] }): Promise<void> {
|
||||
if (!this.operation.idle) throw new Error("AgentHarness is busy");
|
||||
this.operation.idle = false;
|
||||
this.operation.liveOperationId = randomUUID();
|
||||
const expanded = this.expandSkillCommand(expandPromptTemplate(text, this.promptTemplates));
|
||||
let messages: AgentMessage[] = [createUserMessage(expanded, options?.images)];
|
||||
if (this.conversation.nextTurnQueue.length > 0) {
|
||||
messages = [messages[0]!, ...this.conversation.nextTurnQueue];
|
||||
this.conversation.nextTurnQueue = [];
|
||||
await this.emitQueueUpdate();
|
||||
}
|
||||
this.agent.state.systemPrompt = buildSystemPrompt(this.conversation.systemPromptInputs, this.env.cwd);
|
||||
const beforeResult = await this.emitHook(
|
||||
"before_agent_start",
|
||||
{
|
||||
type: "before_agent_start",
|
||||
prompt: expanded,
|
||||
images: options?.images,
|
||||
systemPrompt: this.agent.state.systemPrompt,
|
||||
systemPromptInputs: this.conversation.systemPromptInputs,
|
||||
},
|
||||
this.agent.signal,
|
||||
);
|
||||
if (beforeResult?.messages) messages = [...beforeResult.messages, ...messages];
|
||||
if (beforeResult?.systemPrompt) this.agent.state.systemPrompt = beforeResult.systemPrompt;
|
||||
await this.agent.prompt(messages);
|
||||
}
|
||||
|
||||
async skill(name: string, args?: string): Promise<void> {
|
||||
const skill = this.skills.find((candidate) => candidate.name === name);
|
||||
if (!skill) throw new Error(`Unknown skill: ${name}`);
|
||||
let content = readFileSync(skill.filePath, "utf8");
|
||||
content = content.replace(/^---\s*\n[\s\S]*?\n---\s*\n?/, "").trim();
|
||||
const prompt = args ? `${content}\n\n${args}` : content;
|
||||
await this.prompt(prompt);
|
||||
}
|
||||
|
||||
steer(message: AgentMessage): void {
|
||||
if (this.operation.idle) throw new Error("Cannot steer while idle");
|
||||
this.operation.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);
|
||||
this.agent.followUp(message);
|
||||
void this.emitQueueUpdate();
|
||||
}
|
||||
|
||||
nextTurn(message: AgentMessage): void {
|
||||
this.conversation.nextTurnQueue.push(message);
|
||||
void this.emitQueueUpdate();
|
||||
}
|
||||
|
||||
async appendMessage(message: AgentMessage): Promise<void> {
|
||||
if (this.operation.idle) {
|
||||
await this.sessionTree.appendMessage(message);
|
||||
await this.syncFromTree();
|
||||
} else {
|
||||
this.operation.pendingMutations.appendMessages.push(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 (!model) throw new Error("No model set for compaction");
|
||||
const auth = await this.requestAuth?.(model);
|
||||
if (!auth) throw new Error("No auth available for compaction");
|
||||
const branchEntries = await this.sessionTree.getBranch();
|
||||
const preparation = prepareCompaction(branchEntries, DEFAULT_COMPACTION_SETTINGS);
|
||||
if (!preparation) throw new Error("Nothing to compact");
|
||||
const hookResult = await this.emitHook("session_before_compact", {
|
||||
type: "session_before_compact",
|
||||
preparation,
|
||||
branchEntries,
|
||||
customInstructions,
|
||||
signal: new AbortController().signal,
|
||||
});
|
||||
if ((hookResult as SessionBeforeCompactResult | undefined)?.cancel) throw new Error("Compaction cancelled");
|
||||
const provided = (hookResult as SessionBeforeCompactResult | undefined)?.compaction;
|
||||
const result =
|
||||
provided ??
|
||||
(await compact(
|
||||
preparation,
|
||||
model,
|
||||
auth.apiKey,
|
||||
auth.headers,
|
||||
customInstructions,
|
||||
undefined,
|
||||
this.conversation.thinkingLevel,
|
||||
));
|
||||
const entryId = await this.sessionTree.appendCompaction(
|
||||
result.summary,
|
||||
result.firstKeptEntryId,
|
||||
result.tokensBefore,
|
||||
result.details,
|
||||
provided !== undefined,
|
||||
);
|
||||
const entry = await this.sessionTree.getEntry(entryId);
|
||||
await this.syncFromTree();
|
||||
if (entry?.type === "compaction") {
|
||||
await this.emitOwn({ type: "session_compact", compactionEntry: entry, fromHook: provided !== undefined });
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
async navigateTree(
|
||||
targetId: string,
|
||||
options?: { summarize?: boolean; customInstructions?: string; replaceInstructions?: boolean; label?: string },
|
||||
): Promise<NavigateTreeResult> {
|
||||
if (!this.operation.idle) throw new Error("navigateTree() requires idle harness");
|
||||
const oldLeafId = await this.sessionTree.getLeafId();
|
||||
if (oldLeafId === targetId) return { cancelled: false };
|
||||
const targetEntry = await this.sessionTree.getEntry(targetId);
|
||||
if (!targetEntry) throw new Error(`Entry ${targetId} not found`);
|
||||
const { entries, commonAncestorId } = await collectEntriesForBranchSummary(this.sessionTree, oldLeafId, targetId);
|
||||
const preparation = {
|
||||
targetId,
|
||||
oldLeafId,
|
||||
commonAncestorId,
|
||||
entriesToSummarize: entries,
|
||||
userWantsSummary: options?.summarize ?? false,
|
||||
customInstructions: options?.customInstructions,
|
||||
replaceInstructions: options?.replaceInstructions,
|
||||
label: options?.label,
|
||||
};
|
||||
const signal = new AbortController().signal;
|
||||
const hookResult = await this.emitHook("session_before_tree", {
|
||||
type: "session_before_tree",
|
||||
preparation,
|
||||
signal,
|
||||
});
|
||||
const typedHook = hookResult as SessionBeforeTreeResult | undefined;
|
||||
if (typedHook?.cancel) return { cancelled: true };
|
||||
let summaryEntry: any | undefined;
|
||||
let summaryText: string | undefined = typedHook?.summary?.summary;
|
||||
let summaryDetails: unknown = typedHook?.summary?.details;
|
||||
if (!summaryText && options?.summarize && entries.length > 0) {
|
||||
const model = this.conversation.model;
|
||||
if (!model) throw new Error("No model set for branch summary");
|
||||
const auth = await this.requestAuth?.(model);
|
||||
if (!auth) throw new Error("No auth available for branch summary");
|
||||
const branchSummary = await generateBranchSummary(entries, {
|
||||
model,
|
||||
apiKey: auth.apiKey,
|
||||
headers: auth.headers,
|
||||
signal: new AbortController().signal,
|
||||
customInstructions: typedHook?.customInstructions ?? options?.customInstructions,
|
||||
replaceInstructions: typedHook?.replaceInstructions ?? options?.replaceInstructions,
|
||||
});
|
||||
if (branchSummary.aborted) return { cancelled: true };
|
||||
if (branchSummary.error) throw new Error(branchSummary.error);
|
||||
summaryText = branchSummary.summary;
|
||||
summaryDetails = {
|
||||
readFiles: branchSummary.readFiles ?? [],
|
||||
modifiedFiles: branchSummary.modifiedFiles ?? [],
|
||||
};
|
||||
}
|
||||
let editorText: string | undefined;
|
||||
let newLeafId: string | null;
|
||||
if (targetEntry.type === "message" && targetEntry.message.role === "user") {
|
||||
newLeafId = targetEntry.parentId;
|
||||
const content = targetEntry.message.content;
|
||||
editorText =
|
||||
typeof content === "string"
|
||||
? content
|
||||
: content
|
||||
.filter((c): c is { type: "text"; text: string } => c.type === "text")
|
||||
.map((c) => c.text)
|
||||
.join("");
|
||||
} else if (targetEntry.type === "custom_message") {
|
||||
newLeafId = targetEntry.parentId;
|
||||
editorText =
|
||||
typeof targetEntry.content === "string"
|
||||
? targetEntry.content
|
||||
: targetEntry.content
|
||||
.filter((c): c is { type: "text"; text: string } => c.type === "text")
|
||||
.map((c) => c.text)
|
||||
.join("");
|
||||
} else {
|
||||
newLeafId = targetId;
|
||||
}
|
||||
await this.sessionTree.moveTo(newLeafId);
|
||||
if (summaryText) {
|
||||
const summaryId = await this.sessionTree.appendBranchSummary(
|
||||
newLeafId ?? "root",
|
||||
summaryText,
|
||||
summaryDetails,
|
||||
typedHook?.summary !== undefined,
|
||||
);
|
||||
summaryEntry = await this.sessionTree.getEntry(summaryId);
|
||||
}
|
||||
await this.syncFromTree();
|
||||
await this.emitOwn({
|
||||
type: "session_tree",
|
||||
newLeafId: await this.sessionTree.getLeafId(),
|
||||
oldLeafId,
|
||||
summaryEntry,
|
||||
fromHook: typedHook?.summary !== undefined,
|
||||
});
|
||||
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;
|
||||
this.agent.state.model = model;
|
||||
await this.sessionTree.appendModelChange(model.provider, model.id);
|
||||
await this.emitOwn({ type: "model_select", model, previousModel, source: "set" });
|
||||
} else {
|
||||
this.operation.pendingMutations.model = model;
|
||||
}
|
||||
}
|
||||
|
||||
async setThinkingLevel(level: ThinkingLevel): Promise<void> {
|
||||
if (this.operation.idle) {
|
||||
const previousLevel = this.conversation.thinkingLevel;
|
||||
this.conversation.thinkingLevel = level;
|
||||
this.agent.state.thinkingLevel = level;
|
||||
await this.sessionTree.appendThinkingLevelChange(level);
|
||||
await this.emitOwn({ type: "thinking_level_select", level, previousLevel });
|
||||
} else {
|
||||
this.operation.pendingMutations.thinkingLevel = level;
|
||||
}
|
||||
}
|
||||
|
||||
async setActiveTools(toolNames: string[]): Promise<void> {
|
||||
if (this.operation.idle) {
|
||||
this.conversation.activeToolNames = [...toolNames];
|
||||
this.agent.state.tools = toolNames.map((name) => this.toolRegistry.get(name)).filter(Boolean) as any;
|
||||
} else {
|
||||
this.operation.pendingMutations.activeToolNames = [...toolNames];
|
||||
}
|
||||
}
|
||||
|
||||
async setSystemPromptInputs(inputs: SystemPromptInputs): Promise<void> {
|
||||
if (this.operation.idle) {
|
||||
this.conversation.systemPromptInputs = inputs;
|
||||
this.agent.state.systemPrompt = buildSystemPrompt(this.conversation.systemPromptInputs, this.env.cwd);
|
||||
} else {
|
||||
this.operation.pendingMutations.systemPromptInputs = inputs;
|
||||
}
|
||||
}
|
||||
|
||||
async abort(): Promise<AbortResult> {
|
||||
this.operation.abortRequested = true;
|
||||
const clearedSteer = [...this.operation.steerQueue];
|
||||
const clearedFollowUp = [...this.operation.followUpQueue];
|
||||
this.operation.steerQueue = [];
|
||||
this.operation.followUpQueue = [];
|
||||
this.agent.clearAllQueues();
|
||||
await this.emitQueueUpdate();
|
||||
this.agent.abort();
|
||||
await this.agent.waitForIdle();
|
||||
await this.emitOwn({ type: "abort", clearedSteer, clearedFollowUp });
|
||||
return { clearedSteer, clearedFollowUp };
|
||||
}
|
||||
|
||||
async waitForIdle(): Promise<void> {
|
||||
await this.agent.waitForIdle();
|
||||
}
|
||||
|
||||
subscribe(
|
||||
listener: (
|
||||
event: AgentEvent | import("./types.js").AgentHarnessOwnEvent,
|
||||
signal?: AbortSignal,
|
||||
) => Promise<void> | void,
|
||||
): () => void {
|
||||
this.listeners.add(listener);
|
||||
return () => this.listeners.delete(listener);
|
||||
}
|
||||
|
||||
on<TType extends keyof AgentHarnessEventResultMap>(
|
||||
type: TType,
|
||||
handler: (
|
||||
event: Extract<import("./types.js").AgentHarnessOwnEvent, { type: TType }>,
|
||||
ctx: AgentHarnessContext,
|
||||
) => Promise<AgentHarnessEventResultMap[TType]> | AgentHarnessEventResultMap[TType],
|
||||
): () => void {
|
||||
let handlers = this.hooks.get(type);
|
||||
if (!handlers) {
|
||||
handlers = new Set();
|
||||
this.hooks.set(type, handlers);
|
||||
}
|
||||
handlers.add(handler as any);
|
||||
return () => handlers!.delete(handler as any);
|
||||
}
|
||||
|
||||
private expandSkillCommand(text: string): string {
|
||||
if (!text.startsWith("/skill:")) return text;
|
||||
const spaceIndex = text.indexOf(" ");
|
||||
const skillName = spaceIndex === -1 ? text.slice(7) : text.slice(7, spaceIndex);
|
||||
const args = spaceIndex === -1 ? "" : text.slice(spaceIndex + 1).trim();
|
||||
const skill = this.skills.find((candidate) => candidate.name === skillName);
|
||||
if (!skill) return text;
|
||||
let content = readFileSync(skill.filePath, "utf8");
|
||||
content = content.replace(/^---\s*\n[\s\S]*?\n---\s*\n?/, "").trim();
|
||||
return args ? `${content}\n\n${args}` : content;
|
||||
}
|
||||
}
|
||||
733
packages/agent/src/harness/compaction.ts
Normal file
733
packages/agent/src/harness/compaction.ts
Normal file
@@ -0,0 +1,733 @@
|
||||
import { type AssistantMessage, completeSimple, type Model, type Usage } from "@mariozechner/pi-ai";
|
||||
import type { AgentMessage, ThinkingLevel } from "../types.js";
|
||||
import {
|
||||
convertToLlm,
|
||||
createBranchSummaryMessage,
|
||||
createCompactionSummaryMessage,
|
||||
createCustomMessage,
|
||||
} from "./messages.js";
|
||||
import { buildSessionContext } from "./session-tree.js";
|
||||
import type {
|
||||
BranchSummaryResult,
|
||||
CompactionEntry,
|
||||
CompactionPreparation,
|
||||
CompactionSettings,
|
||||
CompactResult,
|
||||
FileOperations,
|
||||
GenerateBranchSummaryOptions,
|
||||
SessionTreeEntry,
|
||||
} from "./types.js";
|
||||
|
||||
export const DEFAULT_COMPACTION_SETTINGS: CompactionSettings = {
|
||||
enabled: true,
|
||||
reserveTokens: 16384,
|
||||
keepRecentTokens: 20000,
|
||||
};
|
||||
|
||||
export interface CompactionDetails {
|
||||
readFiles: string[];
|
||||
modifiedFiles: string[];
|
||||
}
|
||||
|
||||
export interface BranchSummaryDetails {
|
||||
readFiles: string[];
|
||||
modifiedFiles: string[];
|
||||
}
|
||||
|
||||
export const SUMMARIZATION_SYSTEM_PROMPT = `You are a context summarization assistant. Your task is to read a conversation between a user and an AI coding assistant, then produce a structured summary following the exact format specified.
|
||||
|
||||
Do NOT continue the conversation. Do NOT respond to any questions in the conversation. ONLY output the structured summary.`;
|
||||
|
||||
const SUMMARIZATION_PROMPT = `The messages above are a conversation to summarize. Create a structured context checkpoint summary that another LLM will use to continue the work.
|
||||
|
||||
Use this EXACT format:
|
||||
|
||||
## Goal
|
||||
[What is the user trying to accomplish? Can be multiple items if the session covers different tasks.]
|
||||
|
||||
## Constraints & Preferences
|
||||
- [Any constraints, preferences, or requirements mentioned by user]
|
||||
- [Or "(none)" if none were mentioned]
|
||||
|
||||
## Progress
|
||||
### Done
|
||||
- [x] [Completed tasks/changes]
|
||||
|
||||
### In Progress
|
||||
- [ ] [Current work]
|
||||
|
||||
### Blocked
|
||||
- [Issues preventing progress, if any]
|
||||
|
||||
## Key Decisions
|
||||
- **[Decision]**: [Brief rationale]
|
||||
|
||||
## Next Steps
|
||||
1. [Ordered list of what should happen next]
|
||||
|
||||
## Critical Context
|
||||
- [Any data, examples, or references needed to continue]
|
||||
- [Or "(none)" if not applicable]
|
||||
|
||||
Keep each section concise. Preserve exact file paths, function names, and error messages.`;
|
||||
|
||||
const UPDATE_SUMMARIZATION_PROMPT = `The messages above are NEW conversation messages to incorporate into the existing summary provided in <previous-summary> tags.
|
||||
|
||||
Update the existing structured summary with new information. RULES:
|
||||
- PRESERVE all existing information from the previous summary
|
||||
- ADD new progress, decisions, and context from the new messages
|
||||
- UPDATE the Progress section: move items from "In Progress" to "Done" when completed
|
||||
- UPDATE "Next Steps" based on what was accomplished
|
||||
- PRESERVE exact file paths, function names, and error messages
|
||||
- If something is no longer relevant, you may remove it
|
||||
|
||||
Use this EXACT format:
|
||||
|
||||
## Goal
|
||||
[Preserve existing goals, add new ones if the task expanded]
|
||||
|
||||
## Constraints & Preferences
|
||||
- [Preserve existing, add new ones discovered]
|
||||
|
||||
## Progress
|
||||
### Done
|
||||
- [x] [Include previously done items AND newly completed items]
|
||||
|
||||
### In Progress
|
||||
- [ ] [Current work - update based on progress]
|
||||
|
||||
### Blocked
|
||||
- [Current blockers - remove if resolved]
|
||||
|
||||
## Key Decisions
|
||||
- **[Decision]**: [Brief rationale] (preserve all previous, add new)
|
||||
|
||||
## Next Steps
|
||||
1. [Update based on current state]
|
||||
|
||||
## Critical Context
|
||||
- [Preserve important context, add new if needed]
|
||||
|
||||
Keep each section concise. Preserve exact file paths, function names, and error messages.`;
|
||||
|
||||
const TURN_PREFIX_SUMMARIZATION_PROMPT = `This is the PREFIX of a turn that was too large to keep. The SUFFIX (recent work) is retained.
|
||||
|
||||
Summarize the prefix to provide context for the retained suffix:
|
||||
|
||||
## Original Request
|
||||
[What did the user ask for in this turn?]
|
||||
|
||||
## Early Progress
|
||||
- [Key decisions and work done in the prefix]
|
||||
|
||||
## Context for Suffix
|
||||
- [Information needed to understand the retained recent work]
|
||||
|
||||
Be concise. Focus on what's needed to understand the kept suffix.`;
|
||||
|
||||
const BRANCH_SUMMARY_PREAMBLE = `The user explored a different conversation branch before returning here.
|
||||
Summary of that exploration:
|
||||
|
||||
`;
|
||||
|
||||
const BRANCH_SUMMARY_PROMPT = `Create a structured summary of this conversation branch for context when returning later.
|
||||
|
||||
Use this EXACT format:
|
||||
|
||||
## Goal
|
||||
[What was the user trying to accomplish in this branch?]
|
||||
|
||||
## Constraints & Preferences
|
||||
- [Any constraints, preferences, or requirements mentioned]
|
||||
- [Or "(none)" if none were mentioned]
|
||||
|
||||
## Progress
|
||||
### Done
|
||||
- [x] [Completed tasks/changes]
|
||||
|
||||
### In Progress
|
||||
- [ ] [Work that was started but not finished]
|
||||
|
||||
### Blocked
|
||||
- [Issues preventing progress, if any]
|
||||
|
||||
## Key Decisions
|
||||
- **[Decision]**: [Brief rationale]
|
||||
|
||||
## Next Steps
|
||||
1. [What should happen next to continue this work]
|
||||
|
||||
Keep each section concise. Preserve exact file paths, function names, and error messages.`;
|
||||
|
||||
export function createFileOps(): FileOperations {
|
||||
return { read: new Set(), written: new Set(), edited: new Set() };
|
||||
}
|
||||
|
||||
export function extractFileOpsFromMessage(message: AgentMessage, fileOps: FileOperations): void {
|
||||
if (message.role !== "assistant") return;
|
||||
if (!("content" in message) || !Array.isArray(message.content)) return;
|
||||
for (const block of message.content) {
|
||||
if (typeof block !== "object" || block === null || !("type" in block) || block.type !== "toolCall") continue;
|
||||
if (!("arguments" in block) || !("name" in block)) continue;
|
||||
const args = block.arguments as Record<string, unknown> | undefined;
|
||||
if (!args) continue;
|
||||
const path = typeof args.path === "string" ? args.path : undefined;
|
||||
if (!path) continue;
|
||||
switch (block.name) {
|
||||
case "read":
|
||||
fileOps.read.add(path);
|
||||
break;
|
||||
case "write":
|
||||
fileOps.written.add(path);
|
||||
break;
|
||||
case "edit":
|
||||
fileOps.edited.add(path);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
export function computeFileLists(fileOps: FileOperations): { readFiles: string[]; modifiedFiles: string[] } {
|
||||
const modified = new Set([...fileOps.edited, ...fileOps.written]);
|
||||
const readOnly = [...fileOps.read].filter((f) => !modified.has(f)).sort();
|
||||
const modifiedFiles = [...modified].sort();
|
||||
return { readFiles: readOnly, modifiedFiles };
|
||||
}
|
||||
|
||||
export function formatFileOperations(readFiles: string[], modifiedFiles: string[]): string {
|
||||
const sections: string[] = [];
|
||||
if (readFiles.length > 0) sections.push(`<read-files>\n${readFiles.join("\n")}\n</read-files>`);
|
||||
if (modifiedFiles.length > 0) sections.push(`<modified-files>\n${modifiedFiles.join("\n")}\n</modified-files>`);
|
||||
if (sections.length === 0) return "";
|
||||
return `\n\n${sections.join("\n\n")}`;
|
||||
}
|
||||
|
||||
function truncateForSummary(text: string, maxChars: number): string {
|
||||
if (text.length <= maxChars) return text;
|
||||
const truncatedChars = text.length - maxChars;
|
||||
return `${text.slice(0, maxChars)}\n\n[... ${truncatedChars} more characters truncated]`;
|
||||
}
|
||||
|
||||
export function serializeConversation(messages: ReturnType<typeof convertToLlm>): string {
|
||||
const parts: string[] = [];
|
||||
for (const msg of messages) {
|
||||
if (msg.role === "user") {
|
||||
const content =
|
||||
typeof msg.content === "string"
|
||||
? msg.content
|
||||
: msg.content
|
||||
.filter((c): c is { type: "text"; text: string } => c.type === "text")
|
||||
.map((c) => c.text)
|
||||
.join("");
|
||||
if (content) parts.push(`[User]: ${content}`);
|
||||
} else if (msg.role === "assistant") {
|
||||
const textParts: string[] = [];
|
||||
const thinkingParts: string[] = [];
|
||||
const toolCalls: string[] = [];
|
||||
for (const block of msg.content) {
|
||||
if (block.type === "text") textParts.push(block.text);
|
||||
else if (block.type === "thinking") thinkingParts.push(block.thinking);
|
||||
else if (block.type === "toolCall") {
|
||||
const argsStr = Object.entries(block.arguments as Record<string, unknown>)
|
||||
.map(([k, v]) => `${k}=${JSON.stringify(v)}`)
|
||||
.join(", ");
|
||||
toolCalls.push(`${block.name}(${argsStr})`);
|
||||
}
|
||||
}
|
||||
if (thinkingParts.length > 0) parts.push(`[Assistant thinking]: ${thinkingParts.join("\n")}`);
|
||||
if (textParts.length > 0) parts.push(`[Assistant]: ${textParts.join("\n")}`);
|
||||
if (toolCalls.length > 0) parts.push(`[Assistant tool calls]: ${toolCalls.join("; ")}`);
|
||||
} else if (msg.role === "toolResult") {
|
||||
const content = msg.content
|
||||
.filter((c): c is { type: "text"; text: string } => c.type === "text")
|
||||
.map((c) => c.text)
|
||||
.join("");
|
||||
if (content) parts.push(`[Tool result]: ${truncateForSummary(content, 2000)}`);
|
||||
}
|
||||
}
|
||||
return parts.join("\n\n");
|
||||
}
|
||||
|
||||
export function calculateContextTokens(usage: Usage): number {
|
||||
return usage.totalTokens || usage.input + usage.output + usage.cacheRead + usage.cacheWrite;
|
||||
}
|
||||
|
||||
function getAssistantUsage(msg: AgentMessage): Usage | undefined {
|
||||
if (msg.role === "assistant" && "usage" in msg) {
|
||||
const assistantMsg = msg as AssistantMessage;
|
||||
if (assistantMsg.stopReason !== "aborted" && assistantMsg.stopReason !== "error" && assistantMsg.usage) {
|
||||
return assistantMsg.usage;
|
||||
}
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
export function estimateTokens(message: AgentMessage): number {
|
||||
let chars = 0;
|
||||
switch (message.role) {
|
||||
case "user": {
|
||||
const content = (message as { content: string | Array<{ type: string; text?: string }> }).content;
|
||||
if (typeof content === "string") chars = content.length;
|
||||
else if (Array.isArray(content)) {
|
||||
for (const block of content) if (block.type === "text" && block.text) chars += block.text.length;
|
||||
}
|
||||
return Math.ceil(chars / 4);
|
||||
}
|
||||
case "assistant": {
|
||||
const assistant = message as AssistantMessage;
|
||||
for (const block of assistant.content) {
|
||||
if (block.type === "text") chars += block.text.length;
|
||||
else if (block.type === "thinking") chars += block.thinking.length;
|
||||
else if (block.type === "toolCall") chars += block.name.length + JSON.stringify(block.arguments).length;
|
||||
}
|
||||
return Math.ceil(chars / 4);
|
||||
}
|
||||
case "custom":
|
||||
case "toolResult": {
|
||||
if (typeof message.content === "string") chars = message.content.length;
|
||||
else {
|
||||
for (const block of message.content) {
|
||||
if (block.type === "text" && block.text) chars += block.text.length;
|
||||
if (block.type === "image") chars += 4800;
|
||||
}
|
||||
}
|
||||
return Math.ceil(chars / 4);
|
||||
}
|
||||
case "bashExecution":
|
||||
chars = message.command.length + message.output.length;
|
||||
return Math.ceil(chars / 4);
|
||||
case "branchSummary":
|
||||
case "compactionSummary":
|
||||
chars = message.summary.length;
|
||||
return Math.ceil(chars / 4);
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
function getLastAssistantUsageInfo(messages: AgentMessage[]): { usage: Usage; index: number } | undefined {
|
||||
for (let i = messages.length - 1; i >= 0; i--) {
|
||||
const usage = getAssistantUsage(messages[i]!);
|
||||
if (usage) return { usage, index: i };
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
export function estimateContextTokens(messages: AgentMessage[]): {
|
||||
tokens: number;
|
||||
usageTokens: number;
|
||||
trailingTokens: number;
|
||||
lastUsageIndex: number | null;
|
||||
} {
|
||||
const usageInfo = getLastAssistantUsageInfo(messages);
|
||||
if (!usageInfo) {
|
||||
let estimated = 0;
|
||||
for (const message of messages) estimated += estimateTokens(message);
|
||||
return { tokens: estimated, usageTokens: 0, trailingTokens: estimated, lastUsageIndex: null };
|
||||
}
|
||||
const usageTokens = calculateContextTokens(usageInfo.usage);
|
||||
let trailingTokens = 0;
|
||||
for (let i = usageInfo.index + 1; i < messages.length; i++) trailingTokens += estimateTokens(messages[i]!);
|
||||
return { tokens: usageTokens + trailingTokens, usageTokens, trailingTokens, lastUsageIndex: usageInfo.index };
|
||||
}
|
||||
|
||||
export function shouldCompact(contextTokens: number, contextWindow: number, settings: CompactionSettings): boolean {
|
||||
if (!settings.enabled) return false;
|
||||
return contextTokens > contextWindow - settings.reserveTokens;
|
||||
}
|
||||
|
||||
function findValidCutPoints(entries: SessionTreeEntry[], startIndex: number, endIndex: number): number[] {
|
||||
const cutPoints: number[] = [];
|
||||
for (let i = startIndex; i < endIndex; i++) {
|
||||
const entry = entries[i]!;
|
||||
switch (entry.type) {
|
||||
case "message": {
|
||||
switch (entry.message.role) {
|
||||
case "bashExecution":
|
||||
case "custom":
|
||||
case "branchSummary":
|
||||
case "compactionSummary":
|
||||
case "user":
|
||||
case "assistant":
|
||||
cutPoints.push(i);
|
||||
break;
|
||||
case "toolResult":
|
||||
break;
|
||||
}
|
||||
break;
|
||||
}
|
||||
case "branch_summary":
|
||||
case "custom_message":
|
||||
cutPoints.push(i);
|
||||
break;
|
||||
default:
|
||||
break;
|
||||
}
|
||||
}
|
||||
return cutPoints;
|
||||
}
|
||||
|
||||
export function findTurnStartIndex(entries: SessionTreeEntry[], entryIndex: number, startIndex: number): number {
|
||||
for (let i = entryIndex; i >= startIndex; i--) {
|
||||
const entry = entries[i]!;
|
||||
if (entry.type === "branch_summary" || entry.type === "custom_message") return i;
|
||||
if (entry.type === "message") {
|
||||
const role = entry.message.role;
|
||||
if (role === "user" || role === "bashExecution") return i;
|
||||
}
|
||||
}
|
||||
return -1;
|
||||
}
|
||||
|
||||
export function findCutPoint(
|
||||
entries: SessionTreeEntry[],
|
||||
startIndex: number,
|
||||
endIndex: number,
|
||||
keepRecentTokens: number,
|
||||
) {
|
||||
const cutPoints = findValidCutPoints(entries, startIndex, endIndex);
|
||||
if (cutPoints.length === 0) return { firstKeptEntryIndex: startIndex, turnStartIndex: -1, isSplitTurn: false };
|
||||
let accumulatedTokens = 0;
|
||||
let cutIndex = cutPoints[0]!;
|
||||
for (let i = endIndex - 1; i >= startIndex; i--) {
|
||||
const entry = entries[i]!;
|
||||
if (entry.type !== "message") continue;
|
||||
accumulatedTokens += estimateTokens(entry.message);
|
||||
if (accumulatedTokens >= keepRecentTokens) {
|
||||
for (const cp of cutPoints) {
|
||||
if (cp >= i) {
|
||||
cutIndex = cp;
|
||||
break;
|
||||
}
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
while (cutIndex > startIndex) {
|
||||
const prevEntry = entries[cutIndex - 1]!;
|
||||
if (prevEntry.type === "compaction" || prevEntry.type === "message") break;
|
||||
cutIndex--;
|
||||
}
|
||||
const cutEntry = entries[cutIndex]!;
|
||||
const isUserMessage = cutEntry.type === "message" && cutEntry.message.role === "user";
|
||||
const turnStartIndex = isUserMessage ? -1 : findTurnStartIndex(entries, cutIndex, startIndex);
|
||||
return { firstKeptEntryIndex: cutIndex, turnStartIndex, isSplitTurn: !isUserMessage && turnStartIndex !== -1 };
|
||||
}
|
||||
|
||||
function getMessageFromEntry(entry: SessionTreeEntry): AgentMessage | undefined {
|
||||
switch (entry.type) {
|
||||
case "message":
|
||||
return entry.message.role === "toolResult" ? undefined : entry.message;
|
||||
case "custom_message":
|
||||
return createCustomMessage(entry.customType, entry.content, entry.display, entry.details, entry.timestamp);
|
||||
case "branch_summary":
|
||||
return createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp);
|
||||
case "compaction":
|
||||
return createCompactionSummaryMessage(entry.summary, entry.tokensBefore, entry.timestamp);
|
||||
default:
|
||||
return undefined;
|
||||
}
|
||||
}
|
||||
|
||||
function getMessageFromEntryForCompaction(entry: SessionTreeEntry): AgentMessage | undefined {
|
||||
if (entry.type === "compaction") return undefined;
|
||||
return getMessageFromEntry(entry);
|
||||
}
|
||||
|
||||
export async function generateSummary(
|
||||
currentMessages: AgentMessage[],
|
||||
model: Model<any>,
|
||||
reserveTokens: number,
|
||||
apiKey: string,
|
||||
headers?: Record<string, string>,
|
||||
signal?: AbortSignal,
|
||||
customInstructions?: string,
|
||||
previousSummary?: string,
|
||||
thinkingLevel?: ThinkingLevel,
|
||||
): Promise<string> {
|
||||
const maxTokens = Math.floor(0.8 * reserveTokens);
|
||||
let basePrompt = previousSummary ? UPDATE_SUMMARIZATION_PROMPT : SUMMARIZATION_PROMPT;
|
||||
if (customInstructions) basePrompt = `${basePrompt}\n\nAdditional focus: ${customInstructions}`;
|
||||
const llmMessages = convertToLlm(currentMessages);
|
||||
const conversationText = serializeConversation(llmMessages);
|
||||
let promptText = `<conversation>\n${conversationText}\n</conversation>\n\n`;
|
||||
if (previousSummary) promptText += `<previous-summary>\n${previousSummary}\n</previous-summary>\n\n`;
|
||||
promptText += basePrompt;
|
||||
const summarizationMessages = [
|
||||
{ role: "user" as const, content: [{ type: "text" as const, text: promptText }], timestamp: Date.now() },
|
||||
];
|
||||
const completionOptions =
|
||||
model.reasoning && thinkingLevel && thinkingLevel !== "off"
|
||||
? { maxTokens, signal, apiKey, headers, reasoning: thinkingLevel }
|
||||
: { maxTokens, signal, apiKey, headers };
|
||||
const response = await completeSimple(
|
||||
model,
|
||||
{ systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages },
|
||||
completionOptions,
|
||||
);
|
||||
if (response.stopReason === "error")
|
||||
throw new Error(`Summarization failed: ${response.errorMessage || "Unknown error"}`);
|
||||
return response.content
|
||||
.filter((c): c is { type: "text"; text: string } => c.type === "text")
|
||||
.map((c) => c.text)
|
||||
.join("\n");
|
||||
}
|
||||
|
||||
async function generateTurnPrefixSummary(
|
||||
messages: AgentMessage[],
|
||||
model: Model<any>,
|
||||
reserveTokens: number,
|
||||
apiKey: string,
|
||||
headers?: Record<string, string>,
|
||||
signal?: AbortSignal,
|
||||
thinkingLevel?: ThinkingLevel,
|
||||
): Promise<string> {
|
||||
const maxTokens = Math.floor(0.5 * reserveTokens);
|
||||
const llmMessages = convertToLlm(messages);
|
||||
const conversationText = serializeConversation(llmMessages);
|
||||
const promptText = `<conversation>\n${conversationText}\n</conversation>\n\n${TURN_PREFIX_SUMMARIZATION_PROMPT}`;
|
||||
const summarizationMessages = [
|
||||
{ role: "user" as const, content: [{ type: "text" as const, text: promptText }], timestamp: Date.now() },
|
||||
];
|
||||
const completionOptions =
|
||||
model.reasoning && thinkingLevel && thinkingLevel !== "off"
|
||||
? { maxTokens, signal, apiKey, headers, reasoning: thinkingLevel }
|
||||
: { maxTokens, signal, apiKey, headers };
|
||||
const response = await completeSimple(
|
||||
model,
|
||||
{ systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages },
|
||||
completionOptions,
|
||||
);
|
||||
if (response.stopReason === "error")
|
||||
throw new Error(`Turn prefix summarization failed: ${response.errorMessage || "Unknown error"}`);
|
||||
return response.content
|
||||
.filter((c): c is { type: "text"; text: string } => c.type === "text")
|
||||
.map((c) => c.text)
|
||||
.join("\n");
|
||||
}
|
||||
|
||||
export function prepareCompaction(
|
||||
pathEntries: SessionTreeEntry[],
|
||||
settings: CompactionSettings,
|
||||
): CompactionPreparation | undefined {
|
||||
if (pathEntries.length > 0 && pathEntries[pathEntries.length - 1]!.type === "compaction") return undefined;
|
||||
let prevCompactionIndex = -1;
|
||||
for (let i = pathEntries.length - 1; i >= 0; i--) {
|
||||
if (pathEntries[i]!.type === "compaction") {
|
||||
prevCompactionIndex = i;
|
||||
break;
|
||||
}
|
||||
}
|
||||
let previousSummary: string | undefined;
|
||||
let boundaryStart = 0;
|
||||
if (prevCompactionIndex >= 0) {
|
||||
const prevCompaction = pathEntries[prevCompactionIndex] as CompactionEntry;
|
||||
previousSummary = prevCompaction.summary;
|
||||
const firstKeptEntryIndex = pathEntries.findIndex((entry) => entry.id === prevCompaction.firstKeptEntryId);
|
||||
boundaryStart = firstKeptEntryIndex >= 0 ? firstKeptEntryIndex : prevCompactionIndex + 1;
|
||||
}
|
||||
const boundaryEnd = pathEntries.length;
|
||||
const tokensBefore = estimateContextTokens(buildSessionContext(pathEntries).messages).tokens;
|
||||
const cutPoint = findCutPoint(pathEntries, boundaryStart, boundaryEnd, settings.keepRecentTokens);
|
||||
const firstKeptEntry = pathEntries[cutPoint.firstKeptEntryIndex];
|
||||
if (!firstKeptEntry?.id) return undefined;
|
||||
const firstKeptEntryId = firstKeptEntry.id;
|
||||
const historyEnd = cutPoint.isSplitTurn ? cutPoint.turnStartIndex : cutPoint.firstKeptEntryIndex;
|
||||
const messagesToSummarize: AgentMessage[] = [];
|
||||
for (let i = boundaryStart; i < historyEnd; i++) {
|
||||
const msg = getMessageFromEntryForCompaction(pathEntries[i]!);
|
||||
if (msg) messagesToSummarize.push(msg);
|
||||
}
|
||||
const turnPrefixMessages: AgentMessage[] = [];
|
||||
if (cutPoint.isSplitTurn) {
|
||||
for (let i = cutPoint.turnStartIndex; i < cutPoint.firstKeptEntryIndex; i++) {
|
||||
const msg = getMessageFromEntryForCompaction(pathEntries[i]!);
|
||||
if (msg) turnPrefixMessages.push(msg);
|
||||
}
|
||||
}
|
||||
const fileOps = createFileOps();
|
||||
if (prevCompactionIndex >= 0) {
|
||||
const prevCompaction = pathEntries[prevCompactionIndex] as CompactionEntry;
|
||||
if (!prevCompaction.fromHook && prevCompaction.details) {
|
||||
const details = prevCompaction.details as CompactionDetails;
|
||||
if (Array.isArray(details.readFiles)) for (const f of details.readFiles) fileOps.read.add(f);
|
||||
if (Array.isArray(details.modifiedFiles)) for (const f of details.modifiedFiles) fileOps.edited.add(f);
|
||||
}
|
||||
}
|
||||
for (const msg of messagesToSummarize) extractFileOpsFromMessage(msg, fileOps);
|
||||
if (cutPoint.isSplitTurn) for (const msg of turnPrefixMessages) extractFileOpsFromMessage(msg, fileOps);
|
||||
return {
|
||||
firstKeptEntryId,
|
||||
messagesToSummarize,
|
||||
turnPrefixMessages,
|
||||
isSplitTurn: cutPoint.isSplitTurn,
|
||||
tokensBefore,
|
||||
previousSummary,
|
||||
fileOps,
|
||||
settings,
|
||||
};
|
||||
}
|
||||
|
||||
export async function compact(
|
||||
preparation: CompactionPreparation,
|
||||
model: Model<any>,
|
||||
apiKey: string,
|
||||
headers?: Record<string, string>,
|
||||
customInstructions?: string,
|
||||
signal?: AbortSignal,
|
||||
thinkingLevel?: ThinkingLevel,
|
||||
): Promise<CompactResult> {
|
||||
const {
|
||||
firstKeptEntryId,
|
||||
messagesToSummarize,
|
||||
turnPrefixMessages,
|
||||
isSplitTurn,
|
||||
tokensBefore,
|
||||
previousSummary,
|
||||
fileOps,
|
||||
settings,
|
||||
} = preparation;
|
||||
let summary: string;
|
||||
if (isSplitTurn && turnPrefixMessages.length > 0) {
|
||||
const [historyResult, turnPrefixResult] = await Promise.all([
|
||||
messagesToSummarize.length > 0
|
||||
? generateSummary(
|
||||
messagesToSummarize,
|
||||
model,
|
||||
settings.reserveTokens,
|
||||
apiKey,
|
||||
headers,
|
||||
signal,
|
||||
customInstructions,
|
||||
previousSummary,
|
||||
thinkingLevel,
|
||||
)
|
||||
: Promise.resolve("No prior history."),
|
||||
generateTurnPrefixSummary(
|
||||
turnPrefixMessages,
|
||||
model,
|
||||
settings.reserveTokens,
|
||||
apiKey,
|
||||
headers,
|
||||
signal,
|
||||
thinkingLevel,
|
||||
),
|
||||
]);
|
||||
summary = `${historyResult}\n\n---\n\n**Turn Context (split turn):**\n\n${turnPrefixResult}`;
|
||||
} else {
|
||||
summary = await generateSummary(
|
||||
messagesToSummarize,
|
||||
model,
|
||||
settings.reserveTokens,
|
||||
apiKey,
|
||||
headers,
|
||||
signal,
|
||||
customInstructions,
|
||||
previousSummary,
|
||||
thinkingLevel,
|
||||
);
|
||||
}
|
||||
const { readFiles, modifiedFiles } = computeFileLists(fileOps);
|
||||
summary += formatFileOperations(readFiles, modifiedFiles);
|
||||
return { summary, firstKeptEntryId, tokensBefore, details: { readFiles, modifiedFiles } };
|
||||
}
|
||||
|
||||
export function collectEntriesForBranchSummary(
|
||||
session: {
|
||||
getBranch(fromId?: string): Promise<SessionTreeEntry[]>;
|
||||
getEntry(id: string): Promise<SessionTreeEntry | undefined>;
|
||||
},
|
||||
oldLeafId: string | null,
|
||||
targetId: string,
|
||||
): Promise<{ entries: SessionTreeEntry[]; commonAncestorId: string | null }> {
|
||||
return (async () => {
|
||||
if (!oldLeafId) return { entries: [], commonAncestorId: null };
|
||||
const oldPath = new Set((await session.getBranch(oldLeafId)).map((e) => e.id));
|
||||
const targetPath = await session.getBranch(targetId);
|
||||
let commonAncestorId: string | null = null;
|
||||
for (let i = targetPath.length - 1; i >= 0; i--) {
|
||||
if (oldPath.has(targetPath[i]!.id)) {
|
||||
commonAncestorId = targetPath[i]!.id;
|
||||
break;
|
||||
}
|
||||
}
|
||||
const entries: SessionTreeEntry[] = [];
|
||||
let current: string | null = oldLeafId;
|
||||
while (current && current !== commonAncestorId) {
|
||||
const entry = await session.getEntry(current);
|
||||
if (!entry) break;
|
||||
entries.push(entry);
|
||||
current = entry.parentId;
|
||||
}
|
||||
entries.reverse();
|
||||
return { entries, commonAncestorId };
|
||||
})();
|
||||
}
|
||||
|
||||
export function prepareBranchEntries(
|
||||
entries: SessionTreeEntry[],
|
||||
tokenBudget: number = 0,
|
||||
): { messages: AgentMessage[]; fileOps: FileOperations; totalTokens: number } {
|
||||
const messages: AgentMessage[] = [];
|
||||
const fileOps = createFileOps();
|
||||
let totalTokens = 0;
|
||||
for (const entry of entries) {
|
||||
if (entry.type === "branch_summary" && !entry.fromHook && entry.details) {
|
||||
const details = entry.details as BranchSummaryDetails;
|
||||
if (Array.isArray(details.readFiles)) for (const f of details.readFiles) fileOps.read.add(f);
|
||||
if (Array.isArray(details.modifiedFiles)) for (const f of details.modifiedFiles) fileOps.edited.add(f);
|
||||
}
|
||||
}
|
||||
for (let i = entries.length - 1; i >= 0; i--) {
|
||||
const message = getMessageFromEntry(entries[i]!);
|
||||
if (!message) continue;
|
||||
extractFileOpsFromMessage(message, fileOps);
|
||||
const tokens = estimateTokens(message);
|
||||
if (tokenBudget > 0 && totalTokens + tokens > tokenBudget) {
|
||||
const entry = entries[i]!;
|
||||
if ((entry.type === "compaction" || entry.type === "branch_summary") && totalTokens < tokenBudget * 0.9) {
|
||||
messages.unshift(message);
|
||||
totalTokens += tokens;
|
||||
}
|
||||
break;
|
||||
}
|
||||
messages.unshift(message);
|
||||
totalTokens += tokens;
|
||||
}
|
||||
return { messages, fileOps, totalTokens };
|
||||
}
|
||||
|
||||
export async function generateBranchSummary(
|
||||
entries: SessionTreeEntry[],
|
||||
options: GenerateBranchSummaryOptions,
|
||||
): Promise<BranchSummaryResult> {
|
||||
const { model, apiKey, headers, signal, customInstructions, replaceInstructions, reserveTokens = 16384 } = options;
|
||||
const contextWindow = model.contextWindow || 128000;
|
||||
const tokenBudget = contextWindow - reserveTokens;
|
||||
const { messages, fileOps } = prepareBranchEntries(entries, tokenBudget);
|
||||
if (messages.length === 0) return { summary: "No content to summarize" };
|
||||
const llmMessages = convertToLlm(messages);
|
||||
const conversationText = serializeConversation(llmMessages);
|
||||
let instructions: string;
|
||||
if (replaceInstructions && customInstructions) instructions = customInstructions;
|
||||
else if (customInstructions) instructions = `${BRANCH_SUMMARY_PROMPT}\n\nAdditional focus: ${customInstructions}`;
|
||||
else instructions = BRANCH_SUMMARY_PROMPT;
|
||||
const promptText = `<conversation>\n${conversationText}\n</conversation>\n\n${instructions}`;
|
||||
const summarizationMessages = [
|
||||
{ role: "user" as const, content: [{ type: "text" as const, text: promptText }], timestamp: Date.now() },
|
||||
];
|
||||
const response = await completeSimple(
|
||||
model,
|
||||
{ systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages },
|
||||
{ apiKey, headers, signal, maxTokens: 2048 },
|
||||
);
|
||||
if (response.stopReason === "aborted") return { aborted: true };
|
||||
if (response.stopReason === "error") return { error: response.errorMessage || "Summarization failed" };
|
||||
let summary = response.content
|
||||
.filter((c): c is { type: "text"; text: string } => c.type === "text")
|
||||
.map((c) => c.text)
|
||||
.join("\n");
|
||||
summary = BRANCH_SUMMARY_PREAMBLE + summary;
|
||||
const { readFiles, modifiedFiles } = computeFileLists(fileOps);
|
||||
summary += formatFileOperations(readFiles, modifiedFiles);
|
||||
return { summary: summary || "No summary generated", readFiles, modifiedFiles };
|
||||
}
|
||||
273
packages/agent/src/harness/execution-env.ts
Normal file
273
packages/agent/src/harness/execution-env.ts
Normal file
@@ -0,0 +1,273 @@
|
||||
import { spawn, spawnSync } from "node:child_process";
|
||||
import { randomUUID } from "node:crypto";
|
||||
import { existsSync } from "node:fs";
|
||||
import { mkdir, mkdtemp, readdir, readFile, rm, stat, writeFile } from "node:fs/promises";
|
||||
import { tmpdir } from "node:os";
|
||||
import { isAbsolute, join, resolve } from "node:path";
|
||||
import type { ExecutionEnv } from "./types.js";
|
||||
|
||||
function resolvePath(cwd: string, path: string): string {
|
||||
return isAbsolute(path) ? path : resolve(cwd, path);
|
||||
}
|
||||
|
||||
function findBashOnPath(): string | null {
|
||||
if (process.platform === "win32") {
|
||||
try {
|
||||
const result = spawnSync("where", ["bash.exe"], { encoding: "utf-8", timeout: 5000 });
|
||||
if (result.status === 0 && result.stdout) {
|
||||
const firstMatch = result.stdout.trim().split(/\r?\n/)[0];
|
||||
if (firstMatch && existsSync(firstMatch)) {
|
||||
return firstMatch;
|
||||
}
|
||||
}
|
||||
} catch {
|
||||
// Ignore errors.
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
try {
|
||||
const result = spawnSync("which", ["bash"], { encoding: "utf-8", timeout: 5000 });
|
||||
if (result.status === 0 && result.stdout) {
|
||||
const firstMatch = result.stdout.trim().split(/\r?\n/)[0];
|
||||
if (firstMatch) {
|
||||
return firstMatch;
|
||||
}
|
||||
}
|
||||
} catch {
|
||||
// Ignore errors.
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
function getShellConfig(customShellPath?: string): { shell: string; args: string[] } {
|
||||
if (customShellPath) {
|
||||
if (existsSync(customShellPath)) {
|
||||
return { shell: customShellPath, args: ["-c"] };
|
||||
}
|
||||
throw new Error(`Custom shell path not found: ${customShellPath}`);
|
||||
}
|
||||
if (process.platform === "win32") {
|
||||
const candidates: string[] = [];
|
||||
const programFiles = process.env.ProgramFiles;
|
||||
if (programFiles) candidates.push(`${programFiles}\\Git\\bin\\bash.exe`);
|
||||
const programFilesX86 = process.env["ProgramFiles(x86)"];
|
||||
if (programFilesX86) candidates.push(`${programFilesX86}\\Git\\bin\\bash.exe`);
|
||||
for (const candidate of candidates) {
|
||||
if (existsSync(candidate)) {
|
||||
return { shell: candidate, args: ["-c"] };
|
||||
}
|
||||
}
|
||||
const bashOnPath = findBashOnPath();
|
||||
if (bashOnPath) {
|
||||
return { shell: bashOnPath, args: ["-c"] };
|
||||
}
|
||||
throw new Error("No bash shell found");
|
||||
}
|
||||
|
||||
if (existsSync("/bin/bash")) {
|
||||
return { shell: "/bin/bash", args: ["-c"] };
|
||||
}
|
||||
const bashOnPath = findBashOnPath();
|
||||
if (bashOnPath) {
|
||||
return { shell: bashOnPath, args: ["-c"] };
|
||||
}
|
||||
return { shell: "sh", args: ["-c"] };
|
||||
}
|
||||
|
||||
function getShellEnv(baseEnv?: NodeJS.ProcessEnv, extraEnv?: Record<string, string>): NodeJS.ProcessEnv {
|
||||
return {
|
||||
...process.env,
|
||||
...baseEnv,
|
||||
...extraEnv,
|
||||
};
|
||||
}
|
||||
|
||||
function killProcessTree(pid: number): void {
|
||||
if (process.platform === "win32") {
|
||||
try {
|
||||
spawn("taskkill", ["/F", "/T", "/PID", String(pid)], {
|
||||
stdio: "ignore",
|
||||
detached: true,
|
||||
});
|
||||
} catch {
|
||||
// Ignore errors.
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
process.kill(-pid, "SIGKILL");
|
||||
} catch {
|
||||
try {
|
||||
process.kill(pid, "SIGKILL");
|
||||
} catch {
|
||||
// Process already dead.
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
export class NodeExecutionEnv implements ExecutionEnv {
|
||||
cwd: string;
|
||||
private shellPath?: string;
|
||||
private shellEnv?: NodeJS.ProcessEnv;
|
||||
|
||||
constructor(options: { cwd: string; shellPath?: string; shellEnv?: NodeJS.ProcessEnv }) {
|
||||
this.cwd = options.cwd;
|
||||
this.shellPath = options.shellPath;
|
||||
this.shellEnv = options.shellEnv;
|
||||
}
|
||||
|
||||
async exec(
|
||||
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 }> {
|
||||
const cwd = options?.cwd ? resolvePath(this.cwd, options.cwd) : this.cwd;
|
||||
const { shell, args } = getShellConfig(this.shellPath);
|
||||
|
||||
return await new Promise((resolvePromise, reject) => {
|
||||
let stdout = "";
|
||||
let stderr = "";
|
||||
let settled = false;
|
||||
let timedOut = false;
|
||||
const child = spawn(shell, [...args, command], {
|
||||
cwd,
|
||||
detached: process.platform !== "win32",
|
||||
env: getShellEnv(this.shellEnv, options?.env),
|
||||
stdio: ["ignore", "pipe", "pipe"],
|
||||
});
|
||||
|
||||
const timeoutId =
|
||||
typeof options?.timeout === "number"
|
||||
? setTimeout(() => {
|
||||
timedOut = true;
|
||||
if (child.pid) {
|
||||
killProcessTree(child.pid);
|
||||
}
|
||||
}, options.timeout * 1000)
|
||||
: undefined;
|
||||
|
||||
const onAbort = () => {
|
||||
if (child.pid) {
|
||||
killProcessTree(child.pid);
|
||||
}
|
||||
};
|
||||
if (options?.signal) {
|
||||
if (options.signal.aborted) {
|
||||
onAbort();
|
||||
} else {
|
||||
options.signal.addEventListener("abort", onAbort, { once: true });
|
||||
}
|
||||
}
|
||||
|
||||
child.stdout?.setEncoding("utf8");
|
||||
child.stderr?.setEncoding("utf8");
|
||||
child.stdout?.on("data", (chunk: string) => {
|
||||
stdout += chunk;
|
||||
options?.onStdout?.(chunk);
|
||||
});
|
||||
child.stderr?.on("data", (chunk: string) => {
|
||||
stderr += chunk;
|
||||
options?.onStderr?.(chunk);
|
||||
});
|
||||
|
||||
child.on("error", (error) => {
|
||||
if (timeoutId) clearTimeout(timeoutId);
|
||||
if (options?.signal) options.signal.removeEventListener("abort", onAbort);
|
||||
if (settled) return;
|
||||
settled = true;
|
||||
reject(error);
|
||||
});
|
||||
|
||||
child.on("close", (code) => {
|
||||
if (timeoutId) clearTimeout(timeoutId);
|
||||
if (options?.signal) options.signal.removeEventListener("abort", onAbort);
|
||||
if (settled) return;
|
||||
settled = true;
|
||||
if (options?.signal?.aborted) {
|
||||
reject(new Error("aborted"));
|
||||
return;
|
||||
}
|
||||
if (timedOut) {
|
||||
reject(new Error(`timeout:${options?.timeout}`));
|
||||
return;
|
||||
}
|
||||
resolvePromise({ stdout, stderr, exitCode: code ?? 0 });
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
async readTextFile(path: string): Promise<string> {
|
||||
return await readFile(resolvePath(this.cwd, path), "utf8");
|
||||
}
|
||||
|
||||
async readBinaryFile(path: string): Promise<Uint8Array> {
|
||||
return await readFile(resolvePath(this.cwd, path));
|
||||
}
|
||||
|
||||
async writeFile(path: string, content: string | Uint8Array): Promise<void> {
|
||||
const resolved = resolvePath(this.cwd, path);
|
||||
await mkdir(resolve(resolved, ".."), { recursive: true });
|
||||
await writeFile(resolved, content);
|
||||
}
|
||||
|
||||
async stat(
|
||||
path: string,
|
||||
): Promise<{ isFile: boolean; isDirectory: boolean; isSymbolicLink: boolean; size: number; mtime: Date }> {
|
||||
const s = await stat(resolvePath(this.cwd, path));
|
||||
return {
|
||||
isFile: s.isFile(),
|
||||
isDirectory: s.isDirectory(),
|
||||
isSymbolicLink: s.isSymbolicLink(),
|
||||
size: s.size,
|
||||
mtime: s.mtime,
|
||||
};
|
||||
}
|
||||
|
||||
async listDir(path: string): Promise<string[]> {
|
||||
return await readdir(resolvePath(this.cwd, path));
|
||||
}
|
||||
|
||||
async pathExists(path: string): Promise<boolean> {
|
||||
try {
|
||||
await stat(resolvePath(this.cwd, path));
|
||||
return true;
|
||||
} catch {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
async createDir(path: string, options?: { recursive?: boolean }): Promise<void> {
|
||||
await mkdir(resolvePath(this.cwd, path), { recursive: options?.recursive });
|
||||
}
|
||||
|
||||
async remove(path: string, options?: { recursive?: boolean; force?: boolean }): Promise<void> {
|
||||
await rm(resolvePath(this.cwd, path), { recursive: options?.recursive, force: options?.force });
|
||||
}
|
||||
|
||||
async createTempDir(prefix: string = "tmp-"): Promise<string> {
|
||||
return await mkdtemp(join(tmpdir(), prefix));
|
||||
}
|
||||
|
||||
async createTempFile(options?: { prefix?: string; suffix?: string }): Promise<string> {
|
||||
const dir = await this.createTempDir("tmp-");
|
||||
const filePath = join(dir, `${options?.prefix ?? ""}${randomUUID()}${options?.suffix ?? ""}`);
|
||||
await writeFile(filePath, "");
|
||||
return filePath;
|
||||
}
|
||||
|
||||
resolvePath(path: string): string {
|
||||
return resolvePath(this.cwd, path);
|
||||
}
|
||||
|
||||
async cleanup(): Promise<void> {
|
||||
// nothing to clean up for the local node implementation
|
||||
}
|
||||
}
|
||||
164
packages/agent/src/harness/messages.ts
Normal file
164
packages/agent/src/harness/messages.ts
Normal file
@@ -0,0 +1,164 @@
|
||||
import type { ImageContent, Message, TextContent } from "@mariozechner/pi-ai";
|
||||
import type { AgentMessage } from "../types.js";
|
||||
|
||||
export const COMPACTION_SUMMARY_PREFIX = `The conversation history before this point was compacted into the following summary:
|
||||
|
||||
<summary>
|
||||
`;
|
||||
|
||||
export const COMPACTION_SUMMARY_SUFFIX = `
|
||||
</summary>`;
|
||||
|
||||
export const BRANCH_SUMMARY_PREFIX = `The following is a summary of a branch that this conversation came back from:
|
||||
|
||||
<summary>
|
||||
`;
|
||||
|
||||
export const BRANCH_SUMMARY_SUFFIX = `</summary>`;
|
||||
|
||||
export interface BashExecutionMessage {
|
||||
role: "bashExecution";
|
||||
command: string;
|
||||
output: string;
|
||||
exitCode: number | undefined;
|
||||
cancelled: boolean;
|
||||
truncated: boolean;
|
||||
fullOutputPath?: string;
|
||||
timestamp: number;
|
||||
excludeFromContext?: boolean;
|
||||
}
|
||||
|
||||
export interface CustomMessage<T = unknown> {
|
||||
role: "custom";
|
||||
customType: string;
|
||||
content: string | (TextContent | ImageContent)[];
|
||||
display: boolean;
|
||||
details?: T;
|
||||
timestamp: number;
|
||||
}
|
||||
|
||||
export interface BranchSummaryMessage {
|
||||
role: "branchSummary";
|
||||
summary: string;
|
||||
fromId: string;
|
||||
timestamp: number;
|
||||
}
|
||||
|
||||
export interface CompactionSummaryMessage {
|
||||
role: "compactionSummary";
|
||||
summary: string;
|
||||
tokensBefore: number;
|
||||
timestamp: number;
|
||||
}
|
||||
|
||||
declare module "../types.js" {
|
||||
interface CustomAgentMessages {
|
||||
bashExecution: BashExecutionMessage;
|
||||
custom: CustomMessage;
|
||||
branchSummary: BranchSummaryMessage;
|
||||
compactionSummary: CompactionSummaryMessage;
|
||||
}
|
||||
}
|
||||
|
||||
export function bashExecutionToText(msg: BashExecutionMessage): string {
|
||||
let text = `Ran \`${msg.command}\`\n`;
|
||||
if (msg.output) {
|
||||
text += `\`\`\`\n${msg.output}\n\`\`\``;
|
||||
} else {
|
||||
text += "(no output)";
|
||||
}
|
||||
if (msg.cancelled) {
|
||||
text += "\n\n(command cancelled)";
|
||||
} else if (msg.exitCode !== null && msg.exitCode !== undefined && msg.exitCode !== 0) {
|
||||
text += `\n\nCommand exited with code ${msg.exitCode}`;
|
||||
}
|
||||
if (msg.truncated && msg.fullOutputPath) {
|
||||
text += `\n\n[Output truncated. Full output: ${msg.fullOutputPath}]`;
|
||||
}
|
||||
return text;
|
||||
}
|
||||
|
||||
export function createBranchSummaryMessage(summary: string, fromId: string, timestamp: string): BranchSummaryMessage {
|
||||
return {
|
||||
role: "branchSummary",
|
||||
summary,
|
||||
fromId,
|
||||
timestamp: new Date(timestamp).getTime(),
|
||||
};
|
||||
}
|
||||
|
||||
export function createCompactionSummaryMessage(
|
||||
summary: string,
|
||||
tokensBefore: number,
|
||||
timestamp: string,
|
||||
): CompactionSummaryMessage {
|
||||
return {
|
||||
role: "compactionSummary",
|
||||
summary,
|
||||
tokensBefore,
|
||||
timestamp: new Date(timestamp).getTime(),
|
||||
};
|
||||
}
|
||||
|
||||
export function createCustomMessage(
|
||||
customType: string,
|
||||
content: string | (TextContent | ImageContent)[],
|
||||
display: boolean,
|
||||
details: unknown | undefined,
|
||||
timestamp: string,
|
||||
): CustomMessage {
|
||||
return {
|
||||
role: "custom",
|
||||
customType,
|
||||
content,
|
||||
display,
|
||||
details,
|
||||
timestamp: new Date(timestamp).getTime(),
|
||||
};
|
||||
}
|
||||
|
||||
export function convertToLlm(messages: AgentMessage[]): Message[] {
|
||||
return messages
|
||||
.map((m): Message | undefined => {
|
||||
switch (m.role) {
|
||||
case "bashExecution":
|
||||
if (m.excludeFromContext) {
|
||||
return undefined;
|
||||
}
|
||||
return {
|
||||
role: "user",
|
||||
content: [{ type: "text", text: bashExecutionToText(m) }],
|
||||
timestamp: m.timestamp,
|
||||
};
|
||||
case "custom": {
|
||||
const content = typeof m.content === "string" ? [{ type: "text" as const, text: m.content }] : m.content;
|
||||
return {
|
||||
role: "user",
|
||||
content,
|
||||
timestamp: m.timestamp,
|
||||
};
|
||||
}
|
||||
case "branchSummary":
|
||||
return {
|
||||
role: "user",
|
||||
content: [{ type: "text" as const, text: BRANCH_SUMMARY_PREFIX + m.summary + BRANCH_SUMMARY_SUFFIX }],
|
||||
timestamp: m.timestamp,
|
||||
};
|
||||
case "compactionSummary":
|
||||
return {
|
||||
role: "user",
|
||||
content: [
|
||||
{ type: "text" as const, text: COMPACTION_SUMMARY_PREFIX + m.summary + COMPACTION_SUMMARY_SUFFIX },
|
||||
],
|
||||
timestamp: m.timestamp,
|
||||
};
|
||||
case "user":
|
||||
case "assistant":
|
||||
case "toolResult":
|
||||
return m;
|
||||
default:
|
||||
return undefined;
|
||||
}
|
||||
})
|
||||
.filter((m): m is Message => m !== undefined);
|
||||
}
|
||||
51
packages/agent/src/harness/prompt-templates.ts
Normal file
51
packages/agent/src/harness/prompt-templates.ts
Normal file
@@ -0,0 +1,51 @@
|
||||
import type { PromptTemplate } from "./types.js";
|
||||
|
||||
export function parseCommandArgs(argsString: string): string[] {
|
||||
const args: string[] = [];
|
||||
let current = "";
|
||||
let inQuote: string | null = null;
|
||||
|
||||
for (let i = 0; i < argsString.length; i++) {
|
||||
const char = argsString[i]!;
|
||||
if (inQuote) {
|
||||
if (char === inQuote) inQuote = null;
|
||||
else current += char;
|
||||
} else if (char === '"' || char === "'") {
|
||||
inQuote = char;
|
||||
} else if (char === " " || char === "\t") {
|
||||
if (current) {
|
||||
args.push(current);
|
||||
current = "";
|
||||
}
|
||||
} else {
|
||||
current += char;
|
||||
}
|
||||
}
|
||||
if (current) args.push(current);
|
||||
return args;
|
||||
}
|
||||
|
||||
export function substituteArgs(content: string, args: string[]): string {
|
||||
let result = content;
|
||||
result = result.replace(/\$(\d+)/g, (_, num: string) => args[parseInt(num, 10) - 1] ?? "");
|
||||
result = result.replace(/\$\{@:(\d+)(?::(\d+))?\}/g, (_, startStr: string, lengthStr?: string) => {
|
||||
let start = parseInt(startStr, 10) - 1;
|
||||
if (start < 0) start = 0;
|
||||
if (lengthStr) return args.slice(start, start + parseInt(lengthStr, 10)).join(" ");
|
||||
return args.slice(start).join(" ");
|
||||
});
|
||||
const allArgs = args.join(" ");
|
||||
result = result.replace(/\$ARGUMENTS/g, allArgs);
|
||||
result = result.replace(/\$@/g, allArgs);
|
||||
return result;
|
||||
}
|
||||
|
||||
export function expandPromptTemplate(text: string, templates: PromptTemplate[]): string {
|
||||
if (!text.startsWith("/")) return text;
|
||||
const spaceIndex = text.indexOf(" ");
|
||||
const commandName = spaceIndex === -1 ? text.slice(1) : text.slice(1, spaceIndex);
|
||||
const argsString = spaceIndex === -1 ? "" : text.slice(spaceIndex + 1);
|
||||
const template = templates.find((candidate) => candidate.name === commandName);
|
||||
if (!template) return text;
|
||||
return substituteArgs(template.content, parseCommandArgs(argsString));
|
||||
}
|
||||
411
packages/agent/src/harness/session-tree.ts
Normal file
411
packages/agent/src/harness/session-tree.ts
Normal file
@@ -0,0 +1,411 @@
|
||||
import { randomUUID } from "node:crypto";
|
||||
import { appendFile, mkdir, readFile, writeFile } from "node:fs/promises";
|
||||
import { dirname, resolve } from "node:path";
|
||||
import type { ImageContent, TextContent } from "@mariozechner/pi-ai";
|
||||
import type { AgentMessage } from "../types.js";
|
||||
import { createBranchSummaryMessage, createCompactionSummaryMessage, createCustomMessage } from "./messages.js";
|
||||
import type {
|
||||
BranchSummaryEntry,
|
||||
CompactionEntry,
|
||||
CustomEntry,
|
||||
CustomMessageEntry,
|
||||
LabelEntry,
|
||||
MessageEntry,
|
||||
ModelChangeEntry,
|
||||
SessionContext,
|
||||
SessionInfoEntry,
|
||||
SessionTree,
|
||||
SessionTreeEntry,
|
||||
SessionTreeStorage,
|
||||
ThinkingLevelChangeEntry,
|
||||
} from "./types.js";
|
||||
|
||||
function generateId(byId: { has(id: string): boolean }): string {
|
||||
for (let i = 0; i < 100; i++) {
|
||||
const id = randomUUID().slice(0, 8);
|
||||
if (!byId.has(id)) return id;
|
||||
}
|
||||
return randomUUID();
|
||||
}
|
||||
|
||||
export function buildSessionContext(entries: SessionTreeEntry[]): SessionContext {
|
||||
let thinkingLevel = "off";
|
||||
let model: { provider: string; modelId: string } | null = null;
|
||||
let compaction: CompactionEntry | null = null;
|
||||
|
||||
for (const entry of entries) {
|
||||
if (entry.type === "thinking_level_change") {
|
||||
thinkingLevel = entry.thinkingLevel;
|
||||
} else if (entry.type === "model_change") {
|
||||
model = { provider: entry.provider, modelId: entry.modelId };
|
||||
} else if (entry.type === "message" && entry.message.role === "assistant") {
|
||||
model = { provider: entry.message.provider, modelId: entry.message.model };
|
||||
} else if (entry.type === "compaction") {
|
||||
compaction = entry;
|
||||
}
|
||||
}
|
||||
|
||||
const messages: AgentMessage[] = [];
|
||||
const appendMessage = (entry: SessionTreeEntry) => {
|
||||
if (entry.type === "message") {
|
||||
messages.push(entry.message);
|
||||
} else if (entry.type === "custom_message") {
|
||||
messages.push(
|
||||
createCustomMessage(entry.customType, entry.content, entry.display, entry.details, entry.timestamp),
|
||||
);
|
||||
} else if (entry.type === "branch_summary" && entry.summary) {
|
||||
messages.push(createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp));
|
||||
}
|
||||
};
|
||||
|
||||
if (compaction) {
|
||||
messages.push(createCompactionSummaryMessage(compaction.summary, compaction.tokensBefore, compaction.timestamp));
|
||||
const compactionIdx = entries.findIndex((e) => e.type === "compaction" && e.id === compaction.id);
|
||||
let foundFirstKept = false;
|
||||
for (let i = 0; i < compactionIdx; i++) {
|
||||
const entry = entries[i]!;
|
||||
if (entry.id === compaction.firstKeptEntryId) foundFirstKept = true;
|
||||
if (foundFirstKept) appendMessage(entry);
|
||||
}
|
||||
for (let i = compactionIdx + 1; i < entries.length; i++) {
|
||||
appendMessage(entries[i]!);
|
||||
}
|
||||
} else {
|
||||
for (const entry of entries) {
|
||||
appendMessage(entry);
|
||||
}
|
||||
}
|
||||
|
||||
return { messages, thinkingLevel, model };
|
||||
}
|
||||
|
||||
interface SessionHeader {
|
||||
type: "session";
|
||||
version: 3;
|
||||
id: string;
|
||||
timestamp: string;
|
||||
cwd: string;
|
||||
parentSession?: string;
|
||||
}
|
||||
|
||||
async function loadJsonlStorage(
|
||||
filePath: string,
|
||||
): Promise<{ header?: SessionHeader; entries: SessionTreeEntry[]; leafId: string | null }> {
|
||||
try {
|
||||
const content = await readFile(filePath, "utf8");
|
||||
const entries: SessionTreeEntry[] = [];
|
||||
let header: SessionHeader | undefined;
|
||||
let leafId: string | null = null;
|
||||
for (const line of content.split("\n")) {
|
||||
if (!line.trim()) continue;
|
||||
try {
|
||||
const record = JSON.parse(line) as SessionHeader | SessionTreeEntry;
|
||||
if (record.type === "session") {
|
||||
header = record as SessionHeader;
|
||||
continue;
|
||||
}
|
||||
entries.push(record as SessionTreeEntry);
|
||||
leafId = (record as SessionTreeEntry).id;
|
||||
} catch {
|
||||
// ignore malformed lines
|
||||
}
|
||||
}
|
||||
return { header, entries, leafId };
|
||||
} catch {
|
||||
return { entries: [], leafId: null };
|
||||
}
|
||||
}
|
||||
|
||||
export class JsonlSessionTreeStorage implements SessionTreeStorage {
|
||||
private filePath: string;
|
||||
private cwd: string;
|
||||
private headerInitialized = false;
|
||||
private cacheLoaded = false;
|
||||
private entries: SessionTreeEntry[] = [];
|
||||
private byId = new Map<string, SessionTreeEntry>();
|
||||
private currentLeafId: string | null = null;
|
||||
|
||||
constructor(filePath: string, options: { cwd: string }) {
|
||||
this.filePath = resolve(filePath);
|
||||
this.cwd = options.cwd;
|
||||
}
|
||||
|
||||
private async ensureParentDir(): Promise<void> {
|
||||
await mkdir(dirname(this.filePath), { recursive: true });
|
||||
}
|
||||
|
||||
private async ensureLoaded(): Promise<void> {
|
||||
if (this.cacheLoaded) {
|
||||
return;
|
||||
}
|
||||
const loaded = await loadJsonlStorage(this.filePath);
|
||||
this.entries = loaded.entries;
|
||||
this.byId = new Map(loaded.entries.map((entry) => [entry.id, entry]));
|
||||
this.currentLeafId = loaded.leafId;
|
||||
this.headerInitialized = loaded.header !== undefined;
|
||||
this.cacheLoaded = true;
|
||||
}
|
||||
|
||||
private async ensureHeader(): Promise<void> {
|
||||
await this.ensureLoaded();
|
||||
if (this.headerInitialized) return;
|
||||
await this.ensureParentDir();
|
||||
const header: SessionHeader = {
|
||||
type: "session",
|
||||
version: 3,
|
||||
id: randomUUID(),
|
||||
timestamp: new Date().toISOString(),
|
||||
cwd: this.cwd,
|
||||
};
|
||||
await writeFile(this.filePath, `${JSON.stringify(header)}\n`);
|
||||
this.headerInitialized = true;
|
||||
}
|
||||
|
||||
async getLeafId(): Promise<string | null> {
|
||||
await this.ensureLoaded();
|
||||
return this.currentLeafId;
|
||||
}
|
||||
|
||||
async setLeafId(leafId: string | null): Promise<void> {
|
||||
await this.ensureLoaded();
|
||||
this.currentLeafId = leafId;
|
||||
}
|
||||
|
||||
async appendEntry(entry: SessionTreeEntry): Promise<void> {
|
||||
await this.ensureHeader();
|
||||
await appendFile(this.filePath, `${JSON.stringify(entry)}\n`);
|
||||
this.entries.push(entry);
|
||||
this.byId.set(entry.id, entry);
|
||||
this.currentLeafId = entry.id;
|
||||
}
|
||||
|
||||
async getEntry(id: string): Promise<SessionTreeEntry | undefined> {
|
||||
await this.ensureLoaded();
|
||||
return this.byId.get(id);
|
||||
}
|
||||
|
||||
async getPathToRoot(leafId: string | null): Promise<SessionTreeEntry[]> {
|
||||
await this.ensureLoaded();
|
||||
if (leafId === null) return [];
|
||||
const path: SessionTreeEntry[] = [];
|
||||
let current = this.byId.get(leafId);
|
||||
while (current) {
|
||||
path.unshift(current);
|
||||
current = current.parentId ? this.byId.get(current.parentId) : undefined;
|
||||
}
|
||||
return path;
|
||||
}
|
||||
|
||||
async getEntries(): Promise<SessionTreeEntry[]> {
|
||||
await this.ensureLoaded();
|
||||
return [...this.entries];
|
||||
}
|
||||
}
|
||||
|
||||
export class InMemorySessionTreeStorage implements SessionTreeStorage {
|
||||
private entries: SessionTreeEntry[];
|
||||
private leafId: string | null;
|
||||
|
||||
constructor(options?: { entries?: SessionTreeEntry[]; leafId?: string | null }) {
|
||||
this.entries = options?.entries ? [...options.entries] : [];
|
||||
this.leafId = options?.leafId ?? this.entries[this.entries.length - 1]?.id ?? null;
|
||||
}
|
||||
|
||||
async getLeafId(): Promise<string | null> {
|
||||
return this.leafId;
|
||||
}
|
||||
|
||||
async setLeafId(leafId: string | null): Promise<void> {
|
||||
this.leafId = leafId;
|
||||
}
|
||||
|
||||
async appendEntry(entry: SessionTreeEntry): Promise<void> {
|
||||
this.entries.push(entry);
|
||||
this.leafId = entry.id;
|
||||
}
|
||||
|
||||
async getEntry(id: string): Promise<SessionTreeEntry | undefined> {
|
||||
return this.entries.find((entry) => entry.id === id);
|
||||
}
|
||||
|
||||
async getPathToRoot(leafId: string | null): Promise<SessionTreeEntry[]> {
|
||||
if (leafId === null) return [];
|
||||
const byId = new Map<string, SessionTreeEntry>(this.entries.map((entry) => [entry.id, entry]));
|
||||
const path: SessionTreeEntry[] = [];
|
||||
let current = byId.get(leafId);
|
||||
while (current) {
|
||||
path.unshift(current);
|
||||
current = current.parentId ? byId.get(current.parentId) : undefined;
|
||||
}
|
||||
return path;
|
||||
}
|
||||
|
||||
async getEntries(): Promise<SessionTreeEntry[]> {
|
||||
return [...this.entries];
|
||||
}
|
||||
}
|
||||
|
||||
export class DefaultSessionTree implements SessionTree {
|
||||
private storage: SessionTreeStorage;
|
||||
|
||||
constructor(storage?: SessionTreeStorage) {
|
||||
this.storage = storage ?? new InMemorySessionTreeStorage();
|
||||
}
|
||||
|
||||
getLeafId(): Promise<string | null> {
|
||||
return this.storage.getLeafId();
|
||||
}
|
||||
|
||||
getEntry(id: string): Promise<SessionTreeEntry | undefined> {
|
||||
return this.storage.getEntry(id);
|
||||
}
|
||||
|
||||
getEntries(): Promise<SessionTreeEntry[]> {
|
||||
return this.storage.getEntries();
|
||||
}
|
||||
|
||||
async getBranch(fromId?: string): Promise<SessionTreeEntry[]> {
|
||||
const leafId = fromId ?? (await this.storage.getLeafId());
|
||||
return this.storage.getPathToRoot(leafId);
|
||||
}
|
||||
|
||||
async buildContext(): Promise<SessionContext> {
|
||||
return buildSessionContext(await this.getBranch());
|
||||
}
|
||||
|
||||
private async makeEntryId(): Promise<string> {
|
||||
const entries = await this.storage.getEntries();
|
||||
return generateId(new Set(entries.map((entry) => entry.id)));
|
||||
}
|
||||
|
||||
private async appendTypedEntry<TEntry extends SessionTreeEntry>(entry: TEntry): Promise<string> {
|
||||
await this.storage.appendEntry(entry);
|
||||
return entry.id;
|
||||
}
|
||||
|
||||
async appendMessage(message: AgentMessage): Promise<string> {
|
||||
return this.appendTypedEntry({
|
||||
type: "message",
|
||||
id: await this.makeEntryId(),
|
||||
parentId: await this.storage.getLeafId(),
|
||||
timestamp: new Date().toISOString(),
|
||||
message,
|
||||
} satisfies MessageEntry);
|
||||
}
|
||||
|
||||
async appendThinkingLevelChange(thinkingLevel: string): Promise<string> {
|
||||
return this.appendTypedEntry({
|
||||
type: "thinking_level_change",
|
||||
id: await this.makeEntryId(),
|
||||
parentId: await this.storage.getLeafId(),
|
||||
timestamp: new Date().toISOString(),
|
||||
thinkingLevel,
|
||||
} satisfies ThinkingLevelChangeEntry);
|
||||
}
|
||||
|
||||
async appendModelChange(provider: string, modelId: string): Promise<string> {
|
||||
return this.appendTypedEntry({
|
||||
type: "model_change",
|
||||
id: await this.makeEntryId(),
|
||||
parentId: await this.storage.getLeafId(),
|
||||
timestamp: new Date().toISOString(),
|
||||
provider,
|
||||
modelId,
|
||||
} satisfies ModelChangeEntry);
|
||||
}
|
||||
|
||||
async appendCompaction<T = unknown>(
|
||||
summary: string,
|
||||
firstKeptEntryId: string,
|
||||
tokensBefore: number,
|
||||
details?: T,
|
||||
fromHook?: boolean,
|
||||
): Promise<string> {
|
||||
return this.appendTypedEntry({
|
||||
type: "compaction",
|
||||
id: await this.makeEntryId(),
|
||||
parentId: await this.storage.getLeafId(),
|
||||
timestamp: new Date().toISOString(),
|
||||
summary,
|
||||
firstKeptEntryId,
|
||||
tokensBefore,
|
||||
details,
|
||||
fromHook,
|
||||
} satisfies CompactionEntry<T>);
|
||||
}
|
||||
|
||||
async appendBranchSummary<T = unknown>(
|
||||
fromId: string,
|
||||
summary: string,
|
||||
details?: T,
|
||||
fromHook?: boolean,
|
||||
): Promise<string> {
|
||||
return this.appendTypedEntry({
|
||||
type: "branch_summary",
|
||||
id: await this.makeEntryId(),
|
||||
parentId: await this.storage.getLeafId(),
|
||||
timestamp: new Date().toISOString(),
|
||||
fromId,
|
||||
summary,
|
||||
details,
|
||||
fromHook,
|
||||
} satisfies BranchSummaryEntry<T>);
|
||||
}
|
||||
|
||||
async appendCustomEntry(customType: string, data?: unknown): Promise<string> {
|
||||
return this.appendTypedEntry({
|
||||
type: "custom",
|
||||
id: await this.makeEntryId(),
|
||||
parentId: await this.storage.getLeafId(),
|
||||
timestamp: new Date().toISOString(),
|
||||
customType,
|
||||
data,
|
||||
} satisfies CustomEntry);
|
||||
}
|
||||
|
||||
async appendCustomMessageEntry<T = unknown>(
|
||||
customType: string,
|
||||
content: string | (TextContent | ImageContent)[],
|
||||
display: boolean,
|
||||
details?: T,
|
||||
): Promise<string> {
|
||||
return this.appendTypedEntry({
|
||||
type: "custom_message",
|
||||
id: await this.makeEntryId(),
|
||||
parentId: await this.storage.getLeafId(),
|
||||
timestamp: new Date().toISOString(),
|
||||
customType,
|
||||
content,
|
||||
display,
|
||||
details,
|
||||
} satisfies CustomMessageEntry<T>);
|
||||
}
|
||||
|
||||
async appendLabelChange(targetId: string, label: string | undefined): Promise<string> {
|
||||
return this.appendTypedEntry({
|
||||
type: "label",
|
||||
id: await this.makeEntryId(),
|
||||
parentId: await this.storage.getLeafId(),
|
||||
timestamp: new Date().toISOString(),
|
||||
targetId,
|
||||
label,
|
||||
} satisfies LabelEntry);
|
||||
}
|
||||
|
||||
async appendSessionInfo(name: string): Promise<string> {
|
||||
return this.appendTypedEntry({
|
||||
type: "session_info",
|
||||
id: await this.makeEntryId(),
|
||||
parentId: await this.storage.getLeafId(),
|
||||
timestamp: new Date().toISOString(),
|
||||
name: name.trim(),
|
||||
} satisfies SessionInfoEntry);
|
||||
}
|
||||
|
||||
async moveTo(entryId: string | null): Promise<void> {
|
||||
if (entryId !== null && !(await this.storage.getEntry(entryId))) {
|
||||
throw new Error(`Entry ${entryId} not found`);
|
||||
}
|
||||
await this.storage.setLeafId(entryId);
|
||||
}
|
||||
}
|
||||
561
packages/agent/src/harness/types.ts
Normal file
561
packages/agent/src/harness/types.ts
Normal file
@@ -0,0 +1,561 @@
|
||||
import type { ImageContent, Model, TextContent } from "@mariozechner/pi-ai";
|
||||
import type { Agent, AgentEvent, AgentMessage, ThinkingLevel } from "../index.js";
|
||||
|
||||
export type SourceScope = "user" | "project" | "temporary";
|
||||
export type SourceOrigin = "package" | "top-level";
|
||||
|
||||
export interface SourceInfo {
|
||||
path: string;
|
||||
source: string;
|
||||
scope: SourceScope;
|
||||
origin: SourceOrigin;
|
||||
baseDir?: string;
|
||||
}
|
||||
|
||||
export interface Skill {
|
||||
name: string;
|
||||
description: string;
|
||||
filePath: string;
|
||||
baseDir: string;
|
||||
sourceInfo: SourceInfo;
|
||||
disableModelInvocation: boolean;
|
||||
}
|
||||
|
||||
export interface PromptTemplate {
|
||||
name: string;
|
||||
description: string;
|
||||
argumentHint?: string;
|
||||
content: string;
|
||||
sourceInfo: SourceInfo;
|
||||
filePath: string;
|
||||
}
|
||||
|
||||
export interface SystemPromptInputs {
|
||||
basePrompt?: string;
|
||||
appendPrompt?: string;
|
||||
contextFiles?: Array<{ path: string; content: string }>;
|
||||
skills?: Skill[];
|
||||
}
|
||||
|
||||
export interface ExecutionEnvExecOptions {
|
||||
cwd?: string;
|
||||
env?: Record<string, string>;
|
||||
timeout?: number;
|
||||
signal?: AbortSignal;
|
||||
onStdout?: (chunk: string) => void;
|
||||
onStderr?: (chunk: string) => void;
|
||||
}
|
||||
|
||||
export interface ExecutionEnv {
|
||||
cwd: string;
|
||||
|
||||
exec(
|
||||
command: string,
|
||||
options?: ExecutionEnvExecOptions,
|
||||
): Promise<{ stdout: string; stderr: string; exitCode: number }>;
|
||||
|
||||
readTextFile(path: string): Promise<string>;
|
||||
readBinaryFile(path: string): Promise<Uint8Array>;
|
||||
writeFile(path: string, content: string | Uint8Array): Promise<void>;
|
||||
stat(path: string): Promise<{
|
||||
isFile: boolean;
|
||||
isDirectory: boolean;
|
||||
isSymbolicLink: boolean;
|
||||
size: number;
|
||||
mtime: Date;
|
||||
}>;
|
||||
listDir(path: string): Promise<string[]>;
|
||||
pathExists(path: string): Promise<boolean>;
|
||||
createDir(path: string, options?: { recursive?: boolean }): Promise<void>;
|
||||
remove(path: string, options?: { recursive?: boolean; force?: boolean }): Promise<void>;
|
||||
createTempDir(prefix?: string): Promise<string>;
|
||||
createTempFile(options?: { prefix?: string; suffix?: string }): Promise<string>;
|
||||
|
||||
resolvePath(path: string): string;
|
||||
cleanup(): Promise<void>;
|
||||
}
|
||||
|
||||
export interface SessionTreeEntryBase {
|
||||
type: string;
|
||||
id: string;
|
||||
parentId: string | null;
|
||||
timestamp: string;
|
||||
}
|
||||
|
||||
export interface MessageEntry extends SessionTreeEntryBase {
|
||||
type: "message";
|
||||
message: AgentMessage;
|
||||
}
|
||||
|
||||
export interface ThinkingLevelChangeEntry extends SessionTreeEntryBase {
|
||||
type: "thinking_level_change";
|
||||
thinkingLevel: string;
|
||||
}
|
||||
|
||||
export interface ModelChangeEntry extends SessionTreeEntryBase {
|
||||
type: "model_change";
|
||||
provider: string;
|
||||
modelId: string;
|
||||
}
|
||||
|
||||
export interface CompactionEntry<T = unknown> extends SessionTreeEntryBase {
|
||||
type: "compaction";
|
||||
summary: string;
|
||||
firstKeptEntryId: string;
|
||||
tokensBefore: number;
|
||||
details?: T;
|
||||
fromHook?: boolean;
|
||||
}
|
||||
|
||||
export interface BranchSummaryEntry<T = unknown> extends SessionTreeEntryBase {
|
||||
type: "branch_summary";
|
||||
fromId: string;
|
||||
summary: string;
|
||||
details?: T;
|
||||
fromHook?: boolean;
|
||||
}
|
||||
|
||||
export interface CustomEntry<T = unknown> extends SessionTreeEntryBase {
|
||||
type: "custom";
|
||||
customType: string;
|
||||
data?: T;
|
||||
}
|
||||
|
||||
export interface CustomMessageEntry<T = unknown> extends SessionTreeEntryBase {
|
||||
type: "custom_message";
|
||||
customType: string;
|
||||
content: string | (TextContent | ImageContent)[];
|
||||
details?: T;
|
||||
display: boolean;
|
||||
}
|
||||
|
||||
export interface LabelEntry extends SessionTreeEntryBase {
|
||||
type: "label";
|
||||
targetId: string;
|
||||
label: string | undefined;
|
||||
}
|
||||
|
||||
export interface SessionInfoEntry extends SessionTreeEntryBase {
|
||||
type: "session_info";
|
||||
name?: string;
|
||||
}
|
||||
|
||||
export type SessionTreeEntry =
|
||||
| MessageEntry
|
||||
| ThinkingLevelChangeEntry
|
||||
| ModelChangeEntry
|
||||
| CompactionEntry
|
||||
| BranchSummaryEntry
|
||||
| CustomEntry
|
||||
| CustomMessageEntry
|
||||
| LabelEntry
|
||||
| SessionInfoEntry;
|
||||
|
||||
export interface SessionContext {
|
||||
messages: AgentMessage[];
|
||||
thinkingLevel: string;
|
||||
model: { provider: string; modelId: string } | null;
|
||||
}
|
||||
|
||||
export interface SessionTreeStorage {
|
||||
getLeafId(): Promise<string | null>;
|
||||
setLeafId(leafId: string | null): Promise<void>;
|
||||
appendEntry(entry: SessionTreeEntry): Promise<void>;
|
||||
getEntry(id: string): Promise<SessionTreeEntry | undefined>;
|
||||
getPathToRoot(leafId: string | null): Promise<SessionTreeEntry[]>;
|
||||
getEntries(): Promise<SessionTreeEntry[]>;
|
||||
}
|
||||
|
||||
export interface SessionTree {
|
||||
getLeafId(): Promise<string | null>;
|
||||
getEntry(id: string): Promise<SessionTreeEntry | undefined>;
|
||||
getEntries(): Promise<SessionTreeEntry[]>;
|
||||
getBranch(fromId?: string): Promise<SessionTreeEntry[]>;
|
||||
buildContext(): Promise<SessionContext>;
|
||||
|
||||
appendMessage(message: AgentMessage): Promise<string>;
|
||||
appendThinkingLevelChange(thinkingLevel: string): Promise<string>;
|
||||
appendModelChange(provider: string, modelId: string): Promise<string>;
|
||||
appendCompaction<T = unknown>(
|
||||
summary: string,
|
||||
firstKeptEntryId: string,
|
||||
tokensBefore: number,
|
||||
details?: T,
|
||||
fromHook?: boolean,
|
||||
): Promise<string>;
|
||||
appendBranchSummary<T = unknown>(fromId: string, summary: string, details?: T, fromHook?: boolean): Promise<string>;
|
||||
appendCustomEntry(customType: string, data?: unknown): Promise<string>;
|
||||
appendCustomMessageEntry<T = unknown>(
|
||||
customType: string,
|
||||
content: string | (TextContent | ImageContent)[],
|
||||
display: boolean,
|
||||
details?: T,
|
||||
): Promise<string>;
|
||||
appendLabelChange(targetId: string, label: string | undefined): Promise<string>;
|
||||
appendSessionInfo(name: string): Promise<string>;
|
||||
|
||||
moveTo(entryId: string | null): Promise<void>;
|
||||
}
|
||||
|
||||
export interface AgentHarnessPendingMutations {
|
||||
appendMessages: AgentMessage[];
|
||||
model?: Model<any>;
|
||||
thinkingLevel?: ThinkingLevel;
|
||||
activeToolNames?: string[];
|
||||
systemPromptInputs?: SystemPromptInputs;
|
||||
}
|
||||
|
||||
export interface AgentHarnessConversationState {
|
||||
sessionTree: SessionTree;
|
||||
model: Model<any> | undefined;
|
||||
thinkingLevel: ThinkingLevel;
|
||||
activeToolNames: string[];
|
||||
systemPromptInputs: SystemPromptInputs;
|
||||
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: Readonly<AgentHarnessConversationState>;
|
||||
operation: Readonly<AgentHarnessOperationState>;
|
||||
abortSignal?: AbortSignal;
|
||||
}
|
||||
|
||||
export interface QueueUpdateEvent {
|
||||
type: "queue_update";
|
||||
steer: AgentMessage[];
|
||||
followUp: AgentMessage[];
|
||||
nextTurn: AgentMessage[];
|
||||
}
|
||||
|
||||
export interface SavePointEvent {
|
||||
type: "save_point";
|
||||
liveOperationId: string;
|
||||
hadPendingMutations: boolean;
|
||||
}
|
||||
|
||||
export interface AbortEvent {
|
||||
type: "abort";
|
||||
clearedSteer: AgentMessage[];
|
||||
clearedFollowUp: AgentMessage[];
|
||||
}
|
||||
|
||||
export interface SettledEvent {
|
||||
type: "settled";
|
||||
nextTurnCount: number;
|
||||
}
|
||||
|
||||
export interface BeforeAgentStartEvent {
|
||||
type: "before_agent_start";
|
||||
prompt: string;
|
||||
images?: ImageContent[];
|
||||
systemPrompt: string;
|
||||
systemPromptInputs: SystemPromptInputs;
|
||||
}
|
||||
|
||||
export interface ContextEvent {
|
||||
type: "context";
|
||||
messages: AgentMessage[];
|
||||
}
|
||||
|
||||
export interface BeforeProviderRequestEvent {
|
||||
type: "before_provider_request";
|
||||
payload: unknown;
|
||||
}
|
||||
|
||||
export interface AfterProviderResponseEvent {
|
||||
type: "after_provider_response";
|
||||
status: number;
|
||||
headers: Record<string, string>;
|
||||
}
|
||||
|
||||
export interface ToolCallEvent {
|
||||
type: "tool_call";
|
||||
toolCallId: string;
|
||||
toolName: string;
|
||||
input: Record<string, unknown>;
|
||||
}
|
||||
|
||||
export interface ToolResultEvent {
|
||||
type: "tool_result";
|
||||
toolCallId: string;
|
||||
toolName: string;
|
||||
input: Record<string, unknown>;
|
||||
content: Array<TextContent | ImageContent>;
|
||||
details: unknown;
|
||||
isError: boolean;
|
||||
}
|
||||
|
||||
export interface SessionBeforeCompactEvent {
|
||||
type: "session_before_compact";
|
||||
preparation: CompactionPreparation;
|
||||
branchEntries: SessionTreeEntry[];
|
||||
customInstructions?: string;
|
||||
signal: AbortSignal;
|
||||
}
|
||||
|
||||
export interface SessionCompactEvent {
|
||||
type: "session_compact";
|
||||
compactionEntry: CompactionEntry;
|
||||
fromHook: boolean;
|
||||
}
|
||||
|
||||
export interface SessionBeforeTreeEvent {
|
||||
type: "session_before_tree";
|
||||
preparation: TreePreparation;
|
||||
signal: AbortSignal;
|
||||
}
|
||||
|
||||
export interface SessionTreeEvent {
|
||||
type: "session_tree";
|
||||
newLeafId: string | null;
|
||||
oldLeafId: string | null;
|
||||
summaryEntry?: BranchSummaryEntry;
|
||||
fromHook?: boolean;
|
||||
}
|
||||
|
||||
export interface ModelSelectEvent {
|
||||
type: "model_select";
|
||||
model: Model<any>;
|
||||
previousModel: Model<any> | undefined;
|
||||
source: "set" | "restore";
|
||||
}
|
||||
|
||||
export interface ThinkingLevelSelectEvent {
|
||||
type: "thinking_level_select";
|
||||
level: ThinkingLevel;
|
||||
previousLevel: ThinkingLevel;
|
||||
}
|
||||
|
||||
export type AgentHarnessOwnEvent =
|
||||
| QueueUpdateEvent
|
||||
| SavePointEvent
|
||||
| AbortEvent
|
||||
| SettledEvent
|
||||
| BeforeAgentStartEvent
|
||||
| ContextEvent
|
||||
| BeforeProviderRequestEvent
|
||||
| AfterProviderResponseEvent
|
||||
| ToolCallEvent
|
||||
| ToolResultEvent
|
||||
| SessionBeforeCompactEvent
|
||||
| SessionCompactEvent
|
||||
| SessionBeforeTreeEvent
|
||||
| SessionTreeEvent
|
||||
| ModelSelectEvent
|
||||
| ThinkingLevelSelectEvent;
|
||||
|
||||
export type AgentHarnessEvent = AgentEvent | AgentHarnessOwnEvent;
|
||||
|
||||
export interface BeforeAgentStartResult {
|
||||
messages?: AgentMessage[];
|
||||
systemPrompt?: string;
|
||||
}
|
||||
|
||||
export interface ContextResult {
|
||||
messages: AgentMessage[];
|
||||
}
|
||||
|
||||
export interface BeforeProviderRequestResult {
|
||||
payload: unknown;
|
||||
}
|
||||
|
||||
export interface ToolCallResult {
|
||||
block?: boolean;
|
||||
reason?: string;
|
||||
}
|
||||
|
||||
export interface ToolResultPatch {
|
||||
content?: Array<TextContent | ImageContent>;
|
||||
details?: unknown;
|
||||
isError?: boolean;
|
||||
terminate?: boolean;
|
||||
}
|
||||
|
||||
export interface SessionBeforeCompactResult {
|
||||
cancel?: boolean;
|
||||
compaction?: CompactResult;
|
||||
}
|
||||
|
||||
export interface SessionBeforeTreeResult {
|
||||
cancel?: boolean;
|
||||
summary?: { summary: string; details?: unknown };
|
||||
customInstructions?: string;
|
||||
replaceInstructions?: boolean;
|
||||
label?: string;
|
||||
}
|
||||
|
||||
export type AgentHarnessEventResultMap = {
|
||||
before_agent_start: BeforeAgentStartResult | undefined;
|
||||
context: ContextResult | undefined;
|
||||
before_provider_request: BeforeProviderRequestResult | undefined;
|
||||
after_provider_response: undefined;
|
||||
tool_call: ToolCallResult | undefined;
|
||||
tool_result: ToolResultPatch | undefined;
|
||||
session_before_compact: SessionBeforeCompactResult | undefined;
|
||||
session_compact: undefined;
|
||||
session_before_tree: SessionBeforeTreeResult | undefined;
|
||||
session_tree: undefined;
|
||||
model_select: undefined;
|
||||
thinking_level_select: undefined;
|
||||
queue_update: undefined;
|
||||
save_point: undefined;
|
||||
abort: undefined;
|
||||
settled: undefined;
|
||||
};
|
||||
|
||||
export interface AgentHarnessPromptOptions {
|
||||
images?: ImageContent[];
|
||||
}
|
||||
|
||||
export interface AbortResult {
|
||||
clearedSteer: AgentMessage[];
|
||||
clearedFollowUp: AgentMessage[];
|
||||
}
|
||||
|
||||
export interface CompactResult {
|
||||
summary: string;
|
||||
firstKeptEntryId: string;
|
||||
tokensBefore: number;
|
||||
details?: unknown;
|
||||
}
|
||||
|
||||
export interface NavigateTreeResult {
|
||||
cancelled: boolean;
|
||||
editorText?: string;
|
||||
summaryEntry?: BranchSummaryEntry;
|
||||
}
|
||||
|
||||
export interface CompactionSettings {
|
||||
enabled: boolean;
|
||||
reserveTokens: number;
|
||||
keepRecentTokens: number;
|
||||
}
|
||||
|
||||
export interface CompactionPreparation {
|
||||
firstKeptEntryId: string;
|
||||
messagesToSummarize: AgentMessage[];
|
||||
turnPrefixMessages: AgentMessage[];
|
||||
isSplitTurn: boolean;
|
||||
tokensBefore: number;
|
||||
previousSummary?: string;
|
||||
fileOps: FileOperations;
|
||||
settings: CompactionSettings;
|
||||
}
|
||||
|
||||
export interface FileOperations {
|
||||
read: Set<string>;
|
||||
written: Set<string>;
|
||||
edited: Set<string>;
|
||||
}
|
||||
|
||||
export interface TreePreparation {
|
||||
targetId: string;
|
||||
oldLeafId: string | null;
|
||||
commonAncestorId: string | null;
|
||||
entriesToSummarize: SessionTreeEntry[];
|
||||
userWantsSummary: boolean;
|
||||
customInstructions?: string;
|
||||
replaceInstructions?: boolean;
|
||||
label?: string;
|
||||
}
|
||||
|
||||
export interface GenerateBranchSummaryOptions {
|
||||
model: Model<any>;
|
||||
apiKey: string;
|
||||
headers?: Record<string, string>;
|
||||
signal: AbortSignal;
|
||||
customInstructions?: string;
|
||||
replaceInstructions?: boolean;
|
||||
reserveTokens?: number;
|
||||
}
|
||||
|
||||
export interface BranchSummaryResult {
|
||||
summary?: string;
|
||||
readFiles?: string[];
|
||||
modifiedFiles?: string[];
|
||||
aborted?: boolean;
|
||||
error?: string;
|
||||
}
|
||||
|
||||
export interface AgentHarnessOptions {
|
||||
agent: Agent;
|
||||
env: ExecutionEnv;
|
||||
sessionTree: SessionTree;
|
||||
promptTemplates?: PromptTemplate[];
|
||||
skills?: Skill[];
|
||||
requestAuth?: (model: Model<any>) => Promise<{ apiKey: string; headers?: Record<string, string> } | undefined>;
|
||||
initialModel?: Model<any>;
|
||||
initialThinkingLevel?: ThinkingLevel;
|
||||
initialActiveToolNames?: string[];
|
||||
initialSystemPromptInputs?: SystemPromptInputs;
|
||||
}
|
||||
|
||||
export interface AgentHarness {
|
||||
readonly agent: Agent;
|
||||
readonly env: ExecutionEnv;
|
||||
readonly conversation: Readonly<AgentHarnessConversationState>;
|
||||
readonly operation: Readonly<AgentHarnessOperationState>;
|
||||
|
||||
prompt(text: string, options?: AgentHarnessPromptOptions): Promise<void>;
|
||||
skill(name: string, args?: string): Promise<void>;
|
||||
|
||||
steer(message: AgentMessage): void;
|
||||
followUp(message: AgentMessage): void;
|
||||
nextTurn(message: AgentMessage): void;
|
||||
|
||||
appendMessage(message: AgentMessage): Promise<void>;
|
||||
|
||||
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 }>;
|
||||
|
||||
compact(customInstructions?: string): Promise<CompactResult>;
|
||||
navigateTree(
|
||||
targetId: string,
|
||||
options?: { summarize?: boolean; customInstructions?: string; replaceInstructions?: boolean; label?: string },
|
||||
): Promise<NavigateTreeResult>;
|
||||
|
||||
setModel(model: Model<any>): Promise<void>;
|
||||
setThinkingLevel(level: ThinkingLevel): Promise<void>;
|
||||
setActiveTools(toolNames: string[]): Promise<void>;
|
||||
setSystemPromptInputs(inputs: SystemPromptInputs): Promise<void>;
|
||||
|
||||
abort(): Promise<AbortResult>;
|
||||
waitForIdle(): Promise<void>;
|
||||
|
||||
subscribe(listener: (event: AgentHarnessEvent, signal?: AbortSignal) => Promise<void> | void): () => void;
|
||||
|
||||
on<TType extends keyof AgentHarnessEventResultMap>(
|
||||
type: TType,
|
||||
handler: (
|
||||
event: Extract<AgentHarnessEvent, { type: TType }>,
|
||||
ctx: AgentHarnessContext,
|
||||
) => Promise<AgentHarnessEventResultMap[TType]> | AgentHarnessEventResultMap[TType],
|
||||
): () => void;
|
||||
}
|
||||
113
packages/agent/src/harness/utils/shell-output.ts
Normal file
113
packages/agent/src/harness/utils/shell-output.ts
Normal file
@@ -0,0 +1,113 @@
|
||||
import { randomBytes } from "node:crypto";
|
||||
import { createWriteStream, type WriteStream } from "node:fs";
|
||||
import { tmpdir } from "node:os";
|
||||
import { join } from "node:path";
|
||||
import type { ExecutionEnv, ExecutionEnvExecOptions } from "../types.js";
|
||||
import { DEFAULT_MAX_BYTES, truncateTail } from "./truncate.js";
|
||||
|
||||
export interface ShellCaptureOptions extends Omit<ExecutionEnvExecOptions, "onStdout" | "onStderr"> {
|
||||
onChunk?: (chunk: string) => void;
|
||||
}
|
||||
|
||||
export interface ShellCaptureResult {
|
||||
output: string;
|
||||
exitCode: number | undefined;
|
||||
cancelled: boolean;
|
||||
truncated: boolean;
|
||||
fullOutputPath?: string;
|
||||
}
|
||||
|
||||
export function sanitizeBinaryOutput(str: string): string {
|
||||
return Array.from(str)
|
||||
.filter((char) => {
|
||||
const code = char.codePointAt(0);
|
||||
if (code === undefined) return false;
|
||||
if (code === 0x09 || code === 0x0a || code === 0x0d) return true;
|
||||
if (code <= 0x1f) return false;
|
||||
if (code >= 0xfff9 && code <= 0xfffb) return false;
|
||||
return true;
|
||||
})
|
||||
.join("");
|
||||
}
|
||||
|
||||
export async function executeShellWithCapture(
|
||||
env: ExecutionEnv,
|
||||
command: string,
|
||||
options?: ShellCaptureOptions,
|
||||
): Promise<ShellCaptureResult> {
|
||||
const outputChunks: string[] = [];
|
||||
let outputBytes = 0;
|
||||
const maxOutputBytes = DEFAULT_MAX_BYTES * 2;
|
||||
|
||||
let tempFilePath: string | undefined;
|
||||
let tempFileStream: WriteStream | undefined;
|
||||
let totalBytes = 0;
|
||||
|
||||
const ensureTempFile = () => {
|
||||
if (tempFilePath) return;
|
||||
const id = randomBytes(8).toString("hex");
|
||||
tempFilePath = join(tmpdir(), `bash-${id}.log`);
|
||||
tempFileStream = createWriteStream(tempFilePath);
|
||||
for (const chunk of outputChunks) {
|
||||
tempFileStream.write(chunk);
|
||||
}
|
||||
};
|
||||
|
||||
const onChunk = (chunk: string) => {
|
||||
totalBytes += Buffer.byteLength(chunk, "utf-8");
|
||||
const text = sanitizeBinaryOutput(chunk).replace(/\r/g, "");
|
||||
if (totalBytes > DEFAULT_MAX_BYTES) {
|
||||
ensureTempFile();
|
||||
}
|
||||
if (tempFileStream) {
|
||||
tempFileStream.write(text);
|
||||
}
|
||||
outputChunks.push(text);
|
||||
outputBytes += text.length;
|
||||
while (outputBytes > maxOutputBytes && outputChunks.length > 1) {
|
||||
const removed = outputChunks.shift()!;
|
||||
outputBytes -= removed.length;
|
||||
}
|
||||
options?.onChunk?.(text);
|
||||
};
|
||||
|
||||
try {
|
||||
const result = await env.exec(command, {
|
||||
...(options ?? {}),
|
||||
onStdout: onChunk,
|
||||
onStderr: onChunk,
|
||||
});
|
||||
const fullOutput = outputChunks.join("");
|
||||
const truncationResult = truncateTail(fullOutput);
|
||||
if (truncationResult.truncated) {
|
||||
ensureTempFile();
|
||||
}
|
||||
tempFileStream?.end();
|
||||
const cancelled = options?.signal?.aborted ?? false;
|
||||
return {
|
||||
output: truncationResult.truncated ? truncationResult.content : fullOutput,
|
||||
exitCode: cancelled ? undefined : result.exitCode,
|
||||
cancelled,
|
||||
truncated: truncationResult.truncated,
|
||||
fullOutputPath: tempFilePath,
|
||||
};
|
||||
} catch (err) {
|
||||
if (options?.signal?.aborted) {
|
||||
const fullOutput = outputChunks.join("");
|
||||
const truncationResult = truncateTail(fullOutput);
|
||||
if (truncationResult.truncated) {
|
||||
ensureTempFile();
|
||||
}
|
||||
tempFileStream?.end();
|
||||
return {
|
||||
output: truncationResult.truncated ? truncationResult.content : fullOutput,
|
||||
exitCode: undefined,
|
||||
cancelled: true,
|
||||
truncated: truncationResult.truncated,
|
||||
fullOutputPath: tempFilePath,
|
||||
};
|
||||
}
|
||||
tempFileStream?.end();
|
||||
throw err;
|
||||
}
|
||||
}
|
||||
265
packages/agent/src/harness/utils/truncate.ts
Normal file
265
packages/agent/src/harness/utils/truncate.ts
Normal file
@@ -0,0 +1,265 @@
|
||||
/**
|
||||
* Shared truncation utilities for tool outputs.
|
||||
*
|
||||
* Truncation is based on two independent limits - whichever is hit first wins:
|
||||
* - Line limit (default: 2000 lines)
|
||||
* - Byte limit (default: 50KB)
|
||||
*
|
||||
* Never returns partial lines (except bash tail truncation edge case).
|
||||
*/
|
||||
|
||||
export const DEFAULT_MAX_LINES = 2000;
|
||||
export const DEFAULT_MAX_BYTES = 50 * 1024; // 50KB
|
||||
export const GREP_MAX_LINE_LENGTH = 500; // Max chars per grep match line
|
||||
|
||||
export interface TruncationResult {
|
||||
/** The truncated content */
|
||||
content: string;
|
||||
/** Whether truncation occurred */
|
||||
truncated: boolean;
|
||||
/** Which limit was hit: "lines", "bytes", or null if not truncated */
|
||||
truncatedBy: "lines" | "bytes" | null;
|
||||
/** Total number of lines in the original content */
|
||||
totalLines: number;
|
||||
/** Total number of bytes in the original content */
|
||||
totalBytes: number;
|
||||
/** Number of complete lines in the truncated output */
|
||||
outputLines: number;
|
||||
/** Number of bytes in the truncated output */
|
||||
outputBytes: number;
|
||||
/** Whether the last line was partially truncated (only for tail truncation edge case) */
|
||||
lastLinePartial: boolean;
|
||||
/** Whether the first line exceeded the byte limit (for head truncation) */
|
||||
firstLineExceedsLimit: boolean;
|
||||
/** The max lines limit that was applied */
|
||||
maxLines: number;
|
||||
/** The max bytes limit that was applied */
|
||||
maxBytes: number;
|
||||
}
|
||||
|
||||
export interface TruncationOptions {
|
||||
/** Maximum number of lines (default: 2000) */
|
||||
maxLines?: number;
|
||||
/** Maximum number of bytes (default: 50KB) */
|
||||
maxBytes?: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Format bytes as human-readable size.
|
||||
*/
|
||||
export function formatSize(bytes: number): string {
|
||||
if (bytes < 1024) {
|
||||
return `${bytes}B`;
|
||||
} else if (bytes < 1024 * 1024) {
|
||||
return `${(bytes / 1024).toFixed(1)}KB`;
|
||||
} else {
|
||||
return `${(bytes / (1024 * 1024)).toFixed(1)}MB`;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Truncate content from the head (keep first N lines/bytes).
|
||||
* Suitable for file reads where you want to see the beginning.
|
||||
*
|
||||
* Never returns partial lines. If first line exceeds byte limit,
|
||||
* returns empty content with firstLineExceedsLimit=true.
|
||||
*/
|
||||
export function truncateHead(content: string, options: TruncationOptions = {}): TruncationResult {
|
||||
const maxLines = options.maxLines ?? DEFAULT_MAX_LINES;
|
||||
const maxBytes = options.maxBytes ?? DEFAULT_MAX_BYTES;
|
||||
|
||||
const totalBytes = Buffer.byteLength(content, "utf-8");
|
||||
const lines = content.split("\n");
|
||||
const totalLines = lines.length;
|
||||
|
||||
// Check if no truncation needed
|
||||
if (totalLines <= maxLines && totalBytes <= maxBytes) {
|
||||
return {
|
||||
content,
|
||||
truncated: false,
|
||||
truncatedBy: null,
|
||||
totalLines,
|
||||
totalBytes,
|
||||
outputLines: totalLines,
|
||||
outputBytes: totalBytes,
|
||||
lastLinePartial: false,
|
||||
firstLineExceedsLimit: false,
|
||||
maxLines,
|
||||
maxBytes,
|
||||
};
|
||||
}
|
||||
|
||||
// Check if first line alone exceeds byte limit
|
||||
const firstLineBytes = Buffer.byteLength(lines[0], "utf-8");
|
||||
if (firstLineBytes > maxBytes) {
|
||||
return {
|
||||
content: "",
|
||||
truncated: true,
|
||||
truncatedBy: "bytes",
|
||||
totalLines,
|
||||
totalBytes,
|
||||
outputLines: 0,
|
||||
outputBytes: 0,
|
||||
lastLinePartial: false,
|
||||
firstLineExceedsLimit: true,
|
||||
maxLines,
|
||||
maxBytes,
|
||||
};
|
||||
}
|
||||
|
||||
// Collect complete lines that fit
|
||||
const outputLinesArr: string[] = [];
|
||||
let outputBytesCount = 0;
|
||||
let truncatedBy: "lines" | "bytes" = "lines";
|
||||
|
||||
for (let i = 0; i < lines.length && i < maxLines; i++) {
|
||||
const line = lines[i];
|
||||
const lineBytes = Buffer.byteLength(line, "utf-8") + (i > 0 ? 1 : 0); // +1 for newline
|
||||
|
||||
if (outputBytesCount + lineBytes > maxBytes) {
|
||||
truncatedBy = "bytes";
|
||||
break;
|
||||
}
|
||||
|
||||
outputLinesArr.push(line);
|
||||
outputBytesCount += lineBytes;
|
||||
}
|
||||
|
||||
// If we exited due to line limit
|
||||
if (outputLinesArr.length >= maxLines && outputBytesCount <= maxBytes) {
|
||||
truncatedBy = "lines";
|
||||
}
|
||||
|
||||
const outputContent = outputLinesArr.join("\n");
|
||||
const finalOutputBytes = Buffer.byteLength(outputContent, "utf-8");
|
||||
|
||||
return {
|
||||
content: outputContent,
|
||||
truncated: true,
|
||||
truncatedBy,
|
||||
totalLines,
|
||||
totalBytes,
|
||||
outputLines: outputLinesArr.length,
|
||||
outputBytes: finalOutputBytes,
|
||||
lastLinePartial: false,
|
||||
firstLineExceedsLimit: false,
|
||||
maxLines,
|
||||
maxBytes,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Truncate content from the tail (keep last N lines/bytes).
|
||||
* Suitable for bash output where you want to see the end (errors, final results).
|
||||
*
|
||||
* May return partial first line if the last line of original content exceeds byte limit.
|
||||
*/
|
||||
export function truncateTail(content: string, options: TruncationOptions = {}): TruncationResult {
|
||||
const maxLines = options.maxLines ?? DEFAULT_MAX_LINES;
|
||||
const maxBytes = options.maxBytes ?? DEFAULT_MAX_BYTES;
|
||||
|
||||
const totalBytes = Buffer.byteLength(content, "utf-8");
|
||||
const lines = content.split("\n");
|
||||
const totalLines = lines.length;
|
||||
|
||||
// Check if no truncation needed
|
||||
if (totalLines <= maxLines && totalBytes <= maxBytes) {
|
||||
return {
|
||||
content,
|
||||
truncated: false,
|
||||
truncatedBy: null,
|
||||
totalLines,
|
||||
totalBytes,
|
||||
outputLines: totalLines,
|
||||
outputBytes: totalBytes,
|
||||
lastLinePartial: false,
|
||||
firstLineExceedsLimit: false,
|
||||
maxLines,
|
||||
maxBytes,
|
||||
};
|
||||
}
|
||||
|
||||
// Work backwards from the end
|
||||
const outputLinesArr: string[] = [];
|
||||
let outputBytesCount = 0;
|
||||
let truncatedBy: "lines" | "bytes" = "lines";
|
||||
let lastLinePartial = false;
|
||||
|
||||
for (let i = lines.length - 1; i >= 0 && outputLinesArr.length < maxLines; i--) {
|
||||
const line = lines[i];
|
||||
const lineBytes = Buffer.byteLength(line, "utf-8") + (outputLinesArr.length > 0 ? 1 : 0); // +1 for newline
|
||||
|
||||
if (outputBytesCount + lineBytes > maxBytes) {
|
||||
truncatedBy = "bytes";
|
||||
// Edge case: if we haven't added ANY lines yet and this line exceeds maxBytes,
|
||||
// take the end of the line (partial)
|
||||
if (outputLinesArr.length === 0) {
|
||||
const truncatedLine = truncateStringToBytesFromEnd(line, maxBytes);
|
||||
outputLinesArr.unshift(truncatedLine);
|
||||
outputBytesCount = Buffer.byteLength(truncatedLine, "utf-8");
|
||||
lastLinePartial = true;
|
||||
}
|
||||
break;
|
||||
}
|
||||
|
||||
outputLinesArr.unshift(line);
|
||||
outputBytesCount += lineBytes;
|
||||
}
|
||||
|
||||
// If we exited due to line limit
|
||||
if (outputLinesArr.length >= maxLines && outputBytesCount <= maxBytes) {
|
||||
truncatedBy = "lines";
|
||||
}
|
||||
|
||||
const outputContent = outputLinesArr.join("\n");
|
||||
const finalOutputBytes = Buffer.byteLength(outputContent, "utf-8");
|
||||
|
||||
return {
|
||||
content: outputContent,
|
||||
truncated: true,
|
||||
truncatedBy,
|
||||
totalLines,
|
||||
totalBytes,
|
||||
outputLines: outputLinesArr.length,
|
||||
outputBytes: finalOutputBytes,
|
||||
lastLinePartial,
|
||||
firstLineExceedsLimit: false,
|
||||
maxLines,
|
||||
maxBytes,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Truncate a string to fit within a byte limit (from the end).
|
||||
* Handles multi-byte UTF-8 characters correctly.
|
||||
*/
|
||||
function truncateStringToBytesFromEnd(str: string, maxBytes: number): string {
|
||||
const buf = Buffer.from(str, "utf-8");
|
||||
if (buf.length <= maxBytes) {
|
||||
return str;
|
||||
}
|
||||
|
||||
// Start from the end, skip maxBytes back
|
||||
let start = buf.length - maxBytes;
|
||||
|
||||
// Find a valid UTF-8 boundary (start of a character)
|
||||
while (start < buf.length && (buf[start] & 0xc0) === 0x80) {
|
||||
start++;
|
||||
}
|
||||
|
||||
return buf.slice(start).toString("utf-8");
|
||||
}
|
||||
|
||||
/**
|
||||
* Truncate a single line to max characters, adding [truncated] suffix.
|
||||
* Used for grep match lines.
|
||||
*/
|
||||
export function truncateLine(
|
||||
line: string,
|
||||
maxChars: number = GREP_MAX_LINE_LENGTH,
|
||||
): { text: string; wasTruncated: boolean } {
|
||||
if (line.length <= maxChars) {
|
||||
return { text: line, wasTruncated: false };
|
||||
}
|
||||
return { text: `${line.slice(0, maxChars)}... [truncated]`, wasTruncated: true };
|
||||
}
|
||||
@@ -2,6 +2,16 @@
|
||||
export * from "./agent.js";
|
||||
// Loop functions
|
||||
export * from "./agent-loop.js";
|
||||
export * from "./harness/agent-harness.js";
|
||||
export * from "./harness/compaction.js";
|
||||
export * from "./harness/execution-env.js";
|
||||
export * from "./harness/messages.js";
|
||||
export * from "./harness/prompt-templates.js";
|
||||
export * from "./harness/session-tree.js";
|
||||
// Harness
|
||||
export * from "./harness/types.js";
|
||||
export * from "./harness/utils/shell-output.js";
|
||||
export * from "./harness/utils/truncate.js";
|
||||
// Proxy utilities
|
||||
export * from "./proxy.js";
|
||||
// Types
|
||||
|
||||
310
packages/agent/test/harness/compaction.test.ts
Normal file
310
packages/agent/test/harness/compaction.test.ts
Normal file
@@ -0,0 +1,310 @@
|
||||
import type { AgentMessage } from "@mariozechner/pi-agent-core";
|
||||
import {
|
||||
type AssistantMessage,
|
||||
type FauxProviderRegistration,
|
||||
fauxAssistantMessage,
|
||||
type Model,
|
||||
registerFauxProvider,
|
||||
type Usage,
|
||||
} from "@mariozechner/pi-ai";
|
||||
import { afterEach, beforeEach, describe, expect, it } from "vitest";
|
||||
import {
|
||||
calculateContextTokens,
|
||||
compact,
|
||||
DEFAULT_COMPACTION_SETTINGS,
|
||||
estimateContextTokens,
|
||||
findCutPoint,
|
||||
generateSummary,
|
||||
prepareCompaction,
|
||||
serializeConversation,
|
||||
shouldCompact,
|
||||
} from "../../src/harness/compaction.js";
|
||||
import { buildSessionContext } from "../../src/harness/session-tree.js";
|
||||
import type {
|
||||
CompactionEntry,
|
||||
CompactionSettings,
|
||||
MessageEntry,
|
||||
ModelChangeEntry,
|
||||
SessionTreeEntry,
|
||||
ThinkingLevelChangeEntry,
|
||||
} from "../../src/harness/types.js";
|
||||
|
||||
let nextId = 0;
|
||||
function createId(): string {
|
||||
return `entry-${nextId++}`;
|
||||
}
|
||||
|
||||
function createMockUsage(input: number, output: number, cacheRead = 0, cacheWrite = 0): Usage {
|
||||
return {
|
||||
input,
|
||||
output,
|
||||
cacheRead,
|
||||
cacheWrite,
|
||||
totalTokens: input + output + cacheRead + cacheWrite,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
};
|
||||
}
|
||||
|
||||
function createUserMessage(text: string): AgentMessage {
|
||||
return {
|
||||
role: "user",
|
||||
content: [{ type: "text", text }],
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
}
|
||||
|
||||
function createAssistantMessage(text: string, usage = createMockUsage(100, 50)): AssistantMessage {
|
||||
return {
|
||||
role: "assistant",
|
||||
content: [{ type: "text", text }],
|
||||
api: "anthropic-messages",
|
||||
provider: "anthropic",
|
||||
model: "claude-sonnet-4-5",
|
||||
usage,
|
||||
stopReason: "stop",
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
}
|
||||
|
||||
function createMessageEntry(message: AgentMessage, parentId: string | null = null): MessageEntry {
|
||||
return {
|
||||
type: "message",
|
||||
id: createId(),
|
||||
parentId,
|
||||
timestamp: new Date().toISOString(),
|
||||
message,
|
||||
};
|
||||
}
|
||||
|
||||
function createCompactionEntry(
|
||||
summary: string,
|
||||
firstKeptEntryId: string,
|
||||
parentId: string | null = null,
|
||||
): CompactionEntry {
|
||||
return {
|
||||
type: "compaction",
|
||||
id: createId(),
|
||||
parentId,
|
||||
timestamp: new Date().toISOString(),
|
||||
summary,
|
||||
firstKeptEntryId,
|
||||
tokensBefore: 1234,
|
||||
};
|
||||
}
|
||||
|
||||
function createThinkingLevelEntry(level: string, parentId: string | null = null): ThinkingLevelChangeEntry {
|
||||
return {
|
||||
type: "thinking_level_change",
|
||||
id: createId(),
|
||||
parentId,
|
||||
timestamp: new Date().toISOString(),
|
||||
thinkingLevel: level,
|
||||
};
|
||||
}
|
||||
|
||||
function createModelChangeEntry(provider: string, modelId: string, parentId: string | null = null): ModelChangeEntry {
|
||||
return {
|
||||
type: "model_change",
|
||||
id: createId(),
|
||||
parentId,
|
||||
timestamp: new Date().toISOString(),
|
||||
provider,
|
||||
modelId,
|
||||
};
|
||||
}
|
||||
|
||||
function createFauxModel(reasoning: boolean): { faux: FauxProviderRegistration; model: Model<string> } {
|
||||
const faux = registerFauxProvider({
|
||||
models: [
|
||||
{
|
||||
id: reasoning ? "reasoning-model" : "non-reasoning-model",
|
||||
reasoning,
|
||||
contextWindow: 200000,
|
||||
maxTokens: 8192,
|
||||
},
|
||||
],
|
||||
});
|
||||
fauxRegistrations.push(faux);
|
||||
return { faux, model: faux.getModel() };
|
||||
}
|
||||
|
||||
const fauxRegistrations: FauxProviderRegistration[] = [];
|
||||
|
||||
afterEach(() => {
|
||||
while (fauxRegistrations.length > 0) {
|
||||
fauxRegistrations.pop()?.unregister();
|
||||
}
|
||||
});
|
||||
|
||||
describe("harness compaction", () => {
|
||||
beforeEach(() => {
|
||||
nextId = 0;
|
||||
});
|
||||
|
||||
it("calculates total context tokens from usage", () => {
|
||||
expect(calculateContextTokens(createMockUsage(1000, 500, 200, 100))).toBe(1800);
|
||||
expect(calculateContextTokens(createMockUsage(0, 0, 0, 0))).toBe(0);
|
||||
});
|
||||
|
||||
it("checks compaction threshold", () => {
|
||||
const settings: CompactionSettings = {
|
||||
enabled: true,
|
||||
reserveTokens: 10000,
|
||||
keepRecentTokens: 20000,
|
||||
};
|
||||
expect(shouldCompact(95000, 100000, settings)).toBe(true);
|
||||
expect(shouldCompact(89000, 100000, settings)).toBe(false);
|
||||
expect(shouldCompact(95000, 100000, { ...settings, enabled: false })).toBe(false);
|
||||
});
|
||||
|
||||
it("finds a cut point based on token differences", () => {
|
||||
const entries: SessionTreeEntry[] = [];
|
||||
let parentId: string | null = null;
|
||||
for (let i = 0; i < 10; i++) {
|
||||
const user = createMessageEntry(createUserMessage(`User ${i}`), parentId);
|
||||
entries.push(user);
|
||||
const assistant = createMessageEntry(
|
||||
createAssistantMessage(`Assistant ${i}`, createMockUsage(0, 100, (i + 1) * 1000, 0)),
|
||||
user.id,
|
||||
);
|
||||
entries.push(assistant);
|
||||
parentId = assistant.id;
|
||||
}
|
||||
|
||||
const result = findCutPoint(entries, 0, entries.length, 2500);
|
||||
expect(entries[result.firstKeptEntryIndex]?.type).toBe("message");
|
||||
});
|
||||
|
||||
it("builds session context with a compaction entry", () => {
|
||||
const u1 = createMessageEntry(createUserMessage("1"));
|
||||
const a1 = createMessageEntry(createAssistantMessage("a"), u1.id);
|
||||
const u2 = createMessageEntry(createUserMessage("2"), a1.id);
|
||||
const a2 = createMessageEntry(createAssistantMessage("b"), u2.id);
|
||||
const compaction = createCompactionEntry("Summary of 1,a,2,b", u2.id, a2.id);
|
||||
const u3 = createMessageEntry(createUserMessage("3"), compaction.id);
|
||||
const a3 = createMessageEntry(createAssistantMessage("c"), u3.id);
|
||||
const loaded = buildSessionContext([u1, a1, u2, a2, compaction, u3, a3]);
|
||||
expect(loaded.messages).toHaveLength(5);
|
||||
expect(loaded.messages[0]?.role).toBe("compactionSummary");
|
||||
});
|
||||
|
||||
it("tracks model and thinking level changes in built context", () => {
|
||||
const user = createMessageEntry(createUserMessage("1"));
|
||||
const modelChange = createModelChangeEntry("openai", "gpt-4", user.id);
|
||||
const assistant = createMessageEntry(createAssistantMessage("a"), modelChange.id);
|
||||
const thinkingChange = createThinkingLevelEntry("high", assistant.id);
|
||||
const loaded = buildSessionContext([user, modelChange, assistant, thinkingChange]);
|
||||
expect(loaded.model).toEqual({ provider: "anthropic", modelId: "claude-sonnet-4-5" });
|
||||
expect(loaded.thinkingLevel).toBe("high");
|
||||
});
|
||||
|
||||
it("prepares compaction using the latest compaction summary as previousSummary", () => {
|
||||
const u1 = createMessageEntry(createUserMessage("user msg 1"));
|
||||
const a1 = createMessageEntry(createAssistantMessage("assistant msg 1"), u1.id);
|
||||
const u2 = createMessageEntry(createUserMessage("user msg 2"), a1.id);
|
||||
const a2 = createMessageEntry(createAssistantMessage("assistant msg 2", createMockUsage(5000, 1000)), u2.id);
|
||||
const compaction1 = createCompactionEntry("First summary", u2.id, a2.id);
|
||||
const u3 = createMessageEntry(createUserMessage("user msg 3"), compaction1.id);
|
||||
const a3 = createMessageEntry(createAssistantMessage("assistant msg 3", createMockUsage(8000, 2000)), u3.id);
|
||||
const pathEntries = [u1, a1, u2, a2, compaction1, u3, a3];
|
||||
const preparation = prepareCompaction(pathEntries, DEFAULT_COMPACTION_SETTINGS);
|
||||
expect(preparation).toBeDefined();
|
||||
expect(preparation?.previousSummary).toBe("First summary");
|
||||
expect(preparation?.firstKeptEntryId).toBeTruthy();
|
||||
expect(preparation?.tokensBefore).toBe(estimateContextTokens(buildSessionContext(pathEntries).messages).tokens);
|
||||
});
|
||||
|
||||
it("serializes conversation with truncated tool results", () => {
|
||||
const longContent = "x".repeat(5000);
|
||||
const messages = convertMessages([
|
||||
{
|
||||
role: "toolResult",
|
||||
toolCallId: "tc1",
|
||||
toolName: "read",
|
||||
content: [{ type: "text", text: longContent }],
|
||||
isError: false,
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
]);
|
||||
const result = serializeConversation(messages);
|
||||
expect(result).toContain("[Tool result]:");
|
||||
expect(result).toContain("[... 3000 more characters truncated]");
|
||||
});
|
||||
|
||||
it("passes reasoning through generateSummary only for reasoning models with thinking enabled", async () => {
|
||||
const messages: AgentMessage[] = [createUserMessage("Summarize this.")];
|
||||
const seenOptions: Array<Record<string, unknown> | undefined> = [];
|
||||
const { faux: fauxReasoning, model: reasoningModel } = createFauxModel(true);
|
||||
fauxReasoning.setResponses([
|
||||
(_context, options) => {
|
||||
seenOptions.push(options as Record<string, unknown> | undefined);
|
||||
return fauxAssistantMessage("## Goal\nTest summary");
|
||||
},
|
||||
]);
|
||||
await generateSummary(
|
||||
messages,
|
||||
reasoningModel,
|
||||
2000,
|
||||
"test-key",
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
"medium",
|
||||
);
|
||||
expect(seenOptions[0]).toMatchObject({ reasoning: "medium", apiKey: "test-key" });
|
||||
|
||||
const { faux: fauxOff, model: offModel } = createFauxModel(true);
|
||||
fauxOff.setResponses([
|
||||
(_context, options) => {
|
||||
seenOptions.push(options as Record<string, unknown> | undefined);
|
||||
return fauxAssistantMessage("## Goal\nTest summary");
|
||||
},
|
||||
]);
|
||||
await generateSummary(messages, offModel, 2000, "test-key", undefined, undefined, undefined, undefined, "off");
|
||||
expect(seenOptions[1]).not.toHaveProperty("reasoning");
|
||||
|
||||
const { faux: fauxNonReasoning, model: nonReasoningModel } = createFauxModel(false);
|
||||
fauxNonReasoning.setResponses([
|
||||
(_context, options) => {
|
||||
seenOptions.push(options as Record<string, unknown> | undefined);
|
||||
return fauxAssistantMessage("## Goal\nTest summary");
|
||||
},
|
||||
]);
|
||||
await generateSummary(
|
||||
messages,
|
||||
nonReasoningModel,
|
||||
2000,
|
||||
"test-key",
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
"medium",
|
||||
);
|
||||
expect(seenOptions[2]).not.toHaveProperty("reasoning");
|
||||
});
|
||||
|
||||
it("returns a compaction result with file details", async () => {
|
||||
const u1 = createMessageEntry(createUserMessage("read a file"));
|
||||
const assistantMessage: AssistantMessage = {
|
||||
...createAssistantMessage("calling tool", createMockUsage(1000, 200)),
|
||||
content: [{ type: "toolCall", id: "tool-1", name: "read", arguments: { path: "src/index.ts" } }],
|
||||
};
|
||||
const a1 = createMessageEntry(assistantMessage, u1.id);
|
||||
const u2 = createMessageEntry(createUserMessage("continue"), a1.id);
|
||||
const a2 = createMessageEntry(createAssistantMessage("done", createMockUsage(4000, 500)), u2.id);
|
||||
const preparation = prepareCompaction([u1, a1, u2, a2], DEFAULT_COMPACTION_SETTINGS);
|
||||
expect(preparation).toBeDefined();
|
||||
const { faux, model } = createFauxModel(false);
|
||||
faux.setResponses([fauxAssistantMessage("## Goal\nTest summary")]);
|
||||
const result = await compact(preparation!, model, "test-key");
|
||||
expect(result.summary.length).toBeGreaterThan(0);
|
||||
expect(result.firstKeptEntryId).toBeTruthy();
|
||||
expect(result.details).toBeDefined();
|
||||
});
|
||||
});
|
||||
|
||||
function convertMessages(messages: any[]): any[] {
|
||||
return messages;
|
||||
}
|
||||
174
packages/agent/test/harness/session-tree.test.ts
Normal file
174
packages/agent/test/harness/session-tree.test.ts
Normal file
@@ -0,0 +1,174 @@
|
||||
import { existsSync, mkdirSync, readFileSync, rmSync } from "node:fs";
|
||||
import { tmpdir } from "node:os";
|
||||
import { join } from "node:path";
|
||||
import type { AgentMessage } from "@mariozechner/pi-agent-core";
|
||||
import { afterEach, describe, expect, it } from "vitest";
|
||||
import {
|
||||
DefaultSessionTree,
|
||||
InMemorySessionTreeStorage,
|
||||
JsonlSessionTreeStorage,
|
||||
} from "../../src/harness/session-tree.js";
|
||||
import type { SessionTreeStorage } from "../../src/harness/types.js";
|
||||
|
||||
function createUserMessage(text: string): AgentMessage {
|
||||
return {
|
||||
role: "user",
|
||||
content: [{ type: "text", text }],
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
}
|
||||
|
||||
function createAssistantMessage(text: string): AgentMessage {
|
||||
return {
|
||||
role: "assistant",
|
||||
content: [{ type: "text", text }],
|
||||
api: "anthropic-messages",
|
||||
provider: "anthropic",
|
||||
model: "claude-sonnet-4-5",
|
||||
usage: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
stopReason: "stop",
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
}
|
||||
|
||||
const tempDirs: string[] = [];
|
||||
|
||||
afterEach(() => {
|
||||
while (tempDirs.length > 0) {
|
||||
const dir = tempDirs.pop()!;
|
||||
if (existsSync(dir)) {
|
||||
rmSync(dir, { recursive: true, force: true });
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
async function runSessionTreeSuite(name: string, createStorage: () => SessionTreeStorage, inspect?: () => void) {
|
||||
describe(name, () => {
|
||||
it("appends messages and builds context in order", async () => {
|
||||
const tree = new DefaultSessionTree(createStorage());
|
||||
await tree.appendMessage(createUserMessage("one"));
|
||||
await tree.appendMessage(createAssistantMessage("two"));
|
||||
const context = await tree.buildContext();
|
||||
expect(context.messages.map((message) => message.role)).toEqual(["user", "assistant"]);
|
||||
});
|
||||
|
||||
it("tracks model and thinking level changes", async () => {
|
||||
const tree = new DefaultSessionTree(createStorage());
|
||||
await tree.appendMessage(createUserMessage("one"));
|
||||
await tree.appendModelChange("openai", "gpt-4.1");
|
||||
await tree.appendThinkingLevelChange("high");
|
||||
const context = await tree.buildContext();
|
||||
expect(context.thinkingLevel).toBe("high");
|
||||
expect(context.model).toEqual({ provider: "openai", modelId: "gpt-4.1" });
|
||||
});
|
||||
|
||||
it("supports branching by moving the leaf and appending a new branch", async () => {
|
||||
const tree = new DefaultSessionTree(createStorage());
|
||||
const user1 = await tree.appendMessage(createUserMessage("one"));
|
||||
const assistant1 = await tree.appendMessage(createAssistantMessage("two"));
|
||||
await tree.appendMessage(createUserMessage("three"));
|
||||
await tree.moveTo(user1);
|
||||
await tree.appendMessage(createAssistantMessage("branched"));
|
||||
const branch = await tree.getBranch();
|
||||
expect(branch.map((entry) => entry.id)).toContain(user1);
|
||||
expect(branch.map((entry) => entry.id)).not.toContain(assistant1);
|
||||
const context = await tree.buildContext();
|
||||
expect(context.messages.map((message) => message.role)).toEqual(["user", "assistant"]);
|
||||
});
|
||||
|
||||
it("supports moving the leaf to root", async () => {
|
||||
const tree = new DefaultSessionTree(createStorage());
|
||||
await tree.appendMessage(createUserMessage("one"));
|
||||
await tree.moveTo(null);
|
||||
expect(await tree.getLeafId()).toBeNull();
|
||||
expect((await tree.buildContext()).messages).toEqual([]);
|
||||
});
|
||||
|
||||
it("reconstructs compaction summaries in context", async () => {
|
||||
const tree = new DefaultSessionTree(createStorage());
|
||||
await tree.appendMessage(createUserMessage("one"));
|
||||
await tree.appendMessage(createAssistantMessage("two"));
|
||||
const user2 = await tree.appendMessage(createUserMessage("three"));
|
||||
await tree.appendMessage(createAssistantMessage("four"));
|
||||
await tree.appendCompaction("summary", user2, 1234);
|
||||
await tree.appendMessage(createUserMessage("five"));
|
||||
const context = await tree.buildContext();
|
||||
expect(context.messages[0]?.role).toBe("compactionSummary");
|
||||
expect(context.messages).toHaveLength(4);
|
||||
});
|
||||
|
||||
it("supports branch summary entries in context", async () => {
|
||||
const tree = new DefaultSessionTree(createStorage());
|
||||
const user1 = await tree.appendMessage(createUserMessage("one"));
|
||||
await tree.appendBranchSummary(user1, "summary text");
|
||||
const context = await tree.buildContext();
|
||||
expect(context.messages[1]?.role).toBe("branchSummary");
|
||||
});
|
||||
|
||||
it("supports custom message entries in context", async () => {
|
||||
const tree = new DefaultSessionTree(createStorage());
|
||||
await tree.appendMessage(createUserMessage("one"));
|
||||
await tree.appendCustomMessageEntry("custom", "hello", true, { ok: true });
|
||||
const context = await tree.buildContext();
|
||||
expect(context.messages[1]?.role).toBe("custom");
|
||||
});
|
||||
|
||||
it("supports labels and session info entries without affecting context", async () => {
|
||||
const tree = new DefaultSessionTree(createStorage());
|
||||
const user1 = await tree.appendMessage(createUserMessage("one"));
|
||||
await tree.appendLabelChange(user1, "checkpoint");
|
||||
await tree.appendSessionInfo("name");
|
||||
const entries = await tree.getEntries();
|
||||
expect(entries.some((entry) => entry.type === "label")).toBe(true);
|
||||
expect(entries.some((entry) => entry.type === "session_info")).toBe(true);
|
||||
expect((await tree.buildContext()).messages).toHaveLength(1);
|
||||
});
|
||||
|
||||
it("persists leaf changes and appended entries via storage", async () => {
|
||||
const storage = createStorage();
|
||||
const tree = new DefaultSessionTree(storage);
|
||||
const user1 = await tree.appendMessage(createUserMessage("one"));
|
||||
await tree.appendMessage(createAssistantMessage("two"));
|
||||
await tree.moveTo(user1);
|
||||
await tree.appendMessage(createAssistantMessage("branched"));
|
||||
const tree2 = new DefaultSessionTree(storage);
|
||||
const context = await tree2.buildContext();
|
||||
expect(context.messages.map((message) => message.role)).toEqual(["user", "assistant"]);
|
||||
inspect?.();
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
runSessionTreeSuite("SessionTree with in-memory storage", () => new InMemorySessionTreeStorage());
|
||||
|
||||
runSessionTreeSuite(
|
||||
"SessionTree with JSONL storage",
|
||||
() => {
|
||||
const dir = join(tmpdir(), `pi-agent-session-tree-${Date.now()}-${Math.random().toString(36).slice(2)}`);
|
||||
mkdirSync(dir, { recursive: true });
|
||||
tempDirs.push(dir);
|
||||
return new JsonlSessionTreeStorage(join(dir, "session.jsonl"), { cwd: dir });
|
||||
},
|
||||
() => {
|
||||
const dir = tempDirs[tempDirs.length - 1]!;
|
||||
const filePath = join(dir, "session.jsonl");
|
||||
const lines = readFileSync(filePath, "utf8").trim().split("\n");
|
||||
expect(lines.length).toBeGreaterThan(1);
|
||||
const header = JSON.parse(lines[0]!);
|
||||
expect(header.type).toBe("session");
|
||||
expect(header.version).toBe(3);
|
||||
for (const line of lines.slice(1)) {
|
||||
const entry = JSON.parse(line);
|
||||
expect(entry.type).not.toBe("entry");
|
||||
expect(entry.type).not.toBe("leaf");
|
||||
expect(typeof entry.id).toBe("string");
|
||||
}
|
||||
},
|
||||
);
|
||||
Reference in New Issue
Block a user