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";
|
export * from "./agent.js";
|
||||||
// Loop functions
|
// Loop functions
|
||||||
export * from "./agent-loop.js";
|
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
|
// Proxy utilities
|
||||||
export * from "./proxy.js";
|
export * from "./proxy.js";
|
||||||
// Types
|
// 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