diff --git a/packages/agent/src/harness/agent-harness.ts b/packages/agent/src/harness/agent-harness.ts new file mode 100644 index 00000000..7eedffc3 --- /dev/null +++ b/packages/agent/src/harness/agent-harness.ts @@ -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(); + private listeners = new Set<(event: AgentHarnessEvent, signal?: AbortSignal) => Promise | void>(); + private hooks = new Map< + keyof AgentHarnessEventResultMap, + Set<(event: any, ctx: AgentHarnessContext) => Promise | 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, + }, + 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, + 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) }; + 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 { + for (const listener of this.listeners) { + await listener(event, signal); + } + } + + private async emitAny(event: AgentHarnessEvent, signal?: AbortSignal): Promise { + for (const listener of this.listeners) { + await listener(event, signal); + } + } + + private async emitHook( + type: TType, + event: Extract, + signal?: AbortSignal, + ): Promise { + 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 { + await this.emitOwn({ + type: "queue_update", + steer: [...this.operation.steerQueue], + followUp: [...this.operation.followUpQueue], + nextTurn: [...this.conversation.nextTurnQueue], + }); + } + + private async syncFromTree(): Promise { + 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 { + 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 { + 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 { + 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 { + 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 { + 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; + 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 { + 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): Promise { + 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 { + 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 { + 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 { + 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 { + 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 { + await this.agent.waitForIdle(); + } + + subscribe( + listener: ( + event: AgentEvent | import("./types.js").AgentHarnessOwnEvent, + signal?: AbortSignal, + ) => Promise | void, + ): () => void { + this.listeners.add(listener); + return () => this.listeners.delete(listener); + } + + on( + type: TType, + handler: ( + event: Extract, + ctx: AgentHarnessContext, + ) => Promise | 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; + } +} diff --git a/packages/agent/src/harness/compaction.ts b/packages/agent/src/harness/compaction.ts new file mode 100644 index 00000000..ea484a48 --- /dev/null +++ b/packages/agent/src/harness/compaction.ts @@ -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 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 | 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(`\n${readFiles.join("\n")}\n`); + if (modifiedFiles.length > 0) sections.push(`\n${modifiedFiles.join("\n")}\n`); + 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): 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) + .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, + reserveTokens: number, + apiKey: string, + headers?: Record, + signal?: AbortSignal, + customInstructions?: string, + previousSummary?: string, + thinkingLevel?: ThinkingLevel, +): Promise { + 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 = `\n${conversationText}\n\n\n`; + if (previousSummary) promptText += `\n${previousSummary}\n\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, + reserveTokens: number, + apiKey: string, + headers?: Record, + signal?: AbortSignal, + thinkingLevel?: ThinkingLevel, +): Promise { + const maxTokens = Math.floor(0.5 * reserveTokens); + const llmMessages = convertToLlm(messages); + const conversationText = serializeConversation(llmMessages); + const promptText = `\n${conversationText}\n\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, + apiKey: string, + headers?: Record, + customInstructions?: string, + signal?: AbortSignal, + thinkingLevel?: ThinkingLevel, +): Promise { + 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; + getEntry(id: string): Promise; + }, + 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 { + 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 = `\n${conversationText}\n\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 }; +} diff --git a/packages/agent/src/harness/execution-env.ts b/packages/agent/src/harness/execution-env.ts new file mode 100644 index 00000000..0419ea1d --- /dev/null +++ b/packages/agent/src/harness/execution-env.ts @@ -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): 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; + 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 { + return await readFile(resolvePath(this.cwd, path), "utf8"); + } + + async readBinaryFile(path: string): Promise { + return await readFile(resolvePath(this.cwd, path)); + } + + async writeFile(path: string, content: string | Uint8Array): Promise { + 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 { + return await readdir(resolvePath(this.cwd, path)); + } + + async pathExists(path: string): Promise { + try { + await stat(resolvePath(this.cwd, path)); + return true; + } catch { + return false; + } + } + + async createDir(path: string, options?: { recursive?: boolean }): Promise { + await mkdir(resolvePath(this.cwd, path), { recursive: options?.recursive }); + } + + async remove(path: string, options?: { recursive?: boolean; force?: boolean }): Promise { + await rm(resolvePath(this.cwd, path), { recursive: options?.recursive, force: options?.force }); + } + + async createTempDir(prefix: string = "tmp-"): Promise { + return await mkdtemp(join(tmpdir(), prefix)); + } + + async createTempFile(options?: { prefix?: string; suffix?: string }): Promise { + 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 { + // nothing to clean up for the local node implementation + } +} diff --git a/packages/agent/src/harness/messages.ts b/packages/agent/src/harness/messages.ts new file mode 100644 index 00000000..131ff530 --- /dev/null +++ b/packages/agent/src/harness/messages.ts @@ -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: + + +`; + +export const COMPACTION_SUMMARY_SUFFIX = ` +`; + +export const BRANCH_SUMMARY_PREFIX = `The following is a summary of a branch that this conversation came back from: + + +`; + +export const BRANCH_SUMMARY_SUFFIX = ``; + +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 { + 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); +} diff --git a/packages/agent/src/harness/prompt-templates.ts b/packages/agent/src/harness/prompt-templates.ts new file mode 100644 index 00000000..a05336d3 --- /dev/null +++ b/packages/agent/src/harness/prompt-templates.ts @@ -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)); +} diff --git a/packages/agent/src/harness/session-tree.ts b/packages/agent/src/harness/session-tree.ts new file mode 100644 index 00000000..4b3855fb --- /dev/null +++ b/packages/agent/src/harness/session-tree.ts @@ -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(); + private currentLeafId: string | null = null; + + constructor(filePath: string, options: { cwd: string }) { + this.filePath = resolve(filePath); + this.cwd = options.cwd; + } + + private async ensureParentDir(): Promise { + await mkdir(dirname(this.filePath), { recursive: true }); + } + + private async ensureLoaded(): Promise { + 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 { + 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 { + await this.ensureLoaded(); + return this.currentLeafId; + } + + async setLeafId(leafId: string | null): Promise { + await this.ensureLoaded(); + this.currentLeafId = leafId; + } + + async appendEntry(entry: SessionTreeEntry): Promise { + 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 { + await this.ensureLoaded(); + return this.byId.get(id); + } + + async getPathToRoot(leafId: string | null): Promise { + 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 { + 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 { + return this.leafId; + } + + async setLeafId(leafId: string | null): Promise { + this.leafId = leafId; + } + + async appendEntry(entry: SessionTreeEntry): Promise { + this.entries.push(entry); + this.leafId = entry.id; + } + + async getEntry(id: string): Promise { + return this.entries.find((entry) => entry.id === id); + } + + async getPathToRoot(leafId: string | null): Promise { + if (leafId === null) return []; + const byId = new Map(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 { + return [...this.entries]; + } +} + +export class DefaultSessionTree implements SessionTree { + private storage: SessionTreeStorage; + + constructor(storage?: SessionTreeStorage) { + this.storage = storage ?? new InMemorySessionTreeStorage(); + } + + getLeafId(): Promise { + return this.storage.getLeafId(); + } + + getEntry(id: string): Promise { + return this.storage.getEntry(id); + } + + getEntries(): Promise { + return this.storage.getEntries(); + } + + async getBranch(fromId?: string): Promise { + const leafId = fromId ?? (await this.storage.getLeafId()); + return this.storage.getPathToRoot(leafId); + } + + async buildContext(): Promise { + return buildSessionContext(await this.getBranch()); + } + + private async makeEntryId(): Promise { + const entries = await this.storage.getEntries(); + return generateId(new Set(entries.map((entry) => entry.id))); + } + + private async appendTypedEntry(entry: TEntry): Promise { + await this.storage.appendEntry(entry); + return entry.id; + } + + async appendMessage(message: AgentMessage): Promise { + 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 { + 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 { + 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( + summary: string, + firstKeptEntryId: string, + tokensBefore: number, + details?: T, + fromHook?: boolean, + ): Promise { + 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); + } + + async appendBranchSummary( + fromId: string, + summary: string, + details?: T, + fromHook?: boolean, + ): Promise { + 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); + } + + async appendCustomEntry(customType: string, data?: unknown): Promise { + return this.appendTypedEntry({ + type: "custom", + id: await this.makeEntryId(), + parentId: await this.storage.getLeafId(), + timestamp: new Date().toISOString(), + customType, + data, + } satisfies CustomEntry); + } + + async appendCustomMessageEntry( + customType: string, + content: string | (TextContent | ImageContent)[], + display: boolean, + details?: T, + ): Promise { + 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); + } + + async appendLabelChange(targetId: string, label: string | undefined): Promise { + 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 { + 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 { + if (entryId !== null && !(await this.storage.getEntry(entryId))) { + throw new Error(`Entry ${entryId} not found`); + } + await this.storage.setLeafId(entryId); + } +} diff --git a/packages/agent/src/harness/types.ts b/packages/agent/src/harness/types.ts new file mode 100644 index 00000000..41f701b3 --- /dev/null +++ b/packages/agent/src/harness/types.ts @@ -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; + 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; + readBinaryFile(path: string): Promise; + writeFile(path: string, content: string | Uint8Array): Promise; + stat(path: string): Promise<{ + isFile: boolean; + isDirectory: boolean; + isSymbolicLink: boolean; + size: number; + mtime: Date; + }>; + listDir(path: string): Promise; + pathExists(path: string): Promise; + createDir(path: string, options?: { recursive?: boolean }): Promise; + remove(path: string, options?: { recursive?: boolean; force?: boolean }): Promise; + createTempDir(prefix?: string): Promise; + createTempFile(options?: { prefix?: string; suffix?: string }): Promise; + + resolvePath(path: string): string; + cleanup(): Promise; +} + +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 extends SessionTreeEntryBase { + type: "compaction"; + summary: string; + firstKeptEntryId: string; + tokensBefore: number; + details?: T; + fromHook?: boolean; +} + +export interface BranchSummaryEntry extends SessionTreeEntryBase { + type: "branch_summary"; + fromId: string; + summary: string; + details?: T; + fromHook?: boolean; +} + +export interface CustomEntry extends SessionTreeEntryBase { + type: "custom"; + customType: string; + data?: T; +} + +export interface CustomMessageEntry 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; + setLeafId(leafId: string | null): Promise; + appendEntry(entry: SessionTreeEntry): Promise; + getEntry(id: string): Promise; + getPathToRoot(leafId: string | null): Promise; + getEntries(): Promise; +} + +export interface SessionTree { + getLeafId(): Promise; + getEntry(id: string): Promise; + getEntries(): Promise; + getBranch(fromId?: string): Promise; + buildContext(): Promise; + + appendMessage(message: AgentMessage): Promise; + appendThinkingLevelChange(thinkingLevel: string): Promise; + appendModelChange(provider: string, modelId: string): Promise; + appendCompaction( + summary: string, + firstKeptEntryId: string, + tokensBefore: number, + details?: T, + fromHook?: boolean, + ): Promise; + appendBranchSummary(fromId: string, summary: string, details?: T, fromHook?: boolean): Promise; + appendCustomEntry(customType: string, data?: unknown): Promise; + appendCustomMessageEntry( + customType: string, + content: string | (TextContent | ImageContent)[], + display: boolean, + details?: T, + ): Promise; + appendLabelChange(targetId: string, label: string | undefined): Promise; + appendSessionInfo(name: string): Promise; + + moveTo(entryId: string | null): Promise; +} + +export interface AgentHarnessPendingMutations { + appendMessages: AgentMessage[]; + model?: Model; + thinkingLevel?: ThinkingLevel; + activeToolNames?: string[]; + systemPromptInputs?: SystemPromptInputs; +} + +export interface AgentHarnessConversationState { + sessionTree: SessionTree; + model: Model | 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 | undefined; + thinkingLevel: ThinkingLevel; + activeToolNames: string[]; + systemPrompt: string; +} + +export interface AgentHarnessContext { + env: ExecutionEnv; + conversation: Readonly; + operation: Readonly; + 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; +} + +export interface ToolCallEvent { + type: "tool_call"; + toolCallId: string; + toolName: string; + input: Record; +} + +export interface ToolResultEvent { + type: "tool_result"; + toolCallId: string; + toolName: string; + input: Record; + content: Array; + 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; + previousModel: Model | 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; + 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; + written: Set; + edited: Set; +} + +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; + apiKey: string; + headers?: Record; + 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) => Promise<{ apiKey: string; headers?: Record } | undefined>; + initialModel?: Model; + initialThinkingLevel?: ThinkingLevel; + initialActiveToolNames?: string[]; + initialSystemPromptInputs?: SystemPromptInputs; +} + +export interface AgentHarness { + readonly agent: Agent; + readonly env: ExecutionEnv; + readonly conversation: Readonly; + readonly operation: Readonly; + + prompt(text: string, options?: AgentHarnessPromptOptions): Promise; + skill(name: string, args?: string): Promise; + + steer(message: AgentMessage): void; + followUp(message: AgentMessage): void; + nextTurn(message: AgentMessage): void; + + appendMessage(message: AgentMessage): Promise; + + shell( + command: string, + options?: { + cwd?: string; + env?: Record; + timeout?: number; + signal?: AbortSignal; + onStdout?: (chunk: string) => void; + onStderr?: (chunk: string) => void; + }, + ): Promise<{ stdout: string; stderr: string; exitCode: number }>; + + compact(customInstructions?: string): Promise; + navigateTree( + targetId: string, + options?: { summarize?: boolean; customInstructions?: string; replaceInstructions?: boolean; label?: string }, + ): Promise; + + setModel(model: Model): Promise; + setThinkingLevel(level: ThinkingLevel): Promise; + setActiveTools(toolNames: string[]): Promise; + setSystemPromptInputs(inputs: SystemPromptInputs): Promise; + + abort(): Promise; + waitForIdle(): Promise; + + subscribe(listener: (event: AgentHarnessEvent, signal?: AbortSignal) => Promise | void): () => void; + + on( + type: TType, + handler: ( + event: Extract, + ctx: AgentHarnessContext, + ) => Promise | AgentHarnessEventResultMap[TType], + ): () => void; +} diff --git a/packages/agent/src/harness/utils/shell-output.ts b/packages/agent/src/harness/utils/shell-output.ts new file mode 100644 index 00000000..d31d7a06 --- /dev/null +++ b/packages/agent/src/harness/utils/shell-output.ts @@ -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 { + 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 { + 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; + } +} diff --git a/packages/agent/src/harness/utils/truncate.ts b/packages/agent/src/harness/utils/truncate.ts new file mode 100644 index 00000000..18ac5d74 --- /dev/null +++ b/packages/agent/src/harness/utils/truncate.ts @@ -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 }; +} diff --git a/packages/agent/src/index.ts b/packages/agent/src/index.ts index d8ed5b8e..ead51788 100644 --- a/packages/agent/src/index.ts +++ b/packages/agent/src/index.ts @@ -2,6 +2,16 @@ export * from "./agent.js"; // Loop functions export * from "./agent-loop.js"; +export * from "./harness/agent-harness.js"; +export * from "./harness/compaction.js"; +export * from "./harness/execution-env.js"; +export * from "./harness/messages.js"; +export * from "./harness/prompt-templates.js"; +export * from "./harness/session-tree.js"; +// Harness +export * from "./harness/types.js"; +export * from "./harness/utils/shell-output.js"; +export * from "./harness/utils/truncate.js"; // Proxy utilities export * from "./proxy.js"; // Types diff --git a/packages/agent/test/harness/compaction.test.ts b/packages/agent/test/harness/compaction.test.ts new file mode 100644 index 00000000..4c042df7 --- /dev/null +++ b/packages/agent/test/harness/compaction.test.ts @@ -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 } { + 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 | undefined> = []; + const { faux: fauxReasoning, model: reasoningModel } = createFauxModel(true); + fauxReasoning.setResponses([ + (_context, options) => { + seenOptions.push(options as Record | 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 | 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 | 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; +} diff --git a/packages/agent/test/harness/session-tree.test.ts b/packages/agent/test/harness/session-tree.test.ts new file mode 100644 index 00000000..64a6dba4 --- /dev/null +++ b/packages/agent/test/harness/session-tree.test.ts @@ -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"); + } + }, +);