diff --git a/docs/architecture.md b/docs/architecture.md index 328cd292..1cc4eb16 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -41,7 +41,7 @@ someone who wants to read the code, contribute, or fork. - provider payload transforms and protocol quirks (for example, Yutori tool serialization policy) - canonical CUA tool-definition exports - `@onkernel/cua-agent` owns browser execution orchestration: - - `CuaAgent` / `CuaHarness` class wiring around `pi-agent-core` + - `CuaAgent` / `CuaAgentHarness` class wiring around vendored pi agent core - executing canonical CUA tool calls against Kernel browsers - typed executor coverage and translator integration @@ -66,7 +66,7 @@ Dev/test: └── @onkernel/ptywright (PTY-backed TUI regression harness) External: -├── @earendil-works/pi-agent-core # Agent loop, tool execution, streaming, steering +├── vendored pi agent core # Agent loop, tool execution, streaming, steering │ └── @earendil-works/pi-ai # Provider transport (OpenAI Responses, Anthropic Messages, Google GenAI) ├── @earendil-works/pi-coding-agent # bash / read / write / edit / grep / find / ls AgentTools + SessionManager + skills ├── @earendil-works/pi-tui # Terminal, Editor, Image, differential renderer diff --git a/package-lock.json b/package-lock.json index 3091a2cd..d236a9ef 100644 --- a/package-lock.json +++ b/package-lock.json @@ -680,19 +680,6 @@ "url": "https://github.com/sponsors/Borewit" } }, - "node_modules/@earendil-works/pi-agent-core": { - "version": "0.74.0", - "resolved": "https://registry.npmjs.org/@earendil-works/pi-agent-core/-/pi-agent-core-0.74.0.tgz", - "integrity": "sha512-6GMR7/wwjEJ1EsXLWEz03QOWin4AMrJ/AZoMpgm5DJ6GHsF6q6GOhQbj5Zip4dow3vo/TmBAVqM+vmGfrjGAFQ==", - "license": "MIT", - "dependencies": { - "@earendil-works/pi-ai": "^0.74.0", - "typebox": "^1.1.24" - }, - "engines": { - "node": ">=20.0.0" - } - }, "node_modules/@earendil-works/pi-ai": { "version": "0.74.0", "resolved": "https://registry.npmjs.org/@earendil-works/pi-ai/-/pi-ai-0.74.0.tgz", @@ -4632,10 +4619,10 @@ "name": "@onkernel/cua-agent", "version": "0.1.0", "dependencies": { - "@earendil-works/pi-agent-core": "^0.74.0", "@earendil-works/pi-ai": "^0.74.0", "@onkernel/cua-ai": "0.1.0", - "@onkernel/sdk": "0.49.0" + "@onkernel/sdk": "0.49.0", + "typebox": "^1.1.38" }, "devDependencies": { "vitest": "^3.2.4" diff --git a/packages/agent/README.md b/packages/agent/README.md index 0da87ee1..ee6b780b 100644 --- a/packages/agent/README.md +++ b/packages/agent/README.md @@ -1,9 +1,9 @@ # `@onkernel/cua-agent` -Kernel browser computer-use classes built on -[`@earendil-works/pi-agent-core`](https://github.com/earendil-works/pi/tree/main/packages/agent). +Kernel browser computer-use classes built on vendored pi `Agent` and +`AgentHarness` source. -This package keeps pi-agent-core semantics intact and adds browser execution +This package keeps pi agent semantics intact and adds browser execution plumbing for canonical CUA tools. ## Installation @@ -33,45 +33,59 @@ const agent = new CuaAgent({ await agent.prompt("Open news.ycombinator.com and summarize the top story."); ``` -## Quick Start (`CuaHarness`) +## Quick Start (`CuaAgentHarness`) ```ts -import { CuaHarness } from "@onkernel/cua-agent"; +import { CuaAgentHarness, InMemorySessionRepo, NodeExecutionEnv } from "@onkernel/cua-agent"; -const harness = new CuaHarness({ +const sessionRepo = new InMemorySessionRepo(); +const session = await sessionRepo.create({ id: "example" }); + +const harness = new CuaAgentHarness({ browser, client, + env: new NodeExecutionEnv({ cwd: process.cwd() }), model: "openai:gpt-5.5", + session, }); -await harness.prompt("Open example.com and tell me the current URL."); -const transcript = harness.getTranscript(); -console.log("messages in transcript:", transcript.length); +const response = await harness.prompt("Open example.com and tell me the current URL."); +const branch = await session.getBranch(); +const lastAssistant = [...branch] + .reverse() + .flatMap((entry) => + entry.type === "message" && entry.message.role === "assistant" ? [entry.message] : [], + )[0]; +const assistant = lastAssistant ?? response; +const assistantText = assistant.content + .flatMap((block) => (block.type === "text" ? [block.text] : [])) + .join("") + .trim(); +console.log("assistant stopReason:", assistant.stopReason); +console.log("assistant text:", assistantText || "(no text)"); ``` -Use `CuaAgent` when you want direct pi `Agent` control: raw transcript state, +Use `CuaAgent` when you want direct pi `Agent` control: raw message state, lifecycle events, custom streaming, and explicit prompt/continue/queue control. Reach for the harness shape when you want an app layer around the loop: -session/transcript helpers, resource and prompt entry points, provider/auth -hooks, active tool selection, compaction/tree workflows, and higher-level queue -events. `CuaHarness` is the thin CUA version of that shape today: it installs -CUA defaults, delegates runtime methods to the wrapped `Agent`, and adds -`getTranscript()`. +session-backed turns, resource and prompt entry points, provider/auth hooks, +active tool selection, compaction/tree workflows, and higher-level queue events. +`CuaAgentHarness` extends pi `AgentHarness`, installs CUA defaults, and refreshes +provider-specific runtime state when `setModel()` changes models. ## Core Concepts ### Class-First API - `CuaAgent extends Agent` -- `CuaHarness` wraps a pi `Agent` with a harness-style constructor and - delegated runtime methods. +- `CuaAgentHarness extends AgentHarness` Both classes mirror pi constructor shapes and behavior, with minimal additions: - `browser` (Kernel browser response) - `client` (Kernel SDK client) - CUA model refs (`"provider:model"`) accepted where pi expects a concrete model -If `getApiKey` is omitted, both classes default to CUA env var conventions: +If auth callbacks are omitted, both classes default to CUA env var conventions: - OpenAI: `OPENAI_API_KEY` - Anthropic: `ANTHROPIC_OAUTH_TOKEN` or `ANTHROPIC_API_KEY` - Gemini: `GOOGLE_API_KEY` or `GEMINI_API_KEY` @@ -84,6 +98,17 @@ If tools are omitted, the classes install canonical CUA computer tool executors using runtime specs from `@onkernel/cua-ai`. If tools are provided, they are used exactly. +### Model Switching + +`CuaAgent` follows pi `Agent` semantics: assign `agent.state.model` to a +concrete model or CUA model ref. CUA-owned tools and the default system prompt +refresh with the new provider runtime. + +`CuaAgentHarness` follows pi `AgentHarness` semantics: call +`await harness.setModel(model)`. The harness updates its model through pi's +snapshot machinery and refreshes CUA-owned tools and default prompt state for +the next provider request. + ### Tool Composition Use `createCuaComputerTools()` to compose your own tool list from canonical @@ -105,4 +130,4 @@ const tools = [ ``` For full event semantics, steering, follow-up queues, and tool execution -details, see the pi-agent-core README. +details, see the pi agent core source vendored in this package. diff --git a/packages/agent/examples/agent-openai-smoke.ts b/packages/agent/examples/agent-openai-smoke.ts index 38a09cbb..25a2893e 100644 --- a/packages/agent/examples/agent-openai-smoke.ts +++ b/packages/agent/examples/agent-openai-smoke.ts @@ -1,6 +1,7 @@ import Kernel from "@onkernel/sdk"; import { requireCuaEnvApiKeyForModel, type CuaModelRef } from "@onkernel/cua-ai"; import { CuaAgent } from "../src/index"; +import { logAgentEvent, logAssistant } from "./shared/logging"; import { SCENARIOS } from "./shared/scenarios"; const modelRef = (process.env.MODEL_REF as CuaModelRef | undefined) ?? "openai:gpt-5.5"; @@ -19,20 +20,13 @@ async function main(): Promise { initialState: { model: modelRef }, }); - agent.subscribe((event) => { - if (event.type === "tool_execution_start") { - console.log(`[tool:start] ${event.toolName}`); - } - if (event.type === "tool_execution_end") { - console.log(`[tool:end] ${event.toolName} error=${event.isError}`); - } - }); + agent.subscribe(logAgentEvent); const scenario = SCENARIOS[0]!; - console.log(`running scenario: ${scenario.name}`); + console.log(`running scenario: ${scenario.name} model=${modelRef}`); await agent.prompt(scenario.prompt); const assistant = [...agent.state.messages].reverse().find((message) => message.role === "assistant"); - console.log("assistant stopReason:", assistant?.role === "assistant" ? assistant.stopReason : "unknown"); + logAssistant(assistant?.role === "assistant" ? assistant : undefined); } finally { await client.browsers.deleteByID(browser.session_id); } diff --git a/packages/agent/examples/agent-provider-matrix.ts b/packages/agent/examples/agent-provider-matrix.ts index e661d4af..7f6f25b9 100644 --- a/packages/agent/examples/agent-provider-matrix.ts +++ b/packages/agent/examples/agent-provider-matrix.ts @@ -1,6 +1,7 @@ import Kernel from "@onkernel/sdk"; import { requireCuaEnvApiKeyForModel, type CuaModelRef } from "@onkernel/cua-ai"; import { CuaAgent } from "../src/index"; +import { logAgentEvent, logAssistant } from "./shared/logging"; import { SCENARIOS } from "./shared/scenarios"; const modelRef = (process.env.MODEL_REF as CuaModelRef | undefined) ?? "openai:gpt-5.5"; @@ -20,8 +21,11 @@ async function main(): Promise { client, initialState: { model: modelRef }, }); + agent.subscribe(logAgentEvent); console.log(`model=${modelRef} scenario=${scenario.name}`); await agent.prompt(scenario.prompt); + const assistant = [...agent.state.messages].reverse().find((message) => message.role === "assistant"); + logAssistant(assistant?.role === "assistant" ? assistant : undefined); } finally { await client.browsers.deleteByID(browser.session_id); } diff --git a/packages/agent/examples/harness-openai-smoke.ts b/packages/agent/examples/harness-openai-smoke.ts index 3c38edcc..7891efbe 100644 --- a/packages/agent/examples/harness-openai-smoke.ts +++ b/packages/agent/examples/harness-openai-smoke.ts @@ -1,6 +1,7 @@ import Kernel from "@onkernel/sdk"; import { requireCuaEnvApiKeyForModel, type CuaModelRef } from "@onkernel/cua-ai"; -import { CuaHarness } from "../src/index"; +import { CuaAgentHarness, InMemorySessionRepo, NodeExecutionEnv } from "../src/index"; +import { logAgentEvent, logAssistant } from "./shared/logging"; import { SCENARIOS } from "./shared/scenarios"; const modelRef = (process.env.MODEL_REF as CuaModelRef | undefined) ?? "openai:gpt-5.5"; @@ -13,28 +14,28 @@ async function main(): Promise { const browser = await client.browsers.create({ stealth: true }); try { - const harness = new CuaHarness({ + const sessionRepo = new InMemorySessionRepo(); + const session = await sessionRepo.create({ id: "harness-openai-smoke" }); + const harness = new CuaAgentHarness({ browser, client, + env: new NodeExecutionEnv({ cwd: process.cwd() }), model: modelRef, + session, }); - harness.subscribe((event) => { - if (event.type === "tool_execution_start") { - console.log(`[tool:start] ${event.toolName}`); - } - if (event.type === "tool_execution_end") { - console.log(`[tool:end] ${event.toolName} error=${event.isError}`); - } - }); + harness.subscribe(logAgentEvent); const scenario = SCENARIOS[0]!; - console.log(`running scenario: ${scenario.name}`); - await harness.prompt(scenario.prompt); - const transcript = harness.getTranscript(); - const lastAssistant = [...transcript].reverse().find((message) => message.role === "assistant"); - console.log("transcript messages:", transcript.length); - console.log("assistant stopReason:", lastAssistant?.role === "assistant" ? lastAssistant.stopReason : "unknown"); + console.log(`running scenario: ${scenario.name} model=${modelRef}`); + const response = await harness.prompt(scenario.prompt); + const branch = await session.getBranch(); + const lastAssistant = [...branch] + .reverse() + .flatMap((entry) => + entry.type === "message" && entry.message.role === "assistant" ? [entry.message] : [], + )[0]; + logAssistant(lastAssistant ?? response); } finally { await client.browsers.deleteByID(browser.session_id); } diff --git a/packages/agent/examples/harness-provider-matrix.ts b/packages/agent/examples/harness-provider-matrix.ts index 6970632b..abff2d1f 100644 --- a/packages/agent/examples/harness-provider-matrix.ts +++ b/packages/agent/examples/harness-provider-matrix.ts @@ -1,6 +1,7 @@ import Kernel from "@onkernel/sdk"; import { requireCuaEnvApiKeyForModel, type CuaModelRef } from "@onkernel/cua-ai"; -import { CuaHarness } from "../src/index"; +import { CuaAgentHarness, InMemorySessionRepo, NodeExecutionEnv } from "../src/index"; +import { logAgentEvent, logAssistant } from "./shared/logging"; import { SCENARIOS } from "./shared/scenarios"; const modelRef = (process.env.MODEL_REF as CuaModelRef | undefined) ?? "openai:gpt-5.5"; @@ -15,14 +16,25 @@ async function main(): Promise { const scenario = SCENARIOS.find((entry) => entry.name === scenarioName) ?? SCENARIOS[0]!; try { - const harness = new CuaHarness({ + const sessionRepo = new InMemorySessionRepo(); + const session = await sessionRepo.create({ id: `harness-provider-matrix-${scenario.name}` }); + const harness = new CuaAgentHarness({ browser, client, + env: new NodeExecutionEnv({ cwd: process.cwd() }), model: modelRef, + session, }); + harness.subscribe(logAgentEvent); console.log(`model=${modelRef} scenario=${scenario.name}`); - await harness.prompt(scenario.prompt); - console.log("transcript messages:", harness.getTranscript().length); + const response = await harness.prompt(scenario.prompt); + const branch = await session.getBranch(); + const lastAssistant = [...branch] + .reverse() + .flatMap((entry) => + entry.type === "message" && entry.message.role === "assistant" ? [entry.message] : [], + )[0]; + logAssistant(lastAssistant ?? response); } finally { await client.browsers.deleteByID(browser.session_id); } diff --git a/packages/agent/examples/shared/logging.ts b/packages/agent/examples/shared/logging.ts new file mode 100644 index 00000000..39d11702 --- /dev/null +++ b/packages/agent/examples/shared/logging.ts @@ -0,0 +1,37 @@ +import type { AgentEvent, AgentHarnessEvent } from "../../src/index"; + +type AssistantLike = { + content: Array<{ type: string; text?: string }>; + stopReason?: string; +}; + +export function logAgentEvent(event: AgentEvent | AgentHarnessEvent): void { + if (event.type === "tool_execution_start") { + console.log(`[tool:start] ${event.toolName} args=${formatJson(event.args)}`); + return; + } + if (event.type === "tool_execution_end") { + const result = event.result as { details?: unknown } | undefined; + console.log(`[tool:end] ${event.toolName} error=${event.isError}${formatDetails(result?.details)}`); + } +} + +export function logAssistant(assistant: AssistantLike | undefined): void { + const text = + assistant?.content + .flatMap((block) => (block.type === "text" && typeof block.text === "string" ? [block.text] : [])) + .join("") + .trim() ?? ""; + console.log("assistant stopReason:", assistant?.stopReason ?? "unknown"); + console.log("assistant text:", text || "(no text)"); +} + +function formatDetails(details: unknown): string { + if (!details) return ""; + return ` details=${formatJson(details)}`; +} + +function formatJson(value: unknown): string { + const text = JSON.stringify(value, null, 2) ?? "undefined"; + return text.length > 1200 ? `${text.slice(0, 1197)}...` : text; +} diff --git a/packages/agent/package.json b/packages/agent/package.json index cf11139c..f6e59f16 100644 --- a/packages/agent/package.json +++ b/packages/agent/package.json @@ -16,7 +16,9 @@ "dist", "examples", "README.md", - "CHANGELOG.md" + "CHANGELOG.md", + "src/vendor/pi-agent-core/LICENSE", + "src/vendor/pi-agent-core/README.md" ], "scripts": { "build": "tsc -b", @@ -26,10 +28,10 @@ "test": "vitest --run" }, "dependencies": { - "@earendil-works/pi-agent-core": "^0.74.0", "@earendil-works/pi-ai": "^0.74.0", "@onkernel/cua-ai": "0.1.0", - "@onkernel/sdk": "0.49.0" + "@onkernel/sdk": "0.49.0", + "typebox": "^1.1.38" }, "devDependencies": { "vitest": "^3.2.4" diff --git a/packages/agent/scripts/vendor-pi-agent-harness.ts b/packages/agent/scripts/vendor-pi-agent-harness.ts new file mode 100644 index 00000000..4fd4f166 --- /dev/null +++ b/packages/agent/scripts/vendor-pi-agent-harness.ts @@ -0,0 +1,62 @@ +/** + * Refresh the vendored pi agent core files used by `@onkernel/cua-agent`. + * + * The published pi agent package does not currently expose the `AgentHarness` + * APIs this package extends, so we vendor the minimal source set from a pinned + * official `earendil-works/pi` commit with its MIT license. + */ +import { mkdir, writeFile } from "node:fs/promises"; +import { dirname, join } from "node:path"; +import { fileURLToPath } from "node:url"; + +const PI_COMMIT = "40c05f55391663024a6a05ad33249b616a04e7a1"; +const FILES = [ + "agent.ts", + "agent-loop.ts", + "harness/agent-harness.ts", + "harness/compaction/branch-summarization.ts", + "harness/compaction/compaction.ts", + "harness/compaction/utils.ts", + "harness/env/nodejs.ts", + "harness/execution-env.ts", + "harness/messages.ts", + "harness/prompt-templates.ts", + "harness/session/repo/jsonl.ts", + "harness/session/repo/memory.ts", + "harness/session/repo/shared.ts", + "harness/session/session.ts", + "harness/session/storage/jsonl.ts", + "harness/session/storage/memory.ts", + "harness/session/uuid.ts", + "harness/skills.ts", + "harness/system-prompt.ts", + "harness/types.ts", + "harness/utils/shell-output.ts", + "harness/utils/truncate.ts", + "index.ts", + "proxy.ts", + "types.ts", +]; + +const __dirname = dirname(fileURLToPath(import.meta.url)); +const vendorRoot = join(__dirname, "../src/vendor/pi-agent-core"); +const repoRawBase = `https://raw.githubusercontent.com/earendil-works/pi/${PI_COMMIT}`; +const rawBase = `${repoRawBase}/packages/agent/src`; + +for (const file of FILES) { + const response = await fetch(`${rawBase}/${file}`); + if (!response.ok) { + throw new Error(`Failed to fetch ${file}: ${response.status} ${response.statusText}`); + } + const outputPath = join(vendorRoot, file); + await mkdir(dirname(outputPath), { recursive: true }); + await writeFile(outputPath, await response.text()); + console.log(`vendored ${file}`); +} + +const licenseResponse = await fetch(`${repoRawBase}/LICENSE`); +if (!licenseResponse.ok) { + throw new Error(`Failed to fetch LICENSE: ${licenseResponse.status} ${licenseResponse.statusText}`); +} +await writeFile(join(vendorRoot, "LICENSE"), await licenseResponse.text()); +console.log("vendored LICENSE"); diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index 8120855f..9bd76b16 100644 --- a/packages/agent/src/agent.ts +++ b/packages/agent/src/agent.ts @@ -1,152 +1,327 @@ import { Agent, + AgentHarness, + type AgentHarnessOptions, type AgentOptions, + type AgentState, type AgentTool, - type AgentEvent, - type AgentMessage, -} from "@earendil-works/pi-agent-core"; + type PromptTemplate, + type Skill, + type StreamFn, +} from "./vendor/pi-agent-core/index"; import { - type ImageContent, type Api, - type Model, type CuaModelRef, getCuaEnvApiKey, + type Model, resolveCuaRuntimeSpec, + type SimpleStreamOptions, streamSimple, } from "@onkernel/cua-ai"; import type Kernel from "@onkernel/sdk"; import { createCuaComputerTools } from "./tools"; import type { KernelBrowser } from "./translator/translator"; +/** A CUA model reference string or a concrete pi model object. */ +type CuaRuntimeInput = CuaModelRef | Model; + +type CuaRuntimeSpec = ReturnType; + +/** + * Agent state exposed by {@link CuaAgent}. + * + * It is the regular pi `AgentState`, except assigning `state.model` may use a + * CUA model ref such as `"openai:gpt-5.5"`. CUA-owned tools and the default + * system prompt are refreshed to match the new provider runtime. + */ +export interface CuaAgentState extends Omit { + /** The concrete pi model currently used by the underlying agent loop. */ + get model(): Model; + /** Assign a concrete pi model or CUA model ref and refresh CUA runtime defaults. */ + set model(model: CuaRuntimeInput); +} + +/** Initial state for {@link CuaAgent}. */ type CuaAgentInitialState = Omit, "model" | "tools"> & { - model: CuaModelRef | Model; + /** Model to use for the first turn. CUA refs are resolved before pi sees the state. */ + model: CuaRuntimeInput; + /** Optional caller-owned tools. Omit this to install the provider's default CUA tools. */ tools?: AgentTool[]; }; +/** + * Constructor options for {@link CuaAgent}. + * + * `browser` and `client` are used to build the default computer-use tools. + * Everything else follows pi `AgentOptions`, with `initialState.model` + * widened to accept CUA model refs. + */ export type CuaAgentOptions = Omit & { + /** Kernel browser session used by default CUA tools. */ browser: KernelBrowser; + /** Kernel SDK client used by default CUA tools. */ client: Kernel; + /** Initial pi state plus a CUA-aware model value. */ initialState: CuaAgentInitialState; }; -export type CuaHarnessOptions = Omit & { +/** + * Constructor options for {@link CuaAgentHarness}. + * + * The harness keeps pi `AgentHarnessOptions` intact except that `model` + * accepts CUA refs and `browser`/`client` are required to build default + * computer-use tools. Callers provide pi's `env` and `session` directly. + */ +export type CuaAgentHarnessOptions< + TSkill extends Skill = Skill, + TPromptTemplate extends PromptTemplate = PromptTemplate, +> = Omit, "model"> & { + /** Kernel browser session used by default CUA tools. */ browser: KernelBrowser; + /** Kernel SDK client used by default CUA tools. */ client: Kernel; - model: CuaModelRef | Model; - tools?: AgentTool[]; - systemPrompt?: string; + /** Model used by the harness. CUA refs are resolved before pi sees the model. */ + model: CuaRuntimeInput; + /** Optional payload hook composed after the provider-specific CUA payload hook. */ + onPayload?: SimpleStreamOptions["onPayload"]; }; -export class CuaAgent extends Agent { - constructor(options: CuaAgentOptions) { - const { browser, client, initialState, onPayload, streamFn, ...agentOptions } = options; - const runtimeSpec = resolveCuaRuntimeSpec(initialState.model); - const tools = initialState.tools ?? createCuaComputerTools({ browser, client, toolDefinitions: runtimeSpec.toolDefinitions }); - const systemPrompt = initialState.systemPrompt ?? runtimeSpec.defaultSystemPrompt; +/** + * Holds the CUA-specific pieces that have to change when a model changes. + * + * If callers omit `tools` or `systemPrompt`, CUA owns those values and refreshes + * them from `@onkernel/cua-ai` whenever the model changes. If callers pass + * their own tools or prompt, the controller preserves those caller-owned values. + */ +class CuaRuntimeController { + private runtimeSpec: CuaRuntimeSpec; - super({ - ...agentOptions, - getApiKey: agentOptions.getApiKey ?? getCuaEnvApiKey, - streamFn: streamFn ?? streamSimple, - onPayload: composeOnPayload(runtimeSpec.onPayload, onPayload), - initialState: { - ...initialState, - model: runtimeSpec.model, - tools, - systemPrompt, - }, - }); + constructor( + private readonly options: { + browser: KernelBrowser; + client: Kernel; + model: CuaRuntimeInput; + tools?: AgentTool[]; + systemPrompt?: unknown; + onPayload?: SimpleStreamOptions["onPayload"]; + }, + ) { + this.runtimeSpec = resolveCuaRuntimeSpec(options.model); } -} - -export class CuaHarness { - readonly agent: Agent; - constructor(options: CuaHarnessOptions) { - const { browser, client, model, tools, systemPrompt, onPayload, streamFn, ...agentOptions } = options; - const runtimeSpec = resolveCuaRuntimeSpec(model); - const resolvedTools = tools ?? createCuaComputerTools({ browser, client, toolDefinitions: runtimeSpec.toolDefinitions }); - this.agent = new Agent({ - ...agentOptions, - getApiKey: agentOptions.getApiKey ?? getCuaEnvApiKey, - streamFn: streamFn ?? streamSimple, - onPayload: composeOnPayload(runtimeSpec.onPayload, onPayload), - initialState: { - model: runtimeSpec.model, - tools: resolvedTools, - systemPrompt: systemPrompt ?? runtimeSpec.defaultSystemPrompt, - }, - }); + get model(): Model { + return this.runtimeSpec.model; } - get state() { - return this.agent.state; + get ownsTools(): boolean { + return this.options.tools === undefined; } - get steeringMode() { - return this.agent.steeringMode; + get ownsSystemPrompt(): boolean { + return this.options.systemPrompt === undefined; } - set steeringMode(mode: "all" | "one-at-a-time") { - this.agent.steeringMode = mode; + get systemPrompt(): string { + return this.runtimeSpec.defaultSystemPrompt; } - get followUpMode() { - return this.agent.followUpMode; + setModel(model: CuaRuntimeInput): void { + this.runtimeSpec = resolveCuaRuntimeSpec(model); } - set followUpMode(mode: "all" | "one-at-a-time") { - this.agent.followUpMode = mode; + tools(): AgentTool[] { + return ( + this.options.tools ?? + createCuaComputerTools({ + browser: this.options.browser, + client: this.options.client, + toolDefinitions: this.runtimeSpec.toolDefinitions, + }) + ); } - subscribe(listener: (event: AgentEvent, signal: AbortSignal) => Promise | void): () => void { - return this.agent.subscribe(listener); + onPayloadFor(model: CuaRuntimeInput): SimpleStreamOptions["onPayload"] { + const runtimeSpec = resolveCuaRuntimeSpec(model); + return composeOnPayload(runtimeSpec.onPayload, this.options.onPayload); } +} - async prompt(message: AgentMessage | AgentMessage[]): Promise; - async prompt(input: string, images?: ImageContent[]): Promise; - async prompt(input: string | AgentMessage | AgentMessage[], images?: ImageContent[]): Promise { - if (typeof input === "string") { - await this.agent.prompt(input, images); - return; - } - await this.agent.prompt(input); - } +/** + * Pi `Agent` configured for Kernel browser computer use. + * + * Use this class when you want direct access to the lower-level pi agent state, + * queues, event stream, and `state.model` mutation model. It resolves CUA model + * refs, installs provider-appropriate CUA tools by default, and keeps those + * defaults in sync when `agent.state.model` changes. + */ +export class CuaAgent extends Agent { + private readonly runtime: CuaRuntimeController; + private stateProxy?: CuaAgentState; - steer(message: AgentMessage): void { - this.agent.steer(message); - } + constructor(options: CuaAgentOptions) { + const { browser, client, initialState, onPayload, streamFn, prepareNextTurn, ...agentOptions } = options; + const runtime = new CuaRuntimeController({ + browser, + client, + model: initialState.model, + tools: initialState.tools, + systemPrompt: initialState.systemPrompt, + onPayload, + }); + const wrappedStreamFn: StreamFn = (model, context, streamOptions) => + (streamFn ?? streamSimple)(model, context, { + ...streamOptions, + onPayload: runtime.onPayloadFor(model as Model), + }); - followUp(message: AgentMessage): void { - this.agent.followUp(message); - } + super({ + ...agentOptions, + getApiKey: agentOptions.getApiKey ?? getCuaEnvApiKey, + streamFn: wrappedStreamFn, + initialState: { + ...initialState, + model: runtime.model, + tools: runtime.tools(), + systemPrompt: initialState.systemPrompt ?? runtime.systemPrompt, + }, + }); - async continue(): Promise { - await this.agent.continue(); - } + this.runtime = runtime; + /** + * pi calls `prepareNextTurn` between provider requests. Wrapping it lets CUA + * honor any user-provided turn update while also refreshing provider-specific + * defaults if that update changes the model. + */ + this.prepareNextTurn = async (signal: AbortSignal | undefined) => { + const update = await prepareNextTurn?.(signal); + if (update?.model) { + this.applyRuntime(update.model as CuaRuntimeInput); + } + + const state = super.state; + const context = update?.context ?? { + systemPrompt: state.systemPrompt, + messages: state.messages.slice(), + tools: state.tools.slice(), + }; - clearSteeringQueue(): void { - this.agent.clearSteeringQueue(); + return { + ...update, + model: state.model, + context: { + ...context, + systemPrompt: this.runtime.ownsSystemPrompt ? state.systemPrompt : context.systemPrompt, + tools: this.runtime.ownsTools ? state.tools.slice() : context.tools, + }, + }; + }; } - clearFollowUpQueue(): void { - this.agent.clearFollowUpQueue(); + /** + * Return a state proxy so `agent.state.model = "provider:model"` can behave + * like pi's normal mutable state while also re-resolving CUA tools, prompt, + * and payload hooks for the selected provider. + */ + override get state(): CuaAgentState { + if (!this.stateProxy) { + this.stateProxy = new Proxy(super.state, { + set: (target, prop, value, receiver) => { + if (prop === "model") { + this.applyRuntime(value as CuaRuntimeInput); + return true; + } + return Reflect.set(target, prop, value, receiver); + }, + }) as CuaAgentState; + } + return this.stateProxy; } - clearAllQueues(): void { - this.agent.clearAllQueues(); + private applyRuntime(model: CuaRuntimeInput): void { + this.runtime.setModel(model); + const state = super.state; + state.model = this.runtime.model; + if (this.runtime.ownsTools) { + state.tools = this.runtime.tools(); + } + if (this.runtime.ownsSystemPrompt) { + state.systemPrompt = this.runtime.systemPrompt; + } } +} + +/** + * Pi `AgentHarness` configured for Kernel browser computer use. + * + * Use this class when you want pi's higher-level harness APIs for sessions, + * resources, prompt templates, queue events, compaction, and model selection. + * It installs provider CUA tools by default and keeps CUA-owned runtime + * defaults in sync through `setModel()`. + */ +export class CuaAgentHarness< + TSkill extends Skill = Skill, + TPromptTemplate extends PromptTemplate = PromptTemplate, +> extends AgentHarness { + private readonly runtime: CuaRuntimeController; + private requestedActiveToolNames?: string[]; - abort(): void { - this.agent.abort(); + constructor(options: CuaAgentHarnessOptions) { + const { + browser, + client, + model, + tools, + systemPrompt, + getApiKeyAndHeaders, + onPayload, + activeToolNames, + ...harnessOptions + } = options; + const runtime = new CuaRuntimeController({ browser, client, model, tools, systemPrompt, onPayload }); + const resolvedTools = runtime.tools(); + + super({ + ...harnessOptions, + model: runtime.model, + tools: resolvedTools, + systemPrompt: systemPrompt ?? (() => runtime.systemPrompt), + getApiKeyAndHeaders: + getApiKeyAndHeaders ?? + (async (requestModel: Model) => { + const apiKey = getCuaEnvApiKey(requestModel.provider); + return apiKey ? { apiKey } : undefined; + }), + activeToolNames: activeToolNames ?? resolvedTools.map((tool) => tool.name), + }); + + this.runtime = runtime; + this.requestedActiveToolNames = activeToolNames; + this.on("before_provider_payload", async ({ model, payload }: { model: Model; payload: unknown }) => { + const onPayload = this.runtime.onPayloadFor(model as Model); + if (!onPayload) return { payload }; + return { payload: (await onPayload(payload, model)) ?? payload }; + }); } - async waitForIdle(): Promise { - await this.agent.waitForIdle(); + /** + * Mirror pi `AgentHarness.setModel()` while accepting CUA model refs. + * + * The override refreshes CUA-owned tools before delegating to pi so the + * harness snapshot and session model-change entry are written with the + * concrete model selected by `@onkernel/cua-ai`. + */ + override async setModel(model: CuaRuntimeInput): Promise { + this.runtime.setModel(model); + if (this.runtime.ownsTools) { + const tools = this.runtime.tools(); + await super.setTools(tools, this.requestedActiveToolNames ?? tools.map((tool) => tool.name)); + } + await super.setModel(this.runtime.model); } - getTranscript(): AgentMessage[] { - return [...this.agent.state.messages]; + override async setActiveTools(toolNames: string[]): Promise { + await super.setActiveTools(toolNames); + this.requestedActiveToolNames = [...toolNames]; } } diff --git a/packages/agent/src/index.ts b/packages/agent/src/index.ts index f0d2fc50..57a49ccb 100644 --- a/packages/agent/src/index.ts +++ b/packages/agent/src/index.ts @@ -1,4 +1,4 @@ -export * from "@earendil-works/pi-agent-core"; +export * from "./vendor/pi-agent-core/index"; export type { KernelBrowser } from "./translator/translator"; export { createCuaComputerTools } from "./tools"; @@ -10,5 +10,5 @@ export type { SupportedCuaExecutorToolName, } from "./tools"; export { SUPPORTED_CUA_EXECUTOR_TOOL_NAMES } from "./tools"; -export { CuaAgent, CuaHarness } from "./agent"; -export type { CuaAgentOptions, CuaHarnessOptions } from "./agent"; +export { CuaAgent, CuaAgentHarness } from "./agent"; +export type { CuaAgentHarnessOptions, CuaAgentOptions, CuaAgentState } from "./agent"; diff --git a/packages/agent/src/tools.ts b/packages/agent/src/tools.ts index 594fad80..16e63af4 100644 --- a/packages/agent/src/tools.ts +++ b/packages/agent/src/tools.ts @@ -1,5 +1,4 @@ import type Kernel from "@onkernel/sdk"; -import type { AgentTool, AgentToolResult } from "@earendil-works/pi-agent-core"; import type { ImageContent, TextContent, Tool } from "@earendil-works/pi-ai"; import { CUA_BATCH_TOOL_NAME, @@ -10,6 +9,7 @@ import { type CuaNavigationInput, } from "@onkernel/cua-ai"; import { InternalComputerTranslator, type KernelBrowser } from "./translator/translator"; +import type { AgentTool, AgentToolResult } from "./vendor/pi-agent-core/index"; export interface ComputerToolOptions { browser: KernelBrowser; diff --git a/packages/agent/src/translator/translator.ts b/packages/agent/src/translator/translator.ts index 4e1d1bbb..cf1d9476 100644 --- a/packages/agent/src/translator/translator.ts +++ b/packages/agent/src/translator/translator.ts @@ -25,8 +25,8 @@ export class InternalComputerTranslator { async currentUrl(): Promise { await this.runKernelBatch([ - { type: "press_key", press_key: { keys: ["Control_L", "l"] } }, - { type: "press_key", press_key: { keys: ["Control_L", "c"] } }, + keypress(["Control", "l"]), + keypress(["Control", "c"]), ]); const response = await this.client.browsers.computer.readClipboard(this.sessionId); return (response.text ?? "").trim(); @@ -67,18 +67,18 @@ export class InternalComputerTranslator { } if (type === "goto") { pending.push( - { type: "press_key", press_key: { keys: ["Control_L", "l"] } }, + keypress(["Control", "l"]), { type: "type_text", type_text: { text: stringOr(action.url, "") } }, - { type: "press_key", press_key: { keys: ["Enter"] } }, + keypress(["Enter"]), ); continue; } if (type === "back") { - pending.push({ type: "press_key", press_key: { keys: ["Alt_L", "Left"] } }); + pending.push(keypress(["Alt", "Left"])); continue; } if (type === "forward") { - pending.push({ type: "press_key", press_key: { keys: ["Alt_L", "Right"] } }); + pending.push(keypress(["Alt", "Right"])); continue; } pending.push(toSdkAction(type, action)); @@ -133,7 +133,7 @@ function toSdkAction(type: string, action: ModelAction): KernelBatchAction { case "type": return { type: "type_text", type_text: { text: typeof action.text === "string" ? action.text : "" } }; case "keypress": - return { type: "press_key", press_key: { keys: toStringArray(action.keys) } }; + return keypress(toStringArray(action.keys)); case "scroll": return { type: "scroll", @@ -188,6 +188,67 @@ function toStringArray(value: unknown): string[] { return Array.isArray(value) ? value.filter((item): item is string => typeof item === "string") : []; } +function keypress(keys: string[]): KernelBatchAction { + const translated = translateKeys(keys); + const pressedKeys = translated.filter((key) => !isModifierKey(key)); + const holdKeys = pressedKeys.length > 0 ? translated.filter(isModifierKey) : translated.slice(0, -1); + return { + type: "press_key", + press_key: { + keys: pressedKeys.length > 0 ? pressedKeys : translated.slice(-1), + ...(holdKeys.length > 0 ? { hold_keys: holdKeys } : {}), + }, + }; +} + +const KEY_ALIASES: Record = { + ctrl: "Control_L", + control: "Control_L", + control_l: "Control_L", + controlleft: "Control_L", + alt: "Alt_L", + alt_l: "Alt_L", + altleft: "Alt_L", + shift: "Shift_L", + shift_l: "Shift_L", + shiftleft: "Shift_L", + meta: "Super_L", + super: "Super_L", + cmd: "Super_L", + command: "Super_L", + enter: "Return", + return: "Return", + escape: "Escape", + esc: "Escape", + backspace: "BackSpace", + delete: "Delete", + tab: "Tab", + space: "space", + left: "Left", + right: "Right", + up: "Up", + down: "Down", +}; + +function translateKeys(keys: string[]): string[] { + return keys.flatMap((key) => + key + .split("+") + .map((part) => part.trim()) + .filter(Boolean) + .map((part) => { + const alias = KEY_ALIASES[part.replace(/[-\s]/g, "_").toLowerCase()]; + if (alias) return alias; + if (part.length === 1 && part >= "A" && part <= "Z") return part.toLowerCase(); + return part; + }), + ); +} + +function isModifierKey(key: string): boolean { + return key === "Control_L" || key === "Alt_L" || key === "Shift_L" || key === "Super_L"; +} + function toPath(value: unknown): Array<[number, number]> { if (!Array.isArray(value)) return []; return value.map((point) => toPathPoint(point)); diff --git a/packages/agent/src/vendor/pi-agent-core/LICENSE b/packages/agent/src/vendor/pi-agent-core/LICENSE new file mode 100644 index 00000000..b0a8e9b8 --- /dev/null +++ b/packages/agent/src/vendor/pi-agent-core/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025 Mario Zechner + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. \ No newline at end of file diff --git a/packages/agent/src/vendor/pi-agent-core/README.md b/packages/agent/src/vendor/pi-agent-core/README.md new file mode 100644 index 00000000..ee873d6d --- /dev/null +++ b/packages/agent/src/vendor/pi-agent-core/README.md @@ -0,0 +1,15 @@ +# Vendored pi agent core + +These files are copied from `earendil-works/pi` so `@onkernel/cua-agent` can use +pi's `AgentHarness` and `prepareNextTurn` support before the upstream npm +package includes them. + +Source: https://github.com/earendil-works/pi/tree/40c05f55391663024a6a05ad33249b616a04e7a1/packages/agent/src + +License: MIT. See `LICENSE`, copied from the same pinned commit. + +Regenerate with: + +```bash +npx tsx packages/agent/scripts/vendor-pi-agent-harness.ts +``` diff --git a/packages/agent/src/vendor/pi-agent-core/agent-loop.ts b/packages/agent/src/vendor/pi-agent-core/agent-loop.ts new file mode 100644 index 00000000..7226082a --- /dev/null +++ b/packages/agent/src/vendor/pi-agent-core/agent-loop.ts @@ -0,0 +1,718 @@ +/** + * Agent loop that works with AgentMessage throughout. + * Transforms to Message[] only at the LLM call boundary. + */ + +import { + type AssistantMessage, + type Context, + EventStream, + streamSimple, + type ToolResultMessage, + validateToolArguments, +} from "@earendil-works/pi-ai"; +import type { + AgentContext, + AgentEvent, + AgentLoopConfig, + AgentMessage, + AgentTool, + AgentToolCall, + AgentToolResult, + StreamFn, +} from "./types.js"; + +export type AgentEventSink = (event: AgentEvent) => Promise | void; + +/** + * Start an agent loop with a new prompt message. + * The prompt is added to the context and events are emitted for it. + */ +export function agentLoop( + prompts: AgentMessage[], + context: AgentContext, + config: AgentLoopConfig, + signal?: AbortSignal, + streamFn?: StreamFn, +): EventStream { + const stream = createAgentStream(); + + void runAgentLoop( + prompts, + context, + config, + async (event) => { + stream.push(event); + }, + signal, + streamFn, + ).then((messages) => { + stream.end(messages); + }); + + return stream; +} + +/** + * Continue an agent loop from the current context without adding a new message. + * Used for retries - context already has user message or tool results. + * + * **Important:** The last message in context must convert to a `user` or `toolResult` message + * via `convertToLlm`. If it doesn't, the LLM provider will reject the request. + * This cannot be validated here since `convertToLlm` is only called once per turn. + */ +export function agentLoopContinue( + context: AgentContext, + config: AgentLoopConfig, + signal?: AbortSignal, + streamFn?: StreamFn, +): EventStream { + if (context.messages.length === 0) { + throw new Error("Cannot continue: no messages in context"); + } + + if (context.messages[context.messages.length - 1].role === "assistant") { + throw new Error("Cannot continue from message role: assistant"); + } + + const stream = createAgentStream(); + + void runAgentLoopContinue( + context, + config, + async (event) => { + stream.push(event); + }, + signal, + streamFn, + ).then((messages) => { + stream.end(messages); + }); + + return stream; +} + +export async function runAgentLoop( + prompts: AgentMessage[], + context: AgentContext, + config: AgentLoopConfig, + emit: AgentEventSink, + signal?: AbortSignal, + streamFn?: StreamFn, +): Promise { + const newMessages: AgentMessage[] = [...prompts]; + const currentContext: AgentContext = { + ...context, + messages: [...context.messages, ...prompts], + }; + + await emit({ type: "agent_start" }); + await emit({ type: "turn_start" }); + for (const prompt of prompts) { + await emit({ type: "message_start", message: prompt }); + await emit({ type: "message_end", message: prompt }); + } + + await runLoop(currentContext, newMessages, config, signal, emit, streamFn); + return newMessages; +} + +export async function runAgentLoopContinue( + context: AgentContext, + config: AgentLoopConfig, + emit: AgentEventSink, + signal?: AbortSignal, + streamFn?: StreamFn, +): Promise { + if (context.messages.length === 0) { + throw new Error("Cannot continue: no messages in context"); + } + + if (context.messages[context.messages.length - 1].role === "assistant") { + throw new Error("Cannot continue from message role: assistant"); + } + + const newMessages: AgentMessage[] = []; + const currentContext: AgentContext = { ...context }; + + await emit({ type: "agent_start" }); + await emit({ type: "turn_start" }); + + await runLoop(currentContext, newMessages, config, signal, emit, streamFn); + return newMessages; +} + +function createAgentStream(): EventStream { + return new EventStream( + (event: AgentEvent) => event.type === "agent_end", + (event: AgentEvent) => (event.type === "agent_end" ? event.messages : []), + ); +} + +/** + * Main loop logic shared by agentLoop and agentLoopContinue. + */ +async function runLoop( + initialContext: AgentContext, + newMessages: AgentMessage[], + initialConfig: AgentLoopConfig, + signal: AbortSignal | undefined, + emit: AgentEventSink, + streamFn?: StreamFn, +): Promise { + let currentContext = initialContext; + let config = initialConfig; + let firstTurn = true; + // Check for steering messages at start (user may have typed while waiting) + let pendingMessages: AgentMessage[] = (await config.getSteeringMessages?.()) || []; + + // Outer loop: continues when queued follow-up messages arrive after agent would stop + while (true) { + let hasMoreToolCalls = true; + + // Inner loop: process tool calls and steering messages + while (hasMoreToolCalls || pendingMessages.length > 0) { + if (!firstTurn) { + await emit({ type: "turn_start" }); + } else { + firstTurn = false; + } + + // Process pending messages (inject before next assistant response) + if (pendingMessages.length > 0) { + for (const message of pendingMessages) { + await emit({ type: "message_start", message }); + await emit({ type: "message_end", message }); + currentContext.messages.push(message); + newMessages.push(message); + } + pendingMessages = []; + } + + // Stream assistant response + const message = await streamAssistantResponse(currentContext, config, signal, emit, streamFn); + newMessages.push(message); + + if (message.stopReason === "error" || message.stopReason === "aborted") { + await emit({ type: "turn_end", message, toolResults: [] }); + await emit({ type: "agent_end", messages: newMessages }); + return; + } + + // Check for tool calls + const toolCalls = message.content.filter((c) => c.type === "toolCall"); + + const toolResults: ToolResultMessage[] = []; + hasMoreToolCalls = false; + if (toolCalls.length > 0) { + const executedToolBatch = await executeToolCalls(currentContext, message, config, signal, emit); + toolResults.push(...executedToolBatch.messages); + hasMoreToolCalls = !executedToolBatch.terminate; + + for (const result of toolResults) { + currentContext.messages.push(result); + newMessages.push(result); + } + } + + await emit({ type: "turn_end", message, toolResults }); + + const nextTurnContext = { + message, + toolResults, + context: currentContext, + newMessages, + }; + const nextTurnSnapshot = await config.prepareNextTurn?.(nextTurnContext); + if (nextTurnSnapshot) { + currentContext = nextTurnSnapshot.context ?? currentContext; + config = { + ...config, + model: nextTurnSnapshot.model ?? config.model, + reasoning: + nextTurnSnapshot.thinkingLevel === undefined + ? config.reasoning + : nextTurnSnapshot.thinkingLevel === "off" + ? undefined + : nextTurnSnapshot.thinkingLevel, + }; + } + + if ( + await config.shouldStopAfterTurn?.({ + message, + toolResults, + context: currentContext, + newMessages, + }) + ) { + await emit({ type: "agent_end", messages: newMessages }); + return; + } + + pendingMessages = (await config.getSteeringMessages?.()) || []; + } + + // Agent would stop here. Check for follow-up messages. + const followUpMessages = (await config.getFollowUpMessages?.()) || []; + if (followUpMessages.length > 0) { + // Set as pending so inner loop processes them + pendingMessages = followUpMessages; + continue; + } + + // No more messages, exit + break; + } + + await emit({ type: "agent_end", messages: newMessages }); +} + +/** + * Stream an assistant response from the LLM. + * This is where AgentMessage[] gets transformed to Message[] for the LLM. + */ +async function streamAssistantResponse( + context: AgentContext, + config: AgentLoopConfig, + signal: AbortSignal | undefined, + emit: AgentEventSink, + streamFn?: StreamFn, +): Promise { + // Apply context transform if configured (AgentMessage[] → AgentMessage[]) + let messages = context.messages; + if (config.transformContext) { + messages = await config.transformContext(messages, signal); + } + + // Convert to LLM-compatible messages (AgentMessage[] → Message[]) + const llmMessages = await config.convertToLlm(messages); + + // Build LLM context + const llmContext: Context = { + systemPrompt: context.systemPrompt, + messages: llmMessages, + tools: context.tools, + }; + + const streamFunction = streamFn || streamSimple; + + // Resolve API key (important for expiring tokens) + const resolvedApiKey = + (config.getApiKey ? await config.getApiKey(config.model.provider) : undefined) || config.apiKey; + + const response = await streamFunction(config.model, llmContext, { + ...config, + apiKey: resolvedApiKey, + signal, + }); + + let partialMessage: AssistantMessage | null = null; + let addedPartial = false; + + for await (const event of response) { + switch (event.type) { + case "start": + partialMessage = event.partial; + context.messages.push(partialMessage); + addedPartial = true; + await emit({ type: "message_start", message: { ...partialMessage } }); + break; + + case "text_start": + case "text_delta": + case "text_end": + case "thinking_start": + case "thinking_delta": + case "thinking_end": + case "toolcall_start": + case "toolcall_delta": + case "toolcall_end": + if (partialMessage) { + partialMessage = event.partial; + context.messages[context.messages.length - 1] = partialMessage; + await emit({ + type: "message_update", + assistantMessageEvent: event, + message: { ...partialMessage }, + }); + } + break; + + case "done": + case "error": { + const finalMessage = await response.result(); + if (addedPartial) { + context.messages[context.messages.length - 1] = finalMessage; + } else { + context.messages.push(finalMessage); + } + if (!addedPartial) { + await emit({ type: "message_start", message: { ...finalMessage } }); + } + await emit({ type: "message_end", message: finalMessage }); + return finalMessage; + } + } + } + + const finalMessage = await response.result(); + if (addedPartial) { + context.messages[context.messages.length - 1] = finalMessage; + } else { + context.messages.push(finalMessage); + await emit({ type: "message_start", message: { ...finalMessage } }); + } + await emit({ type: "message_end", message: finalMessage }); + return finalMessage; +} + +/** + * Execute tool calls from an assistant message. + */ +async function executeToolCalls( + currentContext: AgentContext, + assistantMessage: AssistantMessage, + config: AgentLoopConfig, + signal: AbortSignal | undefined, + emit: AgentEventSink, +): Promise { + const toolCalls = assistantMessage.content.filter((c) => c.type === "toolCall"); + const hasSequentialToolCall = toolCalls.some( + (tc) => currentContext.tools?.find((t) => t.name === tc.name)?.executionMode === "sequential", + ); + if (config.toolExecution === "sequential" || hasSequentialToolCall) { + return executeToolCallsSequential(currentContext, assistantMessage, toolCalls, config, signal, emit); + } + return executeToolCallsParallel(currentContext, assistantMessage, toolCalls, config, signal, emit); +} + +type ExecutedToolCallBatch = { + messages: ToolResultMessage[]; + terminate: boolean; +}; + +async function executeToolCallsSequential( + currentContext: AgentContext, + assistantMessage: AssistantMessage, + toolCalls: AgentToolCall[], + config: AgentLoopConfig, + signal: AbortSignal | undefined, + emit: AgentEventSink, +): Promise { + const finalizedCalls: FinalizedToolCallOutcome[] = []; + const messages: ToolResultMessage[] = []; + + for (const toolCall of toolCalls) { + await emit({ + type: "tool_execution_start", + toolCallId: toolCall.id, + toolName: toolCall.name, + args: toolCall.arguments, + }); + + const preparation = await prepareToolCall(currentContext, assistantMessage, toolCall, config, signal); + let finalized: FinalizedToolCallOutcome; + if (preparation.kind === "immediate") { + finalized = { + toolCall, + result: preparation.result, + isError: preparation.isError, + }; + } else { + const executed = await executePreparedToolCall(preparation, signal, emit); + finalized = await finalizeExecutedToolCall( + currentContext, + assistantMessage, + preparation, + executed, + config, + signal, + ); + } + + await emitToolExecutionEnd(finalized, emit); + const toolResultMessage = createToolResultMessage(finalized); + await emitToolResultMessage(toolResultMessage, emit); + finalizedCalls.push(finalized); + messages.push(toolResultMessage); + } + + return { + messages, + terminate: shouldTerminateToolBatch(finalizedCalls), + }; +} + +async function executeToolCallsParallel( + currentContext: AgentContext, + assistantMessage: AssistantMessage, + toolCalls: AgentToolCall[], + config: AgentLoopConfig, + signal: AbortSignal | undefined, + emit: AgentEventSink, +): Promise { + const finalizedCalls: FinalizedToolCallEntry[] = []; + + for (const toolCall of toolCalls) { + await emit({ + type: "tool_execution_start", + toolCallId: toolCall.id, + toolName: toolCall.name, + args: toolCall.arguments, + }); + + const preparation = await prepareToolCall(currentContext, assistantMessage, toolCall, config, signal); + if (preparation.kind === "immediate") { + const finalized = { + toolCall, + result: preparation.result, + isError: preparation.isError, + } satisfies FinalizedToolCallOutcome; + await emitToolExecutionEnd(finalized, emit); + finalizedCalls.push(finalized); + continue; + } + + finalizedCalls.push(async () => { + const executed = await executePreparedToolCall(preparation, signal, emit); + const finalized = await finalizeExecutedToolCall( + currentContext, + assistantMessage, + preparation, + executed, + config, + signal, + ); + await emitToolExecutionEnd(finalized, emit); + return finalized; + }); + } + + const orderedFinalizedCalls = await Promise.all( + finalizedCalls.map((entry) => (typeof entry === "function" ? entry() : Promise.resolve(entry))), + ); + const messages: ToolResultMessage[] = []; + for (const finalized of orderedFinalizedCalls) { + const toolResultMessage = createToolResultMessage(finalized); + await emitToolResultMessage(toolResultMessage, emit); + messages.push(toolResultMessage); + } + + return { + messages, + terminate: shouldTerminateToolBatch(orderedFinalizedCalls), + }; +} + +type PreparedToolCall = { + kind: "prepared"; + toolCall: AgentToolCall; + tool: AgentTool; + args: unknown; +}; + +type ImmediateToolCallOutcome = { + kind: "immediate"; + result: AgentToolResult; + isError: boolean; +}; + +type ExecutedToolCallOutcome = { + result: AgentToolResult; + isError: boolean; +}; + +type FinalizedToolCallOutcome = { + toolCall: AgentToolCall; + result: AgentToolResult; + isError: boolean; +}; + +type FinalizedToolCallEntry = FinalizedToolCallOutcome | (() => Promise); + +function shouldTerminateToolBatch(finalizedCalls: FinalizedToolCallOutcome[]): boolean { + return finalizedCalls.length > 0 && finalizedCalls.every((finalized) => finalized.result.terminate === true); +} + +function prepareToolCallArguments(tool: AgentTool, toolCall: AgentToolCall): AgentToolCall { + if (!tool.prepareArguments) { + return toolCall; + } + const preparedArguments = tool.prepareArguments(toolCall.arguments); + if (preparedArguments === toolCall.arguments) { + return toolCall; + } + return { + ...toolCall, + arguments: preparedArguments as Record, + }; +} + +async function prepareToolCall( + currentContext: AgentContext, + assistantMessage: AssistantMessage, + toolCall: AgentToolCall, + config: AgentLoopConfig, + signal: AbortSignal | undefined, +): Promise { + const tool = currentContext.tools?.find((t) => t.name === toolCall.name); + if (!tool) { + return { + kind: "immediate", + result: createErrorToolResult(`Tool ${toolCall.name} not found`), + isError: true, + }; + } + + try { + const preparedToolCall = prepareToolCallArguments(tool, toolCall); + const validatedArgs = validateToolArguments(tool, preparedToolCall); + if (config.beforeToolCall) { + const beforeResult = await config.beforeToolCall( + { + assistantMessage, + toolCall, + args: validatedArgs, + context: currentContext, + }, + signal, + ); + if (beforeResult?.block) { + return { + kind: "immediate", + result: createErrorToolResult(beforeResult.reason || "Tool execution was blocked"), + isError: true, + }; + } + } + return { + kind: "prepared", + toolCall, + tool, + args: validatedArgs, + }; + } catch (error) { + return { + kind: "immediate", + result: createErrorToolResult(error instanceof Error ? error.message : String(error)), + isError: true, + }; + } +} + +async function executePreparedToolCall( + prepared: PreparedToolCall, + signal: AbortSignal | undefined, + emit: AgentEventSink, +): Promise { + const updateEvents: Promise[] = []; + + try { + const result = await prepared.tool.execute( + prepared.toolCall.id, + prepared.args as never, + signal, + (partialResult) => { + updateEvents.push( + Promise.resolve( + emit({ + type: "tool_execution_update", + toolCallId: prepared.toolCall.id, + toolName: prepared.toolCall.name, + args: prepared.toolCall.arguments, + partialResult, + }), + ), + ); + }, + ); + await Promise.all(updateEvents); + return { result, isError: false }; + } catch (error) { + await Promise.all(updateEvents); + return { + result: createErrorToolResult(error instanceof Error ? error.message : String(error)), + isError: true, + }; + } +} + +async function finalizeExecutedToolCall( + currentContext: AgentContext, + assistantMessage: AssistantMessage, + prepared: PreparedToolCall, + executed: ExecutedToolCallOutcome, + config: AgentLoopConfig, + signal: AbortSignal | undefined, +): Promise { + let result = executed.result; + let isError = executed.isError; + + if (config.afterToolCall) { + try { + const afterResult = await config.afterToolCall( + { + assistantMessage, + toolCall: prepared.toolCall, + args: prepared.args, + result, + isError, + context: currentContext, + }, + signal, + ); + if (afterResult) { + result = { + content: afterResult.content ?? result.content, + details: afterResult.details ?? result.details, + terminate: afterResult.terminate ?? result.terminate, + }; + isError = afterResult.isError ?? isError; + } + } catch (error) { + result = createErrorToolResult(error instanceof Error ? error.message : String(error)); + isError = true; + } + } + + return { + toolCall: prepared.toolCall, + result, + isError, + }; +} + +function createErrorToolResult(message: string): AgentToolResult { + return { + content: [{ type: "text", text: message }], + details: {}, + }; +} + +async function emitToolExecutionEnd(finalized: FinalizedToolCallOutcome, emit: AgentEventSink): Promise { + await emit({ + type: "tool_execution_end", + toolCallId: finalized.toolCall.id, + toolName: finalized.toolCall.name, + result: finalized.result, + isError: finalized.isError, + }); +} + +function createToolResultMessage(finalized: FinalizedToolCallOutcome): ToolResultMessage { + return { + role: "toolResult", + toolCallId: finalized.toolCall.id, + toolName: finalized.toolCall.name, + content: finalized.result.content, + details: finalized.result.details, + isError: finalized.isError, + timestamp: Date.now(), + }; +} + +async function emitToolResultMessage(toolResultMessage: ToolResultMessage, emit: AgentEventSink): Promise { + await emit({ type: "message_start", message: toolResultMessage }); + await emit({ type: "message_end", message: toolResultMessage }); +} diff --git a/packages/agent/src/vendor/pi-agent-core/agent.ts b/packages/agent/src/vendor/pi-agent-core/agent.ts new file mode 100644 index 00000000..6eafd030 --- /dev/null +++ b/packages/agent/src/vendor/pi-agent-core/agent.ts @@ -0,0 +1,553 @@ +import { + type ImageContent, + type Message, + type Model, + type SimpleStreamOptions, + streamSimple, + type TextContent, + type ThinkingBudgets, + type Transport, +} from "@earendil-works/pi-ai"; +import { runAgentLoop, runAgentLoopContinue } from "./agent-loop.js"; +import type { + AfterToolCallContext, + AfterToolCallResult, + AgentContext, + AgentEvent, + AgentLoopConfig, + AgentLoopTurnUpdate, + AgentMessage, + AgentState, + AgentTool, + BeforeToolCallContext, + BeforeToolCallResult, + StreamFn, + ToolExecutionMode, +} from "./types.js"; + +function defaultConvertToLlm(messages: AgentMessage[]): Message[] { + return messages.filter( + (message) => message.role === "user" || message.role === "assistant" || message.role === "toolResult", + ); +} + +const EMPTY_USAGE = { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, +}; + +const DEFAULT_MODEL = { + id: "unknown", + name: "unknown", + api: "unknown", + provider: "unknown", + baseUrl: "", + reasoning: false, + input: [], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 0, + maxTokens: 0, +} satisfies Model; + +export type QueueMode = "all" | "one-at-a-time"; + +type MutableAgentState = Omit & { + isStreaming: boolean; + streamingMessage?: AgentMessage; + pendingToolCalls: Set; + errorMessage?: string; +}; + +function createMutableAgentState( + initialState?: Partial>, +): MutableAgentState { + let tools = initialState?.tools?.slice() ?? []; + let messages = initialState?.messages?.slice() ?? []; + + return { + systemPrompt: initialState?.systemPrompt ?? "", + model: initialState?.model ?? DEFAULT_MODEL, + thinkingLevel: initialState?.thinkingLevel ?? "off", + get tools() { + return tools; + }, + set tools(nextTools: AgentTool[]) { + tools = nextTools.slice(); + }, + get messages() { + return messages; + }, + set messages(nextMessages: AgentMessage[]) { + messages = nextMessages.slice(); + }, + isStreaming: false, + streamingMessage: undefined, + pendingToolCalls: new Set(), + errorMessage: undefined, + }; +} + +/** Options for constructing an {@link Agent}. */ +export interface AgentOptions { + initialState?: Partial>; + convertToLlm?: (messages: AgentMessage[]) => Message[] | Promise; + transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise; + streamFn?: StreamFn; + getApiKey?: (provider: string) => Promise | string | undefined; + onPayload?: SimpleStreamOptions["onPayload"]; + onResponse?: SimpleStreamOptions["onResponse"]; + beforeToolCall?: (context: BeforeToolCallContext, signal?: AbortSignal) => Promise; + afterToolCall?: (context: AfterToolCallContext, signal?: AbortSignal) => Promise; + prepareNextTurn?: ( + signal?: AbortSignal, + ) => Promise | AgentLoopTurnUpdate | undefined; + steeringMode?: QueueMode; + followUpMode?: QueueMode; + sessionId?: string; + thinkingBudgets?: ThinkingBudgets; + transport?: Transport; + maxRetryDelayMs?: number; + toolExecution?: ToolExecutionMode; +} + +class PendingMessageQueue { + private messages: AgentMessage[] = []; + + constructor(public mode: QueueMode) {} + + enqueue(message: AgentMessage): void { + this.messages.push(message); + } + + hasItems(): boolean { + return this.messages.length > 0; + } + + drain(): AgentMessage[] { + if (this.mode === "all") { + const drained = this.messages.slice(); + this.messages = []; + return drained; + } + + const first = this.messages[0]; + if (!first) { + return []; + } + this.messages = this.messages.slice(1); + return [first]; + } + + clear(): void { + this.messages = []; + } +} + +type ActiveRun = { + promise: Promise; + resolve: () => void; + abortController: AbortController; +}; + +/** + * Stateful wrapper around the low-level agent loop. + * + * `Agent` owns the current transcript, emits lifecycle events, executes tools, + * and exposes queueing APIs for steering and follow-up messages. + */ +export class Agent { + private _state: MutableAgentState; + private readonly listeners = new Set<(event: AgentEvent, signal: AbortSignal) => Promise | void>(); + private readonly steeringQueue: PendingMessageQueue; + private readonly followUpQueue: PendingMessageQueue; + + public convertToLlm: (messages: AgentMessage[]) => Message[] | Promise; + public transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise; + public streamFn: StreamFn; + public getApiKey?: (provider: string) => Promise | string | undefined; + public onPayload?: SimpleStreamOptions["onPayload"]; + public onResponse?: SimpleStreamOptions["onResponse"]; + public beforeToolCall?: ( + context: BeforeToolCallContext, + signal?: AbortSignal, + ) => Promise; + public afterToolCall?: ( + context: AfterToolCallContext, + signal?: AbortSignal, + ) => Promise; + public prepareNextTurn?: ( + signal?: AbortSignal, + ) => Promise | AgentLoopTurnUpdate | undefined; + private activeRun?: ActiveRun; + /** Session identifier forwarded to providers for cache-aware backends. */ + public sessionId?: string; + /** Optional per-level thinking token budgets forwarded to the stream function. */ + public thinkingBudgets?: ThinkingBudgets; + /** Preferred transport forwarded to the stream function. */ + public transport: Transport; + /** Optional cap for provider-requested retry delays. */ + public maxRetryDelayMs?: number; + /** Tool execution strategy for assistant messages that contain multiple tool calls. */ + public toolExecution: ToolExecutionMode; + + constructor(options: AgentOptions = {}) { + this._state = createMutableAgentState(options.initialState); + this.convertToLlm = options.convertToLlm ?? defaultConvertToLlm; + this.transformContext = options.transformContext; + this.streamFn = options.streamFn ?? streamSimple; + this.getApiKey = options.getApiKey; + this.onPayload = options.onPayload; + this.onResponse = options.onResponse; + this.beforeToolCall = options.beforeToolCall; + this.afterToolCall = options.afterToolCall; + this.prepareNextTurn = options.prepareNextTurn; + this.steeringQueue = new PendingMessageQueue(options.steeringMode ?? "one-at-a-time"); + this.followUpQueue = new PendingMessageQueue(options.followUpMode ?? "one-at-a-time"); + this.sessionId = options.sessionId; + this.thinkingBudgets = options.thinkingBudgets; + this.transport = options.transport ?? "auto"; + this.maxRetryDelayMs = options.maxRetryDelayMs; + this.toolExecution = options.toolExecution ?? "parallel"; + } + + /** + * Subscribe to agent lifecycle events. + * + * Listener promises are awaited in subscription order and are included in + * the current run's settlement. Listeners also receive the active abort + * signal for the current run. + * + * `agent_end` is the final emitted event for a run, but the agent does not + * become idle until all awaited listeners for that event have settled. + */ + subscribe(listener: (event: AgentEvent, signal: AbortSignal) => Promise | void): () => void { + this.listeners.add(listener); + return () => this.listeners.delete(listener); + } + + /** + * Current agent state. + * + * Assigning `state.tools` or `state.messages` copies the provided top-level array. + */ + get state(): AgentState { + return this._state; + } + + /** Controls how queued steering messages are drained. */ + set steeringMode(mode: QueueMode) { + this.steeringQueue.mode = mode; + } + + get steeringMode(): QueueMode { + return this.steeringQueue.mode; + } + + /** Controls how queued follow-up messages are drained. */ + set followUpMode(mode: QueueMode) { + this.followUpQueue.mode = mode; + } + + get followUpMode(): QueueMode { + return this.followUpQueue.mode; + } + + /** Queue a message to be injected after the current assistant turn finishes. */ + steer(message: AgentMessage): void { + this.steeringQueue.enqueue(message); + } + + /** Queue a message to run only after the agent would otherwise stop. */ + followUp(message: AgentMessage): void { + this.followUpQueue.enqueue(message); + } + + /** Remove all queued steering messages. */ + clearSteeringQueue(): void { + this.steeringQueue.clear(); + } + + /** Remove all queued follow-up messages. */ + clearFollowUpQueue(): void { + this.followUpQueue.clear(); + } + + /** Remove all queued steering and follow-up messages. */ + clearAllQueues(): void { + this.clearSteeringQueue(); + this.clearFollowUpQueue(); + } + + /** Returns true when either queue still contains pending messages. */ + hasQueuedMessages(): boolean { + return this.steeringQueue.hasItems() || this.followUpQueue.hasItems(); + } + + /** Active abort signal for the current run, if any. */ + get signal(): AbortSignal | undefined { + return this.activeRun?.abortController.signal; + } + + /** Abort the current run, if one is active. */ + abort(): void { + this.activeRun?.abortController.abort(); + } + + /** + * Resolve when the current run and all awaited event listeners have finished. + * + * This resolves after `agent_end` listeners settle. + */ + waitForIdle(): Promise { + return this.activeRun?.promise ?? Promise.resolve(); + } + + /** Clear transcript state, runtime state, and queued messages. */ + reset(): void { + this._state.messages = []; + this._state.isStreaming = false; + this._state.streamingMessage = undefined; + this._state.pendingToolCalls = new Set(); + this._state.errorMessage = undefined; + this.clearFollowUpQueue(); + this.clearSteeringQueue(); + } + + /** Start a new prompt from text, a single message, or a batch of messages. */ + async prompt(message: AgentMessage | AgentMessage[]): Promise; + async prompt(input: string, images?: ImageContent[]): Promise; + async prompt(input: string | AgentMessage | AgentMessage[], images?: ImageContent[]): Promise { + if (this.activeRun) { + throw new Error( + "Agent is already processing a prompt. Use steer() or followUp() to queue messages, or wait for completion.", + ); + } + const messages = this.normalizePromptInput(input, images); + await this.runPromptMessages(messages); + } + + /** Continue from the current transcript. The last message must be a user or tool-result message. */ + async continue(): Promise { + if (this.activeRun) { + throw new Error("Agent is already processing. Wait for completion before continuing."); + } + + const lastMessage = this._state.messages[this._state.messages.length - 1]; + if (!lastMessage) { + throw new Error("No messages to continue from"); + } + + if (lastMessage.role === "assistant") { + const queuedSteering = this.steeringQueue.drain(); + if (queuedSteering.length > 0) { + await this.runPromptMessages(queuedSteering, { skipInitialSteeringPoll: true }); + return; + } + + const queuedFollowUps = this.followUpQueue.drain(); + if (queuedFollowUps.length > 0) { + await this.runPromptMessages(queuedFollowUps); + return; + } + + throw new Error("Cannot continue from message role: assistant"); + } + + await this.runContinuation(); + } + + private normalizePromptInput( + input: string | AgentMessage | AgentMessage[], + images?: ImageContent[], + ): AgentMessage[] { + if (Array.isArray(input)) { + return input; + } + + if (typeof input !== "string") { + return [input]; + } + + const content: Array = [{ type: "text", text: input }]; + if (images && images.length > 0) { + content.push(...images); + } + return [{ role: "user", content, timestamp: Date.now() }]; + } + + private async runPromptMessages( + messages: AgentMessage[], + options: { skipInitialSteeringPoll?: boolean } = {}, + ): Promise { + await this.runWithLifecycle(async (signal) => { + await runAgentLoop( + messages, + this.createContextSnapshot(), + this.createLoopConfig(options), + (event) => this.processEvents(event), + signal, + this.streamFn, + ); + }); + } + + private async runContinuation(): Promise { + await this.runWithLifecycle(async (signal) => { + await runAgentLoopContinue( + this.createContextSnapshot(), + this.createLoopConfig(), + (event) => this.processEvents(event), + signal, + this.streamFn, + ); + }); + } + + private createContextSnapshot(): AgentContext { + return { + systemPrompt: this._state.systemPrompt, + messages: this._state.messages.slice(), + tools: this._state.tools.slice(), + }; + } + + private createLoopConfig(options: { skipInitialSteeringPoll?: boolean } = {}): AgentLoopConfig { + let skipInitialSteeringPoll = options.skipInitialSteeringPoll === true; + return { + model: this._state.model, + reasoning: this._state.thinkingLevel === "off" ? undefined : this._state.thinkingLevel, + sessionId: this.sessionId, + onPayload: this.onPayload, + onResponse: this.onResponse, + transport: this.transport, + thinkingBudgets: this.thinkingBudgets, + maxRetryDelayMs: this.maxRetryDelayMs, + toolExecution: this.toolExecution, + beforeToolCall: this.beforeToolCall, + afterToolCall: this.afterToolCall, + prepareNextTurn: this.prepareNextTurn ? async () => await this.prepareNextTurn?.(this.signal) : undefined, + convertToLlm: this.convertToLlm, + transformContext: this.transformContext, + getApiKey: this.getApiKey, + getSteeringMessages: async () => { + if (skipInitialSteeringPoll) { + skipInitialSteeringPoll = false; + return []; + } + return this.steeringQueue.drain(); + }, + getFollowUpMessages: async () => this.followUpQueue.drain(), + }; + } + + private async runWithLifecycle(executor: (signal: AbortSignal) => Promise): Promise { + if (this.activeRun) { + throw new Error("Agent is already processing."); + } + + const abortController = new AbortController(); + let resolvePromise = () => {}; + const promise = new Promise((resolve) => { + resolvePromise = resolve; + }); + this.activeRun = { promise, resolve: resolvePromise, abortController }; + + this._state.isStreaming = true; + this._state.streamingMessage = undefined; + this._state.errorMessage = undefined; + + try { + await executor(abortController.signal); + } catch (error) { + await this.handleRunFailure(error, abortController.signal.aborted); + } finally { + this.finishRun(); + } + } + + private async handleRunFailure(error: unknown, aborted: boolean): Promise { + const failureMessage = { + role: "assistant", + content: [{ type: "text", text: "" }], + api: this._state.model.api, + provider: this._state.model.provider, + model: this._state.model.id, + usage: EMPTY_USAGE, + stopReason: aborted ? "aborted" : "error", + errorMessage: error instanceof Error ? error.message : String(error), + timestamp: Date.now(), + } satisfies AgentMessage; + await this.processEvents({ type: "message_start", message: failureMessage }); + await this.processEvents({ type: "message_end", message: failureMessage }); + await this.processEvents({ type: "turn_end", message: failureMessage, toolResults: [] }); + await this.processEvents({ type: "agent_end", messages: [failureMessage] }); + } + + private finishRun(): void { + this._state.isStreaming = false; + this._state.streamingMessage = undefined; + this._state.pendingToolCalls = new Set(); + this.activeRun?.resolve(); + this.activeRun = undefined; + } + + /** + * Reduce internal state for a loop event, then await listeners. + * + * `agent_end` only means no further loop events will be emitted. The run is + * considered idle later, after all awaited listeners for `agent_end` finish + * and `finishRun()` clears runtime-owned state. + */ + private async processEvents(event: AgentEvent): Promise { + switch (event.type) { + case "message_start": + this._state.streamingMessage = event.message; + break; + + case "message_update": + this._state.streamingMessage = event.message; + break; + + case "message_end": + this._state.streamingMessage = undefined; + this._state.messages.push(event.message); + break; + + case "tool_execution_start": { + const pendingToolCalls = new Set(this._state.pendingToolCalls); + pendingToolCalls.add(event.toolCallId); + this._state.pendingToolCalls = pendingToolCalls; + break; + } + + case "tool_execution_end": { + const pendingToolCalls = new Set(this._state.pendingToolCalls); + pendingToolCalls.delete(event.toolCallId); + this._state.pendingToolCalls = pendingToolCalls; + break; + } + + case "turn_end": + if (event.message.role === "assistant" && event.message.errorMessage) { + this._state.errorMessage = event.message.errorMessage; + } + break; + + case "agent_end": + this._state.streamingMessage = undefined; + break; + } + + const signal = this.activeRun?.abortController.signal; + if (!signal) { + throw new Error("Agent listener invoked outside active run"); + } + for (const listener of this.listeners) { + await listener(event, signal); + } + } +} diff --git a/packages/agent/src/vendor/pi-agent-core/harness/agent-harness.ts b/packages/agent/src/vendor/pi-agent-core/harness/agent-harness.ts new file mode 100644 index 00000000..879a3599 --- /dev/null +++ b/packages/agent/src/vendor/pi-agent-core/harness/agent-harness.ts @@ -0,0 +1,816 @@ +import { + type AssistantMessage, + type ImageContent, + type Model, + streamSimple, + type UserMessage, +} from "@earendil-works/pi-ai"; +import { Agent, type QueueMode } from "../agent.js"; +import type { AgentEvent, AgentMessage, AgentTool, ThinkingLevel } from "../types.js"; +import { collectEntriesForBranchSummary, generateBranchSummary } from "./compaction/branch-summarization.js"; +import { compact, DEFAULT_COMPACTION_SETTINGS, prepareCompaction } from "./compaction/compaction.js"; +import { formatPromptTemplateInvocation } from "./prompt-templates.js"; +import { formatSkillInvocation } from "./skills.js"; +import type { + AbortResult, + AgentHarnessEvent, + AgentHarnessEventResultMap, + AgentHarnessOptions, + AgentHarnessOwnEvent, + AgentHarnessPhase, + AgentHarnessResources, + AgentHarnessStreamOptions, + AgentHarnessStreamOptionsPatch, + ExecutionEnv, + NavigateTreeResult, + PendingSessionWrite, + PromptTemplate, + Session, + Skill, +} from "./types.js"; + +function createUserMessage(text: string, images?: ImageContent[]): UserMessage { + const content: Array<{ type: "text"; text: string } | ImageContent> = [{ type: "text", text }]; + if (images) content.push(...images); + return { role: "user", content, timestamp: Date.now() }; +} + +function cloneStreamOptions(streamOptions?: AgentHarnessStreamOptions): AgentHarnessStreamOptions { + return { + ...streamOptions, + headers: streamOptions?.headers ? { ...streamOptions.headers } : undefined, + metadata: streamOptions?.metadata ? { ...streamOptions.metadata } : undefined, + }; +} + +function mergeHeaders(...headers: Array | undefined>): Record | undefined { + const merged: Record = {}; + let hasHeaders = false; + for (const entry of headers) { + if (!entry) continue; + Object.assign(merged, entry); + hasHeaders = true; + } + return hasHeaders ? merged : undefined; +} + +function hasOwn(object: object, key: PropertyKey): boolean { + return Object.hasOwn(object, key); +} + +function applyStreamOptionsPatch( + base: AgentHarnessStreamOptions, + patch?: AgentHarnessStreamOptionsPatch, +): AgentHarnessStreamOptions { + const result = cloneStreamOptions(base); + if (!patch) return result; + + if (hasOwn(patch, "transport")) result.transport = patch.transport; + if (hasOwn(patch, "timeoutMs")) result.timeoutMs = patch.timeoutMs; + if (hasOwn(patch, "maxRetries")) result.maxRetries = patch.maxRetries; + if (hasOwn(patch, "maxRetryDelayMs")) result.maxRetryDelayMs = patch.maxRetryDelayMs; + if (hasOwn(patch, "cacheRetention")) result.cacheRetention = patch.cacheRetention; + + if (hasOwn(patch, "headers")) { + if (patch.headers === undefined) { + result.headers = undefined; + } else { + const headers = { ...(result.headers ?? {}) }; + for (const [key, value] of Object.entries(patch.headers)) { + if (value === undefined) delete headers[key]; + else headers[key] = value; + } + result.headers = Object.keys(headers).length > 0 ? headers : undefined; + } + } + + if (hasOwn(patch, "metadata")) { + if (patch.metadata === undefined) { + result.metadata = undefined; + } else { + const metadata = { ...(result.metadata ?? {}) }; + for (const [key, value] of Object.entries(patch.metadata)) { + if (value === undefined) delete metadata[key]; + else metadata[key] = value; + } + result.metadata = Object.keys(metadata).length > 0 ? metadata : undefined; + } + } + + return result; +} + +interface AgentHarnessTurnState< + TSkill extends Skill = Skill, + TPromptTemplate extends PromptTemplate = PromptTemplate, + TTool extends AgentTool = AgentTool, +> { + messages: AgentMessage[]; + resources: AgentHarnessResources; + streamOptions: AgentHarnessStreamOptions; + sessionId: string; + systemPrompt: string; + model: Model; + thinkingLevel: ThinkingLevel; + tools: TTool[]; + activeTools: TTool[]; +} + +export class AgentHarness< + TSkill extends Skill = Skill, + TPromptTemplate extends PromptTemplate = PromptTemplate, + TTool extends AgentTool = AgentTool, +> { + readonly agent: Agent; + readonly env: ExecutionEnv; + private session: Session; + private model: Model; + private thinkingLevel: ThinkingLevel; + private activeToolNames: string[]; + private nextTurnQueue: AgentMessage[] = []; + private phase: AgentHarnessPhase = "idle"; + private steerQueue: UserMessage[] = []; + private followUpQueue: UserMessage[] = []; + private pendingSessionWrites: PendingSessionWrite[] = []; + private resources: AgentHarnessResources; + private streamOptions: AgentHarnessStreamOptions; + private appliedStreamOptions: AgentHarnessStreamOptions = {}; + private appliedSessionId?: string; + private systemPrompt: AgentHarnessOptions["systemPrompt"]; + private getApiKeyAndHeaders?: AgentHarnessOptions["getApiKeyAndHeaders"]; + private tools = new Map(); + private listeners = new Set< + (event: AgentHarnessEvent, signal?: AbortSignal) => Promise | void + >(); + private hooks = new Map Promise | any>>(); + + constructor(options: AgentHarnessOptions) { + this.agent = new Agent({ + initialState: { + model: options.model, + thinkingLevel: options.thinkingLevel, + tools: options.tools ?? [], + }, + streamFn: async (model, context, streamOptions) => { + const auth = await this.getApiKeyAndHeaders?.(model); + const snapshotOptions: AgentHarnessStreamOptions = { + ...this.appliedStreamOptions, + headers: mergeHeaders(this.appliedStreamOptions.headers, auth?.headers), + }; + const requestOptions = await this.emitBeforeProviderRequest( + model, + this.appliedSessionId ?? "", + snapshotOptions, + ); + return streamSimple(model, context, { + cacheRetention: requestOptions.cacheRetention, + headers: requestOptions.headers, + maxRetries: requestOptions.maxRetries, + maxRetryDelayMs: requestOptions.maxRetryDelayMs, + metadata: requestOptions.metadata, + onPayload: async (payload) => await this.emitBeforeProviderPayload(model, payload), + onResponse: async (response) => { + const headers = { ...(response.headers as Record) }; + await this.emitOwn( + { type: "after_provider_response", status: response.status, headers }, + this.agent.signal, + ); + }, + reasoning: streamOptions?.reasoning, + signal: streamOptions?.signal, + sessionId: this.appliedSessionId, + timeoutMs: requestOptions.timeoutMs, + transport: requestOptions.transport, + apiKey: auth?.apiKey, + }); + }, + steeringMode: options.steeringMode, + followUpMode: options.followUpMode, + }); + this.env = options.env; + this.session = options.session; + this.resources = options.resources ?? {}; + this.streamOptions = cloneStreamOptions(options.streamOptions); + this.systemPrompt = options.systemPrompt; + this.getApiKeyAndHeaders = options.getApiKeyAndHeaders; + for (const tool of options.tools ?? []) { + this.tools.set(tool.name, tool); + } + this.model = options.model; + this.thinkingLevel = options.thinkingLevel ?? this.agent.state.thinkingLevel; + this.activeToolNames = options.activeToolNames ?? (options.tools ?? []).map((tool) => tool.name); + this.agent.state.model = this.model; + this.agent.state.thinkingLevel = this.thinkingLevel; + this.agent.transformContext = async (messages) => { + const result = await this.emitHook({ type: "context", messages: [...messages] }); + return result?.messages ?? messages; + }; + this.agent.beforeToolCall = async ({ toolCall, args }) => { + const result = await this.emitHook({ + type: "tool_call", + toolCallId: toolCall.id, + toolName: toolCall.name, + input: args as Record, + }); + return result ? { block: result.block, reason: result.reason } : undefined; + }; + this.agent.afterToolCall = async ({ toolCall, args, result, isError }) => { + const patch = await this.emitHook({ + type: "tool_result", + toolCallId: toolCall.id, + toolName: toolCall.name, + input: args as Record, + content: result.content, + details: result.details, + isError, + }); + return patch + ? { content: patch.content, details: patch.details, isError: patch.isError, terminate: patch.terminate } + : undefined; + }; + this.agent.prepareNextTurn = async () => { + await this.flushPendingSessionWrites(); + const turnState = await this.createTurnState(); + this.applyTurnState(turnState); + return { + context: { + systemPrompt: turnState.systemPrompt, + messages: turnState.messages.slice(), + tools: turnState.activeTools.slice(), + }, + model: turnState.model, + thinkingLevel: turnState.thinkingLevel, + }; + }; + this.agent.subscribe(async (event, signal) => { + await this.handleAgentEvent(event, signal); + }); + } + + 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( + event: Extract, + ): Promise { + const handlers = this.hooks.get(event.type as TType); + if (!handlers || handlers.size === 0) return undefined; + let lastResult: AgentHarnessEventResultMap[TType] | undefined; + for (const handler of handlers) { + const result = await handler(event); + if (result !== undefined) { + lastResult = result; + } + } + return lastResult; + } + + private async emitBeforeProviderRequest( + model: Model, + sessionId: string, + streamOptions: AgentHarnessStreamOptions, + ): Promise { + const handlers = this.hooks.get("before_provider_request"); + let current = cloneStreamOptions(streamOptions); + if (!handlers || handlers.size === 0) return current; + for (const handler of handlers) { + const result = await handler({ + type: "before_provider_request", + model, + sessionId, + streamOptions: cloneStreamOptions(current), + }); + if (result?.streamOptions) { + current = applyStreamOptionsPatch(current, result.streamOptions); + } + } + return current; + } + + private async emitBeforeProviderPayload(model: Model, payload: unknown): Promise { + const handlers = this.hooks.get("before_provider_payload"); + let current = payload; + if (!handlers || handlers.size === 0) return current; + for (const handler of handlers) { + const result = await handler({ type: "before_provider_payload", model, payload: current }); + if (result !== undefined) { + current = result.payload; + } + } + return current; + } + + private async emitQueueUpdate(): Promise { + await this.emitOwn({ + type: "queue_update", + steer: [...this.steerQueue], + followUp: [...this.followUpQueue], + nextTurn: [...this.nextTurnQueue], + }); + } + + private async createTurnState(): Promise> { + const context = await this.session.buildContext(); + const resources = this.getResources(); + const sessionMetadata = await this.session.getMetadata(); + const tools = [...this.tools.values()]; + const activeTools = this.activeToolNames + .map((name) => this.tools.get(name)) + .filter((tool): tool is TTool => tool !== undefined); + let systemPrompt = "You are a helpful assistant."; + if (typeof this.systemPrompt === "string") { + systemPrompt = this.systemPrompt; + } else if (this.systemPrompt) { + systemPrompt = await this.systemPrompt({ + env: this.env, + session: this.session, + model: this.model, + thinkingLevel: this.thinkingLevel, + activeTools, + resources, + }); + } + return { + messages: context.messages, + resources, + streamOptions: cloneStreamOptions(this.streamOptions), + sessionId: sessionMetadata.id, + systemPrompt, + model: this.model, + thinkingLevel: this.thinkingLevel, + tools, + activeTools, + }; + } + + private applyTurnState(turnState: AgentHarnessTurnState): void { + this.agent.state.messages = turnState.messages; + this.appliedStreamOptions = cloneStreamOptions(turnState.streamOptions); + this.appliedSessionId = turnState.sessionId; + this.agent.state.systemPrompt = turnState.systemPrompt; + this.agent.state.model = turnState.model; + this.agent.state.thinkingLevel = turnState.thinkingLevel; + this.agent.state.tools = turnState.activeTools; + } + + private validateToolNames(toolNames: string[]): void { + const missing = toolNames.filter((name) => !this.tools.has(name)); + if (missing.length > 0) throw new Error(`Unknown tool(s): ${missing.join(", ")}`); + } + + private async flushPendingSessionWrites(): Promise { + const writes = this.pendingSessionWrites; + this.pendingSessionWrites = []; + for (const write of writes) { + if (write.type === "message") { + await this.session.appendMessage(write.message); + } else if (write.type === "model_change") { + await this.session.appendModelChange(write.provider, write.modelId); + } else if (write.type === "thinking_level_change") { + await this.session.appendThinkingLevelChange(write.thinkingLevel); + } else if (write.type === "custom") { + await this.session.appendCustomEntry(write.customType, write.data); + } else if (write.type === "custom_message") { + await this.session.appendCustomMessageEntry(write.customType, write.content, write.display, write.details); + } else if (write.type === "label") { + await this.session.appendLabel(write.targetId, write.label); + } else if (write.type === "session_info") { + await this.session.appendSessionName(write.name ?? ""); + } + } + } + + private async handleAgentEvent(event: AgentEvent, signal?: AbortSignal): Promise { + await this.emitAny(event, signal); + if (event.type === "message_start" && event.message.role === "user") { + const steerIndex = this.steerQueue.indexOf(event.message); + if (steerIndex !== -1) { + this.steerQueue.splice(steerIndex, 1); + await this.emitQueueUpdate(); + } else { + const followUpIndex = this.followUpQueue.indexOf(event.message); + if (followUpIndex !== -1) { + this.followUpQueue.splice(followUpIndex, 1); + await this.emitQueueUpdate(); + } + } + } + if (event.type === "message_end") { + await this.session.appendMessage(event.message); + } + if (event.type === "turn_end") { + const hadPendingMutations = this.pendingSessionWrites.length > 0; + await this.flushPendingSessionWrites(); + await this.emitOwn({ + type: "save_point", + hadPendingMutations, + }); + } + if (event.type === "agent_end") { + await this.flushPendingSessionWrites(); + this.phase = "idle"; + await this.emitOwn({ type: "settled", nextTurnCount: this.nextTurnQueue.length }, signal); + } + } + + private async executeTurn( + turnState: AgentHarnessTurnState, + text: string, + options?: { images?: ImageContent[] }, + ): Promise { + this.applyTurnState(turnState); + const beforeLength = this.agent.state.messages.length; + let messages: AgentMessage[] = [createUserMessage(text, options?.images)]; + if (this.nextTurnQueue.length > 0) { + messages = [...this.nextTurnQueue, messages[0]!]; + this.nextTurnQueue = []; + await this.emitQueueUpdate(); + } + const beforeResult = await this.emitHook({ + type: "before_agent_start", + prompt: text, + images: options?.images, + systemPrompt: turnState.systemPrompt, + resources: turnState.resources, + }); + if (beforeResult?.messages) messages = [...beforeResult.messages, ...messages]; + if (beforeResult?.systemPrompt) this.agent.state.systemPrompt = beforeResult.systemPrompt; + try { + await this.agent.prompt(messages); + } finally { + await this.flushPendingSessionWrites(); + } + let response: AssistantMessage | undefined; + const newMessages = this.agent.state.messages.slice(beforeLength); + for (let i = newMessages.length - 1; i >= 0; i--) { + const message = newMessages[i]!; + if (message.role === "assistant") { + response = message; + break; + } + } + if (!response) throw new Error("AgentHarness prompt completed without an assistant message"); + return response; + } + + async prompt(text: string, options?: { images?: ImageContent[] }): Promise { + if (this.phase !== "idle") throw new Error("AgentHarness is busy"); + this.phase = "turn"; + try { + const turnState = await this.createTurnState(); + return await this.executeTurn(turnState, text, options); + } catch (error) { + this.phase = "idle"; + throw error; + } + } + + async skill(name: string, additionalInstructions?: string): Promise { + if (this.phase !== "idle") throw new Error("AgentHarness is busy"); + this.phase = "turn"; + try { + const turnState = await this.createTurnState(); + const skill = (turnState.resources.skills ?? []).find((candidate) => candidate.name === name); + if (!skill) throw new Error(`Unknown skill: ${name}`); + return await this.executeTurn(turnState, formatSkillInvocation(skill, additionalInstructions)); + } catch (error) { + this.phase = "idle"; + throw error; + } + } + + async promptFromTemplate(name: string, args: string[] = []): Promise { + if (this.phase !== "idle") throw new Error("AgentHarness is busy"); + this.phase = "turn"; + try { + const turnState = await this.createTurnState(); + const template = (turnState.resources.promptTemplates ?? []).find((candidate) => candidate.name === name); + if (!template) throw new Error(`Unknown prompt template: ${name}`); + return await this.executeTurn(turnState, formatPromptTemplateInvocation(template, args)); + } catch (error) { + this.phase = "idle"; + throw error; + } + } + + steer(text: string, options?: { images?: ImageContent[] }): void { + if (this.phase === "idle") throw new Error("Cannot steer while idle"); + const message = createUserMessage(text, options?.images); + this.steerQueue.push(message); + this.agent.steer(message); + void this.emitQueueUpdate(); + } + + followUp(text: string, options?: { images?: ImageContent[] }): void { + if (this.phase === "idle") throw new Error("Cannot follow up while idle"); + const message = createUserMessage(text, options?.images); + this.followUpQueue.push(message); + this.agent.followUp(message); + void this.emitQueueUpdate(); + } + + nextTurn(text: string, options?: { images?: ImageContent[] }): void { + this.nextTurnQueue.push(createUserMessage(text, options?.images)); + void this.emitQueueUpdate(); + } + + async appendMessage(message: AgentMessage): Promise { + if (this.phase === "idle") { + await this.session.appendMessage(message); + } else { + this.pendingSessionWrites.push({ type: "message", message }); + } + } + + async compact( + customInstructions?: string, + ): Promise<{ summary: string; firstKeptEntryId: string; tokensBefore: number; details?: unknown }> { + if (this.phase !== "idle") throw new Error("compact() requires idle harness"); + this.phase = "compaction"; + const model = this.model; + if (!model) throw new Error("No model set for compaction"); + const auth = await this.getApiKeyAndHeaders?.(model); + if (!auth) throw new Error("No auth available for compaction"); + const branchEntries = await this.session.getBranch(); + const preparation = prepareCompaction(branchEntries, DEFAULT_COMPACTION_SETTINGS); + if (!preparation) throw new Error("Nothing to compact"); + const hookResult = await this.emitHook({ + type: "session_before_compact", + preparation, + branchEntries, + customInstructions, + signal: new AbortController().signal, + }); + if (hookResult?.cancel) { + this.phase = "idle"; + throw new Error("Compaction cancelled"); + } + const provided = hookResult?.compaction; + const result = + provided ?? + (await compact( + preparation, + model, + auth.apiKey, + auth.headers, + customInstructions, + undefined, + this.thinkingLevel, + )); + const entryId = await this.session.appendCompaction( + result.summary, + result.firstKeptEntryId, + result.tokensBefore, + result.details, + provided !== undefined, + ); + const entry = await this.session.getEntry(entryId); + if (entry?.type === "compaction") { + await this.emitOwn({ type: "session_compact", compactionEntry: entry, fromHook: provided !== undefined }); + } + this.phase = "idle"; + return result; + } + + async navigateTree( + targetId: string, + options?: { summarize?: boolean; customInstructions?: string; replaceInstructions?: boolean; label?: string }, + ): Promise { + if (this.phase !== "idle") throw new Error("navigateTree() requires idle harness"); + this.phase = "branch_summary"; + const oldLeafId = await this.session.getLeafId(); + if (oldLeafId === targetId) { + this.phase = "idle"; + return { cancelled: false }; + } + const targetEntry = await this.session.getEntry(targetId); + if (!targetEntry) throw new Error(`Entry ${targetId} not found`); + const { entries, commonAncestorId } = await collectEntriesForBranchSummary(this.session, oldLeafId, targetId); + 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({ + type: "session_before_tree", + preparation, + signal, + }); + if (hookResult?.cancel) { + this.phase = "idle"; + return { cancelled: true }; + } + let summaryEntry: any | undefined; + let summaryText: string | undefined = hookResult?.summary?.summary; + let summaryDetails: unknown = hookResult?.summary?.details; + if (!summaryText && options?.summarize && entries.length > 0) { + const model = this.model; + if (!model) throw new Error("No model set for branch summary"); + const auth = await this.getApiKeyAndHeaders?.(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: hookResult?.customInstructions ?? options?.customInstructions, + replaceInstructions: hookResult?.replaceInstructions ?? options?.replaceInstructions, + }); + if (branchSummary.aborted) { + this.phase = "idle"; + 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 { readonly type: "text"; readonly 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 { readonly type: "text"; readonly text: string } => c.type === "text") + .map((c) => c.text) + .join(""); + } else { + newLeafId = targetId; + } + const summaryId = await this.session.moveTo( + newLeafId, + summaryText + ? { + summary: summaryText, + details: summaryDetails, + fromHook: hookResult?.summary !== undefined, + } + : undefined, + ); + if (summaryId) { + summaryEntry = await this.session.getEntry(summaryId); + } + await this.emitOwn({ + type: "session_tree", + newLeafId: await this.session.getLeafId(), + oldLeafId, + summaryEntry, + fromHook: hookResult?.summary !== undefined, + }); + this.phase = "idle"; + return { cancelled: false, editorText, summaryEntry }; + } + + async setModel(model: Model): Promise { + const previousModel = this.model; + this.model = model; + if (this.phase === "idle") { + this.agent.state.model = model; + await this.session.appendModelChange(model.provider, model.id); + } else { + this.pendingSessionWrites.push({ type: "model_change", provider: model.provider, modelId: model.id }); + } + await this.emitOwn({ type: "model_select", model, previousModel, source: "set" }); + } + + async setThinkingLevel(level: ThinkingLevel): Promise { + const previousLevel = this.thinkingLevel; + this.thinkingLevel = level; + if (this.phase === "idle") { + this.agent.state.thinkingLevel = level; + await this.session.appendThinkingLevelChange(level); + } else { + this.pendingSessionWrites.push({ type: "thinking_level_change", thinkingLevel: level }); + } + await this.emitOwn({ type: "thinking_level_select", level, previousLevel }); + } + + async setActiveTools(toolNames: string[]): Promise { + this.validateToolNames(toolNames); + this.activeToolNames = [...toolNames]; + if (this.phase === "idle") { + this.agent.state.tools = this.activeToolNames.map((name) => this.tools.get(name)!); + } + } + + get steeringMode(): QueueMode { + return this.agent.steeringMode; + } + + set steeringMode(mode: QueueMode) { + this.agent.steeringMode = mode; + } + + get followUpMode(): QueueMode { + return this.agent.followUpMode; + } + + set followUpMode(mode: QueueMode) { + this.agent.followUpMode = mode; + } + + getResources(): AgentHarnessResources { + return { + skills: this.resources.skills?.slice(), + promptTemplates: this.resources.promptTemplates?.slice(), + }; + } + + async setResources(resources: AgentHarnessResources): Promise { + const previousResources = this.getResources(); + this.resources = { + skills: resources.skills?.slice(), + promptTemplates: resources.promptTemplates?.slice(), + }; + await this.emitOwn({ type: "resources_update", resources: this.getResources(), previousResources }); + } + + getStreamOptions(): AgentHarnessStreamOptions { + return cloneStreamOptions(this.streamOptions); + } + + setStreamOptions(streamOptions: AgentHarnessStreamOptions): void { + this.streamOptions = cloneStreamOptions(streamOptions); + } + + async setTools(tools: TTool[], activeToolNames?: string[]): Promise { + this.tools = new Map(tools.map((tool) => [tool.name, tool])); + if (activeToolNames) { + this.validateToolNames(activeToolNames); + this.activeToolNames = [...activeToolNames]; + } else { + this.validateToolNames(this.activeToolNames); + } + if (this.phase === "idle") { + this.agent.state.tools = this.activeToolNames.map((name) => this.tools.get(name)!); + } + } + + async abort(): Promise { + const clearedSteer = [...this.steerQueue]; + const clearedFollowUp = [...this.followUpQueue]; + this.steerQueue = []; + this.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: AgentHarnessEvent, signal?: AbortSignal) => Promise | void, + ): () => void { + this.listeners.add(listener); + return () => this.listeners.delete(listener); + } + + on( + type: TType, + handler: ( + event: Extract, + ) => 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); + } +} diff --git a/packages/agent/src/vendor/pi-agent-core/harness/compaction/branch-summarization.ts b/packages/agent/src/vendor/pi-agent-core/harness/compaction/branch-summarization.ts new file mode 100644 index 00000000..e44bf458 --- /dev/null +++ b/packages/agent/src/vendor/pi-agent-core/harness/compaction/branch-summarization.ts @@ -0,0 +1,361 @@ +/** + * Branch summarization for tree navigation. + * + * When navigating to a different point in the session tree, this generates + * a summary of the branch being left so context isn't lost. + */ + +import type { ImageContent, Model, TextContent } from "@earendil-works/pi-ai"; +import { completeSimple } from "@earendil-works/pi-ai"; +import type { AgentMessage } from "../../types.js"; +import { + convertToLlm, + createBranchSummaryMessage, + createCompactionSummaryMessage, + createCustomMessage, +} from "../messages.js"; +import type { Session, SessionTreeEntry } from "../types.js"; +import { estimateTokens } from "./compaction.js"; +import { + computeFileLists, + createFileOps, + extractFileOpsFromMessage, + type FileOperations, + formatFileOperations, + SUMMARIZATION_SYSTEM_PROMPT, + serializeConversation, +} from "./utils.js"; + +// ============================================================================ +// Types +// ============================================================================ + +export interface BranchSummaryResult { + summary?: string; + readFiles?: string[]; + modifiedFiles?: string[]; + aborted?: boolean; + error?: string; +} + +/** Details stored in BranchSummaryEntry.details for file tracking */ +export interface BranchSummaryDetails { + readFiles: string[]; + modifiedFiles: string[]; +} + +export type { FileOperations } from "./utils.js"; + +export interface BranchPreparation { + /** Messages extracted for summarization, in chronological order */ + messages: AgentMessage[]; + /** File operations extracted from tool calls */ + fileOps: FileOperations; + /** Total estimated tokens in messages */ + totalTokens: number; +} + +export interface CollectEntriesResult { + /** Entries to summarize, in chronological order */ + entries: SessionTreeEntry[]; + /** Common ancestor between old and new position, if any */ + commonAncestorId: string | null; +} + +export interface GenerateBranchSummaryOptions { + /** Model to use for summarization */ + model: Model; + /** API key for the model */ + apiKey: string; + /** Request headers for the model */ + headers?: Record; + /** Abort signal for cancellation */ + signal: AbortSignal; + /** Optional custom instructions for summarization */ + customInstructions?: string; + /** If true, customInstructions replaces the default prompt instead of being appended */ + replaceInstructions?: boolean; + /** Tokens reserved for prompt + LLM response (default 16384) */ + reserveTokens?: number; +} + +// ============================================================================ +// Entry Collection +// ============================================================================ + +/** + * Collect entries that should be summarized when navigating from one position to another. + * + * Walks from oldLeafId back to the common ancestor with targetId, collecting entries + * along the way. Does NOT stop at compaction boundaries - those are included and their + * summaries become context. + * + * @param session - Session manager (read-only access) + * @param oldLeafId - Current position (where we're navigating from) + * @param targetId - Target position (where we're navigating to) + * @returns Entries to summarize and the common ancestor + */ +export async function collectEntriesForBranchSummary( + session: Session, + oldLeafId: string | null, + targetId: string, +): Promise { + // If no old position, nothing to summarize + if (!oldLeafId) { + return { entries: [], commonAncestorId: null }; + } + + // Find common ancestor (deepest node that's on both paths) + const oldPath = new Set((await session.getBranch(oldLeafId)).map((e) => e.id)); + const targetPath = await session.getBranch(targetId); + + // targetPath is root-first, so iterate backwards to find deepest common ancestor + 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; + } + } + + // Collect entries from old leaf back to common ancestor + const entries: SessionTreeEntry[] = []; + let current: string | null = oldLeafId; + + while (current && current !== commonAncestorId) { + const entry = await session.getEntry(current); + if (!entry) break; + entries.push(entry as SessionTreeEntry); + current = entry.parentId; + } + + // Reverse to get chronological order + entries.reverse(); + + return { entries, commonAncestorId }; +} + +// ============================================================================ +// Entry to Message Conversion +// ============================================================================ + +/** + * Extract AgentMessage from a session entry. + * Similar to getMessageFromEntry in compaction.ts but also handles compaction entries. + */ +function getMessageFromEntry(entry: SessionTreeEntry): AgentMessage | undefined { + switch (entry.type) { + case "message": + // Skip tool results - context is in assistant's tool call + if (entry.message.role === "toolResult") return undefined; + return entry.message as AgentMessage; + + case "custom_message": + return createCustomMessage( + entry.customType, + entry.content as string | (TextContent | ImageContent)[], + 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); + + // These don't contribute to conversation content + case "thinking_level_change": + case "model_change": + case "custom": + case "label": + case "session_info": + return undefined; + } +} + +/** + * Prepare entries for summarization with token budget. + * + * Walks entries from NEWEST to OLDEST, adding messages until we hit the token budget. + * This ensures we keep the most recent context when the branch is too long. + * + * Also collects file operations from: + * - Tool calls in assistant messages + * - Existing branch_summary entries' details (for cumulative tracking) + * + * @param entries - Entries in chronological order + * @param tokenBudget - Maximum tokens to include (0 = no limit) + */ +export function prepareBranchEntries(entries: SessionTreeEntry[], tokenBudget: number = 0): BranchPreparation { + const messages: AgentMessage[] = []; + const fileOps = createFileOps(); + let totalTokens = 0; + + // First pass: collect file ops from ALL entries (even if they don't fit in token budget) + // This ensures we capture cumulative file tracking from nested branch summaries + // Only extract from pi-generated summaries (fromHook !== true), not extension-generated ones + 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)) { + // Modified files go into both edited and written for proper deduplication + for (const f of details.modifiedFiles) { + fileOps.edited.add(f); + } + } + } + } + + // Second pass: walk from newest to oldest, adding messages until token budget + for (let i = entries.length - 1; i >= 0; i--) { + const entry = entries[i]; + const message = getMessageFromEntry(entry); + if (!message) continue; + + // Extract file ops from assistant messages (tool calls) + extractFileOpsFromMessage(message, fileOps); + + const tokens = estimateTokens(message); + + // Check budget before adding + if (tokenBudget > 0 && totalTokens + tokens > tokenBudget) { + // If this is a summary entry, try to fit it anyway as it's important context + if (entry.type === "compaction" || entry.type === "branch_summary") { + if (totalTokens < tokenBudget * 0.9) { + messages.unshift(message); + totalTokens += tokens; + } + } + // Stop - we've hit the budget + break; + } + + messages.unshift(message); + totalTokens += tokens; + } + + return { messages, fileOps, totalTokens }; +} + +// ============================================================================ +// Summary Generation +// ============================================================================ + +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.`; + +/** + * Generate a summary of abandoned branch entries. + * + * @param entries - Session entries to summarize (chronological order) + * @param options - Generation options + */ +export async function generateBranchSummary( + entries: SessionTreeEntry[], + options: GenerateBranchSummaryOptions, +): Promise { + const { model, apiKey, headers, signal, customInstructions, replaceInstructions, reserveTokens = 16384 } = options; + + // Token budget = context window minus reserved space for prompt + response + 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" }; + } + + // Transform to LLM-compatible messages, then serialize to text + // Serialization prevents the model from treating it as a conversation to continue + const llmMessages = convertToLlm(messages); + const conversationText = serializeConversation(llmMessages); + + // Build prompt + 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(), + }, + ]; + + // Call LLM for summarization + const response = await completeSimple( + model, + { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages }, + { apiKey, headers, signal, maxTokens: 2048 }, + ); + + // Check if aborted or errored + 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"); + + // Prepend preamble to provide context about the branch summary + summary = BRANCH_SUMMARY_PREAMBLE + summary; + + // Compute file lists and append to summary + const { readFiles, modifiedFiles } = computeFileLists(fileOps); + summary += formatFileOperations(readFiles, modifiedFiles); + + return { + summary: summary || "No summary generated", + readFiles, + modifiedFiles, + }; +} diff --git a/packages/agent/src/vendor/pi-agent-core/harness/compaction/compaction.ts b/packages/agent/src/vendor/pi-agent-core/harness/compaction/compaction.ts new file mode 100644 index 00000000..298dc86d --- /dev/null +++ b/packages/agent/src/vendor/pi-agent-core/harness/compaction/compaction.ts @@ -0,0 +1,854 @@ +/** + * Context compaction for long sessions. + * + * Pure functions for compaction logic. The session manager handles I/O, + * and after compaction the session is reloaded. + */ + +import type { AssistantMessage, ImageContent, Model, TextContent, Usage } from "@earendil-works/pi-ai"; +import { completeSimple } from "@earendil-works/pi-ai"; +import type { AgentMessage, ThinkingLevel } from "../../types.js"; +import { + convertToLlm, + createBranchSummaryMessage, + createCompactionSummaryMessage, + createCustomMessage, +} from "../messages.js"; +import { buildSessionContext } from "../session/session.js"; +import type { CompactionEntry, SessionTreeEntry } from "../types.js"; +import { + computeFileLists, + createFileOps, + extractFileOpsFromMessage, + type FileOperations, + formatFileOperations, + SUMMARIZATION_SYSTEM_PROMPT, + serializeConversation, +} from "./utils.js"; + +// ============================================================================ +// File Operation Tracking +// ============================================================================ + +/** Details stored in CompactionEntry.details for file tracking */ +export interface CompactionDetails { + readFiles: string[]; + modifiedFiles: string[]; +} + +/** + * Extract file operations from messages and previous compaction entries. + */ +function extractFileOperations( + messages: AgentMessage[], + entries: SessionTreeEntry[], + prevCompactionIndex: number, +): FileOperations { + const fileOps = createFileOps(); + + // Collect from previous compaction's details (if pi-generated) + if (prevCompactionIndex >= 0) { + const prevCompaction = entries[prevCompactionIndex] as CompactionEntry; + if (!prevCompaction.fromHook && prevCompaction.details) { + // fromHook field kept for session file compatibility + 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); + } + } + } + + // Extract from tool calls in messages + for (const msg of messages) { + extractFileOpsFromMessage(msg, fileOps); + } + + return fileOps; +} + +// ============================================================================ +// Message Extraction +// ============================================================================ + +/** + * Extract AgentMessage from an entry if it produces one. + * Returns undefined for entries that don't contribute to LLM context. + */ +function getMessageFromEntry(entry: SessionTreeEntry): AgentMessage | undefined { + if (entry.type === "message") { + return entry.message as AgentMessage; + } + if (entry.type === "custom_message") { + return createCustomMessage( + entry.customType, + entry.content as string | (TextContent | ImageContent)[], + entry.display, + entry.details, + entry.timestamp, + ); + } + if (entry.type === "branch_summary") { + return createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp); + } + if (entry.type === "compaction") { + return createCompactionSummaryMessage(entry.summary, entry.tokensBefore, entry.timestamp); + } + return undefined; +} + +function getMessageFromEntryForCompaction(entry: SessionTreeEntry): AgentMessage | undefined { + if (entry.type === "compaction") { + return undefined; + } + return getMessageFromEntry(entry); +} + +/** Result from compact() - SessionManager adds uuid/parentUuid when saving */ +export interface CompactionResult { + summary: string; + firstKeptEntryId: string; + tokensBefore: number; + /** Extension-specific data (e.g., ArtifactIndex, version markers for structured compaction) */ + details?: T; +} + +// ============================================================================ +// Types +// ============================================================================ + +export interface CompactionSettings { + enabled: boolean; + reserveTokens: number; + keepRecentTokens: number; +} + +export const DEFAULT_COMPACTION_SETTINGS: CompactionSettings = { + enabled: true, + reserveTokens: 16384, + keepRecentTokens: 20000, +}; + +// ============================================================================ +// Token calculation +// ============================================================================ + +/** + * Calculate total context tokens from usage. + * Uses the native totalTokens field when available, falls back to computing from components. + */ +export function calculateContextTokens(usage: Usage): number { + return usage.totalTokens || usage.input + usage.output + usage.cacheRead + usage.cacheWrite; +} + +/** + * Get usage from an assistant message if available. + * Skips aborted and error messages as they don't have valid usage data. + */ +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; +} + +/** + * Find the last non-aborted assistant message usage from session entries. + */ +export function getLastAssistantUsage(entries: SessionTreeEntry[]): Usage | undefined { + for (let i = entries.length - 1; i >= 0; i--) { + const entry = entries[i]; + if (entry.type === "message") { + const usage = getAssistantUsage(entry.message as AgentMessage); + if (usage) return usage; + } + } + return undefined; +} + +export interface ContextUsageEstimate { + tokens: number; + usageTokens: number; + trailingTokens: number; + lastUsageIndex: number | null; +} + +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; +} + +/** + * Estimate context tokens from messages, using the last assistant usage when available. + * If there are messages after the last usage, estimate their tokens with estimateTokens. + */ +export function estimateContextTokens(messages: AgentMessage[]): ContextUsageEstimate { + 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, + }; +} + +/** + * Check if compaction should trigger based on context usage. + */ +export function shouldCompact(contextTokens: number, contextWindow: number, settings: CompactionSettings): boolean { + if (!settings.enabled) return false; + return contextTokens > contextWindow - settings.reserveTokens; +} + +// ============================================================================ +// Cut point detection +// ============================================================================ + +/** + * Estimate token count for a message using chars/4 heuristic. + * This is conservative (overestimates tokens). + */ +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; // Estimate images as 4000 chars, or 1200 tokens + } + } + } + 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; +} + +/** + * Find valid cut points: indices of user, assistant, custom, or bashExecution messages. + * Never cut at tool results (they must follow their tool call). + * When we cut at an assistant message with tool calls, its tool results follow it + * and will be kept. + * BashExecutionMessage is treated like a user message (user-initiated context). + */ +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": { + const role = entry.message.role; + switch (role) { + case "bashExecution": + case "custom": + case "branchSummary": + case "compactionSummary": + case "user": + case "assistant": + cutPoints.push(i); + break; + case "toolResult": + break; + } + break; + } + case "thinking_level_change": + case "model_change": + case "compaction": + case "branch_summary": + case "custom": + case "custom_message": + case "label": + case "session_info": + break; + } + + // branch_summary and custom_message are user-role messages, valid cut points + if (entry.type === "branch_summary" || entry.type === "custom_message") { + cutPoints.push(i); + } + } + return cutPoints; +} + +/** + * Find the user message (or bashExecution) that starts the turn containing the given entry index. + * Returns -1 if no turn start found before the index. + * BashExecutionMessage is treated like a user message for turn boundaries. + */ +export function findTurnStartIndex(entries: SessionTreeEntry[], entryIndex: number, startIndex: number): number { + for (let i = entryIndex; i >= startIndex; i--) { + const entry = entries[i]; + // branch_summary and custom_message are user-role messages, can start a turn + 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 interface CutPointResult { + /** Index of first entry to keep */ + firstKeptEntryIndex: number; + /** Index of user message that starts the turn being split, or -1 if not splitting */ + turnStartIndex: number; + /** Whether this cut splits a turn (cut point is not a user message) */ + isSplitTurn: boolean; +} + +/** + * Find the cut point in session entries that keeps approximately `keepRecentTokens`. + * + * Algorithm: Walk backwards from newest, accumulating estimated message sizes. + * Stop when we've accumulated >= keepRecentTokens. Cut at that point. + * + * Can cut at user OR assistant messages (never tool results). When cutting at an + * assistant message with tool calls, its tool results come after and will be kept. + * + * Returns CutPointResult with: + * - firstKeptEntryIndex: the entry index to start keeping from + * - turnStartIndex: if cutting mid-turn, the user message that started that turn + * - isSplitTurn: whether we're cutting in the middle of a turn + * + * Only considers entries between `startIndex` and `endIndex` (exclusive). + */ +export function findCutPoint( + entries: SessionTreeEntry[], + startIndex: number, + endIndex: number, + keepRecentTokens: number, +): CutPointResult { + const cutPoints = findValidCutPoints(entries, startIndex, endIndex); + + if (cutPoints.length === 0) { + return { firstKeptEntryIndex: startIndex, turnStartIndex: -1, isSplitTurn: false }; + } + + // Walk backwards from newest, accumulating estimated message sizes + let accumulatedTokens = 0; + let cutIndex = cutPoints[0]; // Default: keep from first message (not header) + + for (let i = endIndex - 1; i >= startIndex; i--) { + const entry = entries[i]; + if (entry.type !== "message") continue; + + // Estimate this message's size + const messageTokens = estimateTokens(entry.message as AgentMessage); + accumulatedTokens += messageTokens; + + // Check if we've exceeded the budget + if (accumulatedTokens >= keepRecentTokens) { + // Find the closest valid cut point at or after this entry + for (let c = 0; c < cutPoints.length; c++) { + if (cutPoints[c] >= i) { + cutIndex = cutPoints[c]; + break; + } + } + break; + } + } + + // Scan backwards from cutIndex to include any non-message entries (bash, settings, etc.) + while (cutIndex > startIndex) { + const prevEntry = entries[cutIndex - 1]; + // Stop at session header or compaction boundaries + if (prevEntry.type === "compaction") { + break; + } + if (prevEntry.type === "message") { + // Stop if we hit any message + break; + } + // Include this non-message entry (bash, settings change, etc.) + cutIndex--; + } + + // Determine if this is a split turn + 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, + }; +} + +// ============================================================================ +// Summarization +// ============================================================================ + +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.`; + +/** + * Generate a summary of the conversation using the LLM. + * If previousSummary is provided, uses the update prompt to merge. + */ +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.min( + Math.floor(0.8 * reserveTokens), + model.maxTokens > 0 ? model.maxTokens : Number.POSITIVE_INFINITY, + ); + + // Use update prompt if we have a previous summary, otherwise initial prompt + let basePrompt = previousSummary ? UPDATE_SUMMARIZATION_PROMPT : SUMMARIZATION_PROMPT; + if (customInstructions) { + basePrompt = `${basePrompt}\n\nAdditional focus: ${customInstructions}`; + } + + // Serialize conversation to text so model doesn't try to continue it + // Convert to LLM messages first (handles custom types like bashExecution, custom, etc.) + const llmMessages = convertToLlm(currentMessages); + const conversationText = serializeConversation(llmMessages); + + // Build the prompt with conversation wrapped in tags + 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"}`); + } + + const textContent = response.content + .filter((c): c is { type: "text"; text: string } => c.type === "text") + .map((c) => c.text) + .join("\n"); + + return textContent; +} + +// ============================================================================ +// Compaction Preparation (for extensions) +// ============================================================================ + +export interface CompactionPreparation { + /** UUID of first entry to keep */ + firstKeptEntryId: string; + /** Messages that will be summarized and discarded */ + messagesToSummarize: AgentMessage[]; + /** Messages that will be turned into turn prefix summary (if splitting) */ + turnPrefixMessages: AgentMessage[]; + /** Whether this is a split turn (cut point in middle of turn) */ + isSplitTurn: boolean; + tokensBefore: number; + /** Summary from previous compaction, for iterative update */ + previousSummary?: string; + /** File operations extracted from messagesToSummarize */ + fileOps: FileOperations; + /** Compaction settions from settings.jsonl */ + settings: CompactionSettings; +} + +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); + + // Get UUID of first kept entry + const firstKeptEntry = pathEntries[cutPoint.firstKeptEntryIndex]; + if (!firstKeptEntry?.id) { + return undefined; // Session needs migration + } + const firstKeptEntryId = firstKeptEntry.id; + + const historyEnd = cutPoint.isSplitTurn ? cutPoint.turnStartIndex : cutPoint.firstKeptEntryIndex; + + // Messages to summarize (will be discarded after summary) + const messagesToSummarize: AgentMessage[] = []; + for (let i = boundaryStart; i < historyEnd; i++) { + const msg = getMessageFromEntryForCompaction(pathEntries[i]); + if (msg) messagesToSummarize.push(msg); + } + + // Messages for turn prefix summary (if splitting a turn) + 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); + } + } + + // Extract file operations from messages and previous compaction + const fileOps = extractFileOperations(messagesToSummarize, pathEntries, prevCompactionIndex); + + // Also extract file ops from turn prefix if splitting + if (cutPoint.isSplitTurn) { + for (const msg of turnPrefixMessages) { + extractFileOpsFromMessage(msg, fileOps); + } + } + + return { + firstKeptEntryId, + messagesToSummarize, + turnPrefixMessages, + isSplitTurn: cutPoint.isSplitTurn, + tokensBefore, + previousSummary, + fileOps, + settings, + }; +} + +// ============================================================================ +// Main compaction function +// ============================================================================ + +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.`; + +/** + * Generate summaries for compaction using prepared data. + * Returns CompactionResult - SessionManager adds uuid/parentUuid when saving. + * + * @param preparation - Pre-calculated preparation from prepareCompaction() + * @param customInstructions - Optional custom focus for the summary + */ +export { serializeConversation } from "./utils.js"; + +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; + + // Generate summaries (can be parallel if both needed) and merge into one + let summary: string; + + if (isSplitTurn && turnPrefixMessages.length > 0) { + // Generate both summaries in parallel + 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, + ), + ]); + // Merge into single summary + summary = `${historyResult}\n\n---\n\n**Turn Context (split turn):**\n\n${turnPrefixResult}`; + } else { + // Just generate history summary + summary = await generateSummary( + messagesToSummarize, + model, + settings.reserveTokens, + apiKey, + headers, + signal, + customInstructions, + previousSummary, + thinkingLevel, + ); + } + + // Compute file lists and append to summary + const { readFiles, modifiedFiles } = computeFileLists(fileOps); + summary += formatFileOperations(readFiles, modifiedFiles); + + if (!firstKeptEntryId) { + throw new Error("First kept entry has no UUID - session may need migration"); + } + + return { + summary, + firstKeptEntryId, + tokensBefore, + details: { readFiles, modifiedFiles } as CompactionDetails, + }; +} + +/** + * Generate a summary for a turn prefix (when splitting a turn). + */ +async function generateTurnPrefixSummary( + messages: AgentMessage[], + model: Model, + reserveTokens: number, + apiKey: string, + headers?: Record, + signal?: AbortSignal, + thinkingLevel?: ThinkingLevel, +): Promise { + const maxTokens = Math.min( + Math.floor(0.5 * reserveTokens), + model.maxTokens > 0 ? model.maxTokens : Number.POSITIVE_INFINITY, + ); // Smaller budget for turn prefix + 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 response = await completeSimple( + model, + { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages }, + model.reasoning && thinkingLevel && thinkingLevel !== "off" + ? { maxTokens, signal, apiKey, headers, reasoning: thinkingLevel } + : { maxTokens, signal, apiKey, headers }, + ); + + 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"); +} diff --git a/packages/agent/src/vendor/pi-agent-core/harness/compaction/utils.ts b/packages/agent/src/vendor/pi-agent-core/harness/compaction/utils.ts new file mode 100644 index 00000000..35cd1938 --- /dev/null +++ b/packages/agent/src/vendor/pi-agent-core/harness/compaction/utils.ts @@ -0,0 +1,170 @@ +/** + * Shared utilities for compaction and branch summarization. + */ + +import type { Message } from "@earendil-works/pi-ai"; +import type { AgentMessage } from "../../types.js"; + +// ============================================================================ +// File Operation Tracking +// ============================================================================ + +export interface FileOperations { + read: Set; + written: Set; + edited: Set; +} + +export function createFileOps(): FileOperations { + return { + read: new Set(), + written: new Set(), + edited: new Set(), + }; +} + +/** + * Extract file operations from tool calls in an assistant message. + */ +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) continue; + if (!("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; + } + } +} + +/** + * Compute final file lists from file operations. + * Returns readFiles (files only read, not modified) and modifiedFiles. + */ +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 }; +} + +/** + * Format file operations as XML tags for summary. + */ +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")}`; +} + +// ============================================================================ +// Message Serialization +// ============================================================================ + +/** Maximum characters for a tool result in serialized summaries. */ +const TOOL_RESULT_MAX_CHARS = 2000; + +/** + * Truncate text to a maximum character length for summarization. + * Keeps the beginning and appends a truncation marker. + */ +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]`; +} + +/** + * Serialize LLM messages to text for summarization. + * This prevents the model from treating it as a conversation to continue. + * Call convertToLlm() first to handle custom message types. + * + * Tool results are truncated to keep the summarization request within + * reasonable token budgets. Full content is not needed for summarization. + */ +export function serializeConversation(messages: Message[]): 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 args = block.arguments as Record; + const argsStr = Object.entries(args) + .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, TOOL_RESULT_MAX_CHARS)}`); + } + } + } + + return parts.join("\n\n"); +} + +// ============================================================================ +// Summarization System Prompt +// ============================================================================ + +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.`; diff --git a/packages/agent/src/vendor/pi-agent-core/harness/env/nodejs.ts b/packages/agent/src/vendor/pi-agent-core/harness/env/nodejs.ts new file mode 100644 index 00000000..d4b6c0a9 --- /dev/null +++ b/packages/agent/src/vendor/pi-agent-core/harness/env/nodejs.ts @@ -0,0 +1,370 @@ +import { spawn } from "node:child_process"; +import { randomUUID } from "node:crypto"; +import { constants } from "node:fs"; +import { access, lstat, mkdir, mkdtemp, readdir, readFile, realpath, rm, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { isAbsolute, join, resolve } from "node:path"; +import { type ExecutionEnv, FileError, type FileInfo, type FileKind } from "../types.js"; + +function resolvePath(cwd: string, path: string): string { + return isAbsolute(path) ? path : resolve(cwd, path); +} + +function fileKindFromStats(stats: { isFile(): boolean; isDirectory(): boolean; isSymbolicLink(): boolean }): FileKind { + if (stats.isFile()) return "file"; + if (stats.isDirectory()) return "directory"; + if (stats.isSymbolicLink()) return "symlink"; + throw new FileError("invalid", "Unsupported file type"); +} + +function fileInfoFromStats( + path: string, + stats: { isFile(): boolean; isDirectory(): boolean; isSymbolicLink(): boolean; size: number; mtimeMs: number }, +): FileInfo { + return { + name: path.replace(/\/+$/, "").split("/").pop() ?? path, + path, + kind: fileKindFromStats(stats), + size: stats.size, + mtimeMs: stats.mtimeMs, + }; +} + +function isNodeError(error: unknown): error is NodeJS.ErrnoException { + return error instanceof Error && "code" in error; +} + +function toFileError(error: unknown, path?: string): FileError { + if (error instanceof FileError) return error; + if (isNodeError(error)) { + const message = error.message; + switch (error.code) { + case "ENOENT": + return new FileError("not_found", message, path, { cause: error }); + case "EACCES": + case "EPERM": + return new FileError("permission_denied", message, path, { cause: error }); + case "ENOTDIR": + return new FileError("not_directory", message, path, { cause: error }); + case "EISDIR": + return new FileError("is_directory", message, path, { cause: error }); + case "EINVAL": + return new FileError("invalid", message, path, { cause: error }); + } + } + return new FileError("unknown", error instanceof Error ? error.message : String(error), path, { cause: error }); +} + +async function pathExists(path: string): Promise { + try { + await access(path, constants.F_OK); + return true; + } catch { + return false; + } +} + +async function runCommand( + command: string, + args: string[], + timeoutMs: number, +): Promise<{ stdout: string; status: number | null }> { + return await new Promise((resolve) => { + let stdout = ""; + const child = spawn(command, args, { stdio: ["ignore", "pipe", "ignore"] }); + const timeout = setTimeout(() => { + if (child.pid) killProcessTree(child.pid); + }, timeoutMs); + child.stdout?.setEncoding("utf8"); + child.stdout?.on("data", (chunk: string) => { + stdout += chunk; + }); + child.on("error", () => { + clearTimeout(timeout); + resolve({ stdout: "", status: null }); + }); + child.on("close", (status) => { + clearTimeout(timeout); + resolve({ stdout, status }); + }); + }); +} + +async function findBashOnPath(): Promise { + const result = + process.platform === "win32" + ? await runCommand("where", ["bash.exe"], 5000) + : await runCommand("which", ["bash"], 5000); + if (result.status !== 0 || !result.stdout) return null; + const firstMatch = result.stdout.trim().split(/\r?\n/)[0]; + return firstMatch && (await pathExists(firstMatch)) ? firstMatch : null; +} + +async function getShellConfig(customShellPath?: string): Promise<{ shell: string; args: string[] }> { + if (customShellPath) { + if (await pathExists(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 (await pathExists(candidate)) { + return { shell: candidate, args: ["-c"] }; + } + } + const bashOnPath = await findBashOnPath(); + if (bashOnPath) { + return { shell: bashOnPath, args: ["-c"] }; + } + throw new Error("No bash shell found"); + } + + if (await pathExists("/bin/bash")) { + return { shell: "/bin/bash", args: ["-c"] }; + } + const bashOnPath = await 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 } = await 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 { + const resolved = resolvePath(this.cwd, path); + try { + return await readFile(resolved, "utf8"); + } catch (error) { + throw toFileError(error, resolved); + } + } + + async readBinaryFile(path: string): Promise { + const resolved = resolvePath(this.cwd, path); + try { + return await readFile(resolved); + } catch (error) { + throw toFileError(error, resolved); + } + } + + async writeFile(path: string, content: string | Uint8Array): Promise { + const resolved = resolvePath(this.cwd, path); + try { + await mkdir(resolve(resolved, ".."), { recursive: true }); + await writeFile(resolved, content); + } catch (error) { + throw toFileError(error, resolved); + } + } + + async fileInfo(path: string): Promise { + const resolved = resolvePath(this.cwd, path); + try { + return fileInfoFromStats(resolved, await lstat(resolved)); + } catch (error) { + throw toFileError(error, resolved); + } + } + + async listDir(path: string): Promise { + const resolved = resolvePath(this.cwd, path); + try { + const entries = await readdir(resolved, { withFileTypes: true }); + const infos: FileInfo[] = []; + for (const entry of entries) { + const entryPath = resolve(resolved, entry.name); + try { + infos.push(fileInfoFromStats(entryPath, await lstat(entryPath))); + } catch (error) { + if (error instanceof FileError && error.code === "invalid") continue; + throw error; + } + } + return infos; + } catch (error) { + throw toFileError(error, resolved); + } + } + + async realPath(path: string): Promise { + const resolved = resolvePath(this.cwd, path); + try { + return await realpath(resolved); + } catch (error) { + throw toFileError(error, resolved); + } + } + + async exists(path: string): Promise { + try { + await this.fileInfo(path); + return true; + } catch (error) { + if (error instanceof FileError && error.code === "not_found") return false; + throw error; + } + } + + 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 { + const resolved = resolvePath(this.cwd, path); + try { + await rm(resolved, { recursive: options?.recursive ?? false, force: options?.force ?? false }); + } catch (error) { + throw toFileError(error, resolved); + } + } + + 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; + } + + async cleanup(): Promise { + // nothing to clean up for the local node implementation + } +} diff --git a/packages/agent/src/vendor/pi-agent-core/harness/execution-env.ts b/packages/agent/src/vendor/pi-agent-core/harness/execution-env.ts new file mode 100644 index 00000000..786586e5 --- /dev/null +++ b/packages/agent/src/vendor/pi-agent-core/harness/execution-env.ts @@ -0,0 +1,3 @@ +export { NodeExecutionEnv } from "./env/nodejs.js"; +export type { ExecutionEnv, ExecutionEnvExecOptions, FileErrorCode, FileInfo, FileKind } from "./types.js"; +export { FileError } from "./types.js"; diff --git a/packages/agent/src/vendor/pi-agent-core/harness/messages.ts b/packages/agent/src/vendor/pi-agent-core/harness/messages.ts new file mode 100644 index 00000000..615bf1e4 --- /dev/null +++ b/packages/agent/src/vendor/pi-agent-core/harness/messages.ts @@ -0,0 +1,164 @@ +import type { ImageContent, Message, TextContent } from "@earendil-works/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/vendor/pi-agent-core/harness/prompt-templates.ts b/packages/agent/src/vendor/pi-agent-core/harness/prompt-templates.ts new file mode 100644 index 00000000..e77682b8 --- /dev/null +++ b/packages/agent/src/vendor/pi-agent-core/harness/prompt-templates.ts @@ -0,0 +1,224 @@ +import { parse } from "yaml"; +import type { ExecutionEnv, FileInfo, PromptTemplate } from "./types.js"; + +/** Warning produced while loading prompt templates. */ +export interface PromptTemplateDiagnostic { + /** Diagnostic severity. Currently only warnings are emitted. */ + type: "warning"; + /** Human-readable diagnostic message. */ + message: string; + /** Path associated with the diagnostic. */ + path: string; +} + +interface PromptTemplateFrontmatter { + description?: string; + "argument-hint"?: string; + [key: string]: unknown; +} + +/** + * Load prompt templates from one or more paths. + * + * Directory inputs load direct `.md` children non-recursively. File inputs load explicit `.md` files. Missing paths and + * non-markdown files are skipped. Read and parse failures are returned as diagnostics. + */ +export async function loadPromptTemplates( + env: ExecutionEnv, + paths: string | string[], +): Promise<{ promptTemplates: PromptTemplate[]; diagnostics: PromptTemplateDiagnostic[] }> { + const promptTemplates: PromptTemplate[] = []; + const diagnostics: PromptTemplateDiagnostic[] = []; + for (const path of Array.isArray(paths) ? paths : [paths]) { + const info = await safeFileInfo(env, path); + if (!info) continue; + const kind = await resolveKind(env, info); + if (kind === "directory") { + const result = await loadTemplatesFromDir(env, info.path); + promptTemplates.push(...result.promptTemplates); + diagnostics.push(...result.diagnostics); + } else if (kind === "file" && info.name.endsWith(".md")) { + const result = await loadTemplateFromFile(env, info.path); + if (result.promptTemplate) promptTemplates.push(result.promptTemplate); + diagnostics.push(...result.diagnostics); + } + } + return { promptTemplates, diagnostics }; +} + +/** + * Load prompt templates from source-tagged paths. + * + * Source values are preserved exactly and attached to every loaded prompt template and diagnostic. The agent package does + * not interpret source values; applications define their own provenance shape. + */ +export async function loadSourcedPromptTemplates( + env: ExecutionEnv, + inputs: Array<{ path: string; source: TSource }>, + mapPromptTemplate?: (promptTemplate: PromptTemplate, source: TSource) => TPromptTemplate, +): Promise<{ + promptTemplates: Array<{ promptTemplate: TPromptTemplate; source: TSource }>; + diagnostics: Array; +}> { + const promptTemplates: Array<{ promptTemplate: TPromptTemplate; source: TSource }> = []; + const diagnostics: Array = []; + for (const input of inputs) { + const result = await loadPromptTemplates(env, input.path); + for (const promptTemplate of result.promptTemplates) { + promptTemplates.push({ + promptTemplate: mapPromptTemplate + ? mapPromptTemplate(promptTemplate, input.source) + : (promptTemplate as TPromptTemplate), + source: input.source, + }); + } + for (const diagnostic of result.diagnostics) diagnostics.push({ ...diagnostic, source: input.source }); + } + return { promptTemplates, diagnostics }; +} + +async function loadTemplatesFromDir( + env: ExecutionEnv, + dir: string, +): Promise<{ promptTemplates: PromptTemplate[]; diagnostics: PromptTemplateDiagnostic[] }> { + const promptTemplates: PromptTemplate[] = []; + const diagnostics: PromptTemplateDiagnostic[] = []; + let entries: FileInfo[]; + try { + entries = await env.listDir(dir); + } catch (error) { + diagnostics.push({ + type: "warning", + message: errorMessage(error, "failed to list prompt template directory"), + path: dir, + }); + return { promptTemplates, diagnostics }; + } + + for (const entry of entries.sort((a, b) => a.name.localeCompare(b.name))) { + const kind = await resolveKind(env, entry); + if (kind !== "file" || !entry.name.endsWith(".md")) continue; + const result = await loadTemplateFromFile(env, entry.path); + if (result.promptTemplate) promptTemplates.push(result.promptTemplate); + diagnostics.push(...result.diagnostics); + } + return { promptTemplates, diagnostics }; +} + +async function loadTemplateFromFile( + env: ExecutionEnv, + filePath: string, +): Promise<{ promptTemplate: PromptTemplate | null; diagnostics: PromptTemplateDiagnostic[] }> { + const diagnostics: PromptTemplateDiagnostic[] = []; + try { + const rawContent = await env.readTextFile(filePath); + const { frontmatter, body } = parseFrontmatter(rawContent); + const firstLine = body.split("\n").find((line) => line.trim()); + let description = typeof frontmatter.description === "string" ? frontmatter.description : ""; + if (!description && firstLine) { + description = firstLine.slice(0, 60); + if (firstLine.length > 60) description += "..."; + } + return { + promptTemplate: { + name: basenameEnvPath(filePath).replace(/\.md$/i, ""), + description, + content: body, + }, + diagnostics, + }; + } catch (error) { + diagnostics.push({ + type: "warning", + message: errorMessage(error, "failed to load prompt template"), + path: filePath, + }); + return { promptTemplate: null, diagnostics }; + } +} + +async function safeFileInfo(env: ExecutionEnv, path: string): Promise { + try { + return await env.fileInfo(path); + } catch { + return undefined; + } +} + +async function resolveKind(env: ExecutionEnv, info: FileInfo): Promise<"file" | "directory" | undefined> { + if (info.kind === "file" || info.kind === "directory") return info.kind; + try { + const realPath = await env.realPath(info.path); + const target = await env.fileInfo(realPath); + return target.kind === "file" || target.kind === "directory" ? target.kind : undefined; + } catch { + return undefined; + } +} + +function parseFrontmatter>(content: string): { frontmatter: T; body: string } { + const normalized = content.replace(/\r\n/g, "\n").replace(/\r/g, "\n"); + if (!normalized.startsWith("---")) return { frontmatter: {} as T, body: normalized }; + const endIndex = normalized.indexOf("\n---", 3); + if (endIndex === -1) return { frontmatter: {} as T, body: normalized }; + const yamlString = normalized.slice(4, endIndex); + const body = normalized.slice(endIndex + 4).trim(); + return { frontmatter: (parse(yamlString) ?? {}) as T, body }; +} + +function basenameEnvPath(path: string): string { + const normalized = path.replace(/\/+$/, ""); + const slashIndex = normalized.lastIndexOf("/"); + return slashIndex === -1 ? normalized : normalized.slice(slashIndex + 1); +} + +function errorMessage(error: unknown, fallback: string): string { + return error instanceof Error ? error.message : fallback; +} + +/** Parse an argument string using simple shell-style single and double quotes. */ +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; +} + +/** Substitute prompt template placeholders (`$1`, `$@`, `$ARGUMENTS`, `${@:N}`, `${@:N:L}`) with command arguments. */ +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; +} + +/** Format a prompt template invocation with positional arguments. */ +export function formatPromptTemplateInvocation(template: PromptTemplate, args: string[] = []): string { + return substituteArgs(template.content, args); +} diff --git a/packages/agent/src/vendor/pi-agent-core/harness/session/repo/jsonl.ts b/packages/agent/src/vendor/pi-agent-core/harness/session/repo/jsonl.ts new file mode 100644 index 00000000..a4c5ec5b --- /dev/null +++ b/packages/agent/src/vendor/pi-agent-core/harness/session/repo/jsonl.ts @@ -0,0 +1,109 @@ +import { constants } from "node:fs"; +import { access, mkdir, readdir, rm } from "node:fs/promises"; +import { join, resolve } from "node:path"; +import type { + JsonlSessionCreateOptions, + JsonlSessionListOptions, + JsonlSessionMetadata, + JsonlSessionRepoApi, + Session, +} from "../../types.js"; +import { JsonlSessionStorage, loadJsonlSessionMetadata } from "../storage/jsonl.js"; +import { createSessionId, createTimestamp, getEntriesToFork, toSession } from "./shared.js"; + +async function exists(path: string): Promise { + try { + await access(path, constants.F_OK); + return true; + } catch { + return false; + } +} + +function encodeCwd(cwd: string): string { + return `--${cwd.replace(/^[/\\]/, "").replace(/[/\\:]/g, "-")}--`; +} + +export class JsonlSessionRepo implements JsonlSessionRepoApi { + private sessionsRoot: string; + + constructor(options: { sessionsRoot: string }) { + this.sessionsRoot = resolve(options.sessionsRoot); + } + + private getSessionDir(cwd: string): string { + return join(this.sessionsRoot, encodeCwd(cwd)); + } + + private createSessionFilePath(cwd: string, sessionId: string, timestamp: string): string { + return join(this.getSessionDir(cwd), `${timestamp.replace(/[:.]/g, "-")}_${sessionId}.jsonl`); + } + + async create(options: JsonlSessionCreateOptions): Promise> { + await mkdir(this.sessionsRoot, { recursive: true }); + const id = options.id ?? createSessionId(); + const createdAt = createTimestamp(); + const filePath = this.createSessionFilePath(options.cwd, id, createdAt); + const storage = await JsonlSessionStorage.create(filePath, { + cwd: options.cwd, + sessionId: id, + parentSessionPath: options.parentSessionPath, + }); + return toSession(storage); + } + + async open(metadata: JsonlSessionMetadata): Promise> { + if (!(await exists(metadata.path))) { + throw new Error(`Session not found: ${metadata.path}`); + } + const storage = await JsonlSessionStorage.open(metadata.path); + return toSession(storage); + } + + async list(options: JsonlSessionListOptions = {}): Promise { + const dirs = options.cwd ? [this.getSessionDir(options.cwd)] : await this.listSessionDirs(); + const sessions: JsonlSessionMetadata[] = []; + for (const dir of dirs) { + if (!(await exists(dir))) continue; + const files = (await readdir(dir)).filter((file) => file.endsWith(".jsonl")).map((file) => join(dir, file)); + for (const filePath of files) { + try { + sessions.push(await loadJsonlSessionMetadata(filePath)); + } catch { + // Ignore invalid session files when listing a directory. + } + } + } + sessions.sort((a, b) => new Date(b.createdAt).getTime() - new Date(a.createdAt).getTime()); + return sessions; + } + + async delete(metadata: JsonlSessionMetadata): Promise { + await rm(metadata.path, { force: true }); + } + + async fork( + sourceMetadata: JsonlSessionMetadata, + options: JsonlSessionCreateOptions & { entryId?: string; position?: "before" | "at"; id?: string }, + ): Promise> { + const source = await this.open(sourceMetadata); + const forkedEntries = await getEntriesToFork(source.getStorage(), options); + const id = options.id ?? createSessionId(); + const createdAt = createTimestamp(); + const storage = await JsonlSessionStorage.create(this.createSessionFilePath(options.cwd, id, createdAt), { + cwd: options.cwd, + sessionId: id, + parentSessionPath: options.parentSessionPath ?? sourceMetadata.path, + }); + for (const entry of forkedEntries) { + await storage.appendEntry(entry); + } + return toSession(storage); + } + + private async listSessionDirs(): Promise { + if (!(await exists(this.sessionsRoot))) return []; + const entries = await readdir(this.sessionsRoot, { withFileTypes: true }); + return entries.filter((entry) => entry.isDirectory()).map((entry) => join(this.sessionsRoot, entry.name)); + } +} diff --git a/packages/agent/src/vendor/pi-agent-core/harness/session/repo/memory.ts b/packages/agent/src/vendor/pi-agent-core/harness/session/repo/memory.ts new file mode 100644 index 00000000..3846ae2a --- /dev/null +++ b/packages/agent/src/vendor/pi-agent-core/harness/session/repo/memory.ts @@ -0,0 +1,51 @@ +import type { Session, SessionMetadata, SessionRepo } from "../../types.js"; +import { InMemorySessionStorage } from "../storage/memory.js"; +import { createSessionId, createTimestamp, getEntriesToFork, toSession } from "./shared.js"; + +export class InMemorySessionRepo implements SessionRepo { + private sessions = new Map>(); + + async create(options: { id?: string } = {}): Promise> { + const metadata: SessionMetadata = { + id: options.id ?? createSessionId(), + createdAt: createTimestamp(), + }; + const storage = new InMemorySessionStorage({ metadata }); + const session = toSession(storage); + this.sessions.set(metadata.id, session); + return session; + } + + async open(metadata: SessionMetadata): Promise> { + const session = this.sessions.get(metadata.id); + if (!session) { + throw new Error(`Session not found: ${metadata.id}`); + } + return session; + } + + async list(): Promise { + return Promise.all([...this.sessions.values()].map((session) => session.getMetadata())); + } + + async delete(metadata: SessionMetadata): Promise { + this.sessions.delete(metadata.id); + } + + async fork( + sourceMetadata: SessionMetadata, + options: { entryId?: string; position?: "before" | "at"; id?: string }, + ): Promise> { + const source = await this.open(sourceMetadata); + const forkedEntries = await getEntriesToFork(source.getStorage(), options); + const metadata: SessionMetadata = { + id: options.id ?? createSessionId(), + createdAt: createTimestamp(), + }; + const leafId = forkedEntries[forkedEntries.length - 1]?.id ?? null; + const storage = new InMemorySessionStorage({ metadata, entries: forkedEntries, leafId }); + const session = toSession(storage); + this.sessions.set(metadata.id, session); + return session; + } +} diff --git a/packages/agent/src/vendor/pi-agent-core/harness/session/repo/shared.ts b/packages/agent/src/vendor/pi-agent-core/harness/session/repo/shared.ts new file mode 100644 index 00000000..2c628ca9 --- /dev/null +++ b/packages/agent/src/vendor/pi-agent-core/harness/session/repo/shared.ts @@ -0,0 +1,36 @@ +import type { SessionMetadata, SessionStorage, SessionTreeEntry } from "../../types.js"; +import { Session } from "../session.js"; +import { uuidv7 } from "../uuid.js"; + +export function createSessionId(): string { + return uuidv7(); +} + +export function createTimestamp(): string { + return new Date().toISOString(); +} + +export function toSession(storage: SessionStorage): Session { + return new Session(storage); +} + +export async function getEntriesToFork( + storage: SessionStorage, + options: { entryId?: string; position?: "before" | "at" }, +): Promise { + if (!options.entryId) return storage.getEntries(); + const target = await storage.getEntry(options.entryId); + if (!target) { + throw new Error(`Entry ${options.entryId} not found`); + } + let effectiveLeafId: string | null; + if ((options.position ?? "before") === "at") { + effectiveLeafId = target.id; + } else { + if (target.type !== "message" || target.message.role !== "user") { + throw new Error(`Entry ${options.entryId} is not a user message`); + } + effectiveLeafId = target.parentId; + } + return storage.getPathToRoot(effectiveLeafId); +} diff --git a/packages/agent/src/vendor/pi-agent-core/harness/session/session.ts b/packages/agent/src/vendor/pi-agent-core/harness/session/session.ts new file mode 100644 index 00000000..f0b42ecf --- /dev/null +++ b/packages/agent/src/vendor/pi-agent-core/harness/session/session.ts @@ -0,0 +1,251 @@ +import type { ImageContent, TextContent } from "@earendil-works/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, + SessionMetadata, + SessionStorage, + SessionTreeEntry, + ThinkingLevelChangeEntry, +} from "../types.js"; + +export function buildSessionContext(pathEntries: SessionTreeEntry[]): SessionContext { + let thinkingLevel = "off"; + let model: { provider: string; modelId: string } | null = null; + let compaction: CompactionEntry | null = null; + + for (const entry of pathEntries) { + 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 as AgentMessage); + } else if (entry.type === "custom_message") { + messages.push( + createCustomMessage( + entry.customType, + entry.content as string | (TextContent | ImageContent)[], + 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 = pathEntries.findIndex((e) => e.type === "compaction" && e.id === compaction.id); + let foundFirstKept = false; + for (let i = 0; i < compactionIdx; i++) { + const entry = pathEntries[i]!; + if (entry.id === compaction.firstKeptEntryId) foundFirstKept = true; + if (foundFirstKept) appendMessage(entry); + } + for (let i = compactionIdx + 1; i < pathEntries.length; i++) { + appendMessage(pathEntries[i]!); + } + } else { + for (const entry of pathEntries) { + appendMessage(entry); + } + } + + return { messages, thinkingLevel, model }; +} + +export class Session { + private storage: SessionStorage; + + constructor(storage: SessionStorage) { + this.storage = storage; + } + + getMetadata(): Promise { + return this.storage.getMetadata(); + } + + getStorage(): SessionStorage { + return this.storage; + } + + 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()); + } + + getLabel(id: string): Promise { + return this.storage.getLabel(id); + } + + async getSessionName(): Promise { + const entries = await this.storage.findEntries("session_info"); + return entries[entries.length - 1]?.name?.trim() || undefined; + } + + 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.storage.createEntryId(), + 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.storage.createEntryId(), + 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.storage.createEntryId(), + 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.storage.createEntryId(), + parentId: await this.storage.getLeafId(), + timestamp: new Date().toISOString(), + summary, + firstKeptEntryId, + tokensBefore, + details, + fromHook, + } satisfies CompactionEntry); + } + + async appendCustomEntry(customType: string, data?: unknown): Promise { + return this.appendTypedEntry({ + type: "custom", + id: await this.storage.createEntryId(), + 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.storage.createEntryId(), + parentId: await this.storage.getLeafId(), + timestamp: new Date().toISOString(), + customType, + content, + display, + details, + } satisfies CustomMessageEntry); + } + + async appendLabel(targetId: string, label: string | undefined): Promise { + if (!(await this.storage.getEntry(targetId))) { + throw new Error(`Entry ${targetId} not found`); + } + return this.appendTypedEntry({ + type: "label", + id: await this.storage.createEntryId(), + parentId: await this.storage.getLeafId(), + timestamp: new Date().toISOString(), + targetId, + label, + } satisfies LabelEntry); + } + + async appendSessionName(name: string): Promise { + return this.appendTypedEntry({ + type: "session_info", + id: await this.storage.createEntryId(), + parentId: await this.storage.getLeafId(), + timestamp: new Date().toISOString(), + name: name.trim(), + } satisfies SessionInfoEntry); + } + + async moveTo( + entryId: string | null, + summary?: { summary: string; details?: unknown; fromHook?: boolean }, + ): Promise { + if (entryId !== null && !(await this.storage.getEntry(entryId))) { + throw new Error(`Entry ${entryId} not found`); + } + await this.storage.setLeafId(entryId); + if (!summary) return undefined; + return this.appendTypedEntry({ + type: "branch_summary", + id: await this.storage.createEntryId(), + parentId: entryId, + timestamp: new Date().toISOString(), + fromId: entryId ?? "root", + summary: summary.summary, + details: summary.details, + fromHook: summary.fromHook, + } satisfies BranchSummaryEntry); + } +} diff --git a/packages/agent/src/vendor/pi-agent-core/harness/session/storage/jsonl.ts b/packages/agent/src/vendor/pi-agent-core/harness/session/storage/jsonl.ts new file mode 100644 index 00000000..72ce6ce3 --- /dev/null +++ b/packages/agent/src/vendor/pi-agent-core/harness/session/storage/jsonl.ts @@ -0,0 +1,205 @@ +import { randomUUID } from "node:crypto"; +import { createReadStream } from "node:fs"; +import { appendFile, mkdir, readFile, writeFile } from "node:fs/promises"; +import { dirname, resolve } from "node:path"; +import { createInterface } from "node:readline"; +import type { JsonlSessionMetadata, SessionStorage, SessionTreeEntry } from "../../types.js"; + +interface SessionHeader { + type: "session"; + version: 3; + id: string; + timestamp: string; + cwd: string; + parentSession?: string; +} + +function updateLabelCache(labelsById: Map, entry: SessionTreeEntry): void { + if (entry.type !== "label") return; + const label = entry.label?.trim(); + if (label) { + labelsById.set(entry.targetId, label); + } else { + labelsById.delete(entry.targetId); + } +} + +function buildLabelsById(entries: SessionTreeEntry[]): Map { + const labelsById = new Map(); + for (const entry of entries) { + updateLabelCache(labelsById, entry); + } + return labelsById; +} + +function generateEntryId(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(); +} + +function headerToSessionMetadata(header: SessionHeader, path: string): JsonlSessionMetadata { + return { + id: header.id, + createdAt: header.timestamp, + cwd: header.cwd, + path, + parentSessionPath: header.parentSession, + }; +} + +export async function loadJsonlSessionMetadata(filePath: string): Promise { + const stream = createReadStream(filePath, { encoding: "utf8" }); + const lines = createInterface({ input: stream, crlfDelay: Infinity }); + try { + for await (const line of lines) { + if (!line.trim()) break; + try { + const header = JSON.parse(line) as SessionHeader; + return headerToSessionMetadata(header, resolve(filePath)); + } catch { + throw new Error(`Invalid JSONL session file ${filePath}: first line is not a valid session header`); + } + } + throw new Error(`Invalid JSONL session file ${filePath}: missing session header`); + } finally { + lines.close(); + stream.destroy(); + } +} + +async function loadJsonlStorage(filePath: string): Promise<{ + header: SessionHeader; + entries: SessionTreeEntry[]; + leafId: string | null; +}> { + const content = await readFile(filePath, "utf8"); + const lines = content.split("\n").filter((line) => line.trim()); + if (lines.length === 0) { + throw new Error(`Invalid JSONL session file ${filePath}: missing session header`); + } + + let header: SessionHeader; + try { + header = JSON.parse(lines[0]!) as SessionHeader; + } catch { + throw new Error(`Invalid JSONL session file ${filePath}: first line is not a valid session header`); + } + + const entries: SessionTreeEntry[] = []; + let leafId: string | null = null; + for (const line of lines.slice(1)) { + try { + const entry = JSON.parse(line) as SessionTreeEntry; + entries.push(entry); + leafId = entry.id; + } catch { + // ignore malformed entry lines + } + } + return { header, entries, leafId }; +} + +export class JsonlSessionStorage implements SessionStorage { + private readonly filePath: string; + private readonly metadata: JsonlSessionMetadata; + private entries: SessionTreeEntry[]; + private byId: Map; + private labelsById: Map; + private currentLeafId: string | null; + + private constructor(filePath: string, header: SessionHeader, entries: SessionTreeEntry[], leafId: string | null) { + this.filePath = resolve(filePath); + this.metadata = headerToSessionMetadata(header, this.filePath); + this.entries = entries; + this.byId = new Map(entries.map((entry) => [entry.id, entry])); + this.labelsById = buildLabelsById(entries); + this.currentLeafId = leafId; + } + + static async open(filePath: string): Promise { + const resolvedPath = resolve(filePath); + const loaded = await loadJsonlStorage(resolvedPath); + return new JsonlSessionStorage(resolvedPath, loaded.header, loaded.entries, loaded.leafId); + } + + static async create( + filePath: string, + options: { + cwd: string; + sessionId: string; + parentSessionPath?: string; + }, + ): Promise { + const resolvedPath = resolve(filePath); + const header: SessionHeader = { + type: "session", + version: 3, + id: options.sessionId, + timestamp: new Date().toISOString(), + cwd: options.cwd, + parentSession: options.parentSessionPath, + }; + await mkdir(dirname(resolvedPath), { recursive: true }); + await writeFile(resolvedPath, `${JSON.stringify(header)}\n`); + return new JsonlSessionStorage(resolvedPath, header, [], null); + } + + async getMetadata(): Promise { + return this.metadata; + } + + async getLeafId(): Promise { + return this.currentLeafId; + } + + async setLeafId(leafId: string | null): Promise { + if (leafId !== null && !this.byId.has(leafId)) { + throw new Error(`Entry ${leafId} not found`); + } + this.currentLeafId = leafId; + } + + async createEntryId(): Promise { + return generateEntryId(this.byId); + } + + async appendEntry(entry: SessionTreeEntry): Promise { + await appendFile(this.filePath, `${JSON.stringify(entry)}\n`); + this.entries.push(entry); + this.byId.set(entry.id, entry); + updateLabelCache(this.labelsById, entry); + this.currentLeafId = entry.id; + } + + async getEntry(id: string): Promise { + return this.byId.get(id); + } + + async findEntries( + type: TType, + ): Promise>> { + return this.entries.filter((entry): entry is Extract => entry.type === type); + } + + async getLabel(id: string): Promise { + return this.labelsById.get(id); + } + + async getPathToRoot(leafId: string | null): Promise { + 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 { + return [...this.entries]; + } +} diff --git a/packages/agent/src/vendor/pi-agent-core/harness/session/storage/memory.ts b/packages/agent/src/vendor/pi-agent-core/harness/session/storage/memory.ts new file mode 100644 index 00000000..652f633f --- /dev/null +++ b/packages/agent/src/vendor/pi-agent-core/harness/session/storage/memory.ts @@ -0,0 +1,103 @@ +import { randomUUID } from "node:crypto"; +import type { SessionMetadata, SessionStorage, SessionTreeEntry } from "../../types.js"; +import { uuidv7 } from "../uuid.js"; + +function updateLabelCache(labelsById: Map, entry: SessionTreeEntry): void { + if (entry.type !== "label") return; + const label = entry.label?.trim(); + if (label) { + labelsById.set(entry.targetId, label); + } else { + labelsById.delete(entry.targetId); + } +} + +function buildLabelsById(entries: SessionTreeEntry[]): Map { + const labelsById = new Map(); + for (const entry of entries) { + updateLabelCache(labelsById, entry); + } + return labelsById; +} + +function generateEntryId(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 class InMemorySessionStorage implements SessionStorage { + private readonly metadata: SessionMetadata; + private entries: SessionTreeEntry[]; + private byId: Map; + private labelsById: Map; + private leafId: string | null; + + constructor(options?: { entries?: SessionTreeEntry[]; leafId?: string | null; metadata?: SessionMetadata }) { + this.entries = options?.entries ? [...options.entries] : []; + this.byId = new Map(this.entries.map((entry) => [entry.id, entry])); + this.labelsById = buildLabelsById(this.entries); + this.leafId = options?.leafId ?? this.entries[this.entries.length - 1]?.id ?? null; + if (this.leafId !== null && !this.byId.has(this.leafId)) { + throw new Error(`Entry ${this.leafId} not found`); + } + this.metadata = options?.metadata ?? { id: uuidv7(), createdAt: new Date().toISOString() }; + } + + async getMetadata(): Promise { + return this.metadata; + } + + async getLeafId(): Promise { + return this.leafId; + } + + async setLeafId(leafId: string | null): Promise { + if (leafId !== null && !this.byId.has(leafId)) { + throw new Error(`Entry ${leafId} not found`); + } + this.leafId = leafId; + } + + async createEntryId(): Promise { + return generateEntryId(this.byId); + } + + async appendEntry(entry: SessionTreeEntry): Promise { + this.entries.push(entry); + this.byId.set(entry.id, entry); + updateLabelCache(this.labelsById, entry); + this.leafId = entry.id; + } + + async getEntry(id: string): Promise { + return this.byId.get(id); + } + + async findEntries( + type: TType, + ): Promise>> { + return this.entries.filter((entry): entry is Extract => entry.type === type); + } + + async getLabel(id: string): Promise { + return this.labelsById.get(id); + } + + async getPathToRoot(leafId: string | null): Promise { + 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 { + return [...this.entries]; + } +} diff --git a/packages/agent/src/vendor/pi-agent-core/harness/session/uuid.ts b/packages/agent/src/vendor/pi-agent-core/harness/session/uuid.ts new file mode 100644 index 00000000..c7e2896e --- /dev/null +++ b/packages/agent/src/vendor/pi-agent-core/harness/session/uuid.ts @@ -0,0 +1,44 @@ +import { randomBytes } from "node:crypto"; + +let lastTimestamp = -Infinity; +let sequence = 0; + +export function uuidv7(): string { + const random = randomBytes(16); + const timestamp = Date.now(); + + if (timestamp > lastTimestamp) { + sequence = random[6] * 0x1000000 + random[7] * 0x10000 + random[8] * 0x100 + random[9]; + lastTimestamp = timestamp; + } else { + sequence = (sequence + 1) >>> 0; + if (sequence === 0) { + lastTimestamp++; + } + } + + const bytes = new Uint8Array(16); + bytes[0] = (lastTimestamp / 0x10000000000) & 0xff; + bytes[1] = (lastTimestamp / 0x100000000) & 0xff; + bytes[2] = (lastTimestamp / 0x1000000) & 0xff; + bytes[3] = (lastTimestamp / 0x10000) & 0xff; + bytes[4] = (lastTimestamp / 0x100) & 0xff; + bytes[5] = lastTimestamp & 0xff; + bytes[6] = 0x70 | ((sequence >>> 28) & 0x0f); + bytes[7] = (sequence >>> 20) & 0xff; + bytes[8] = 0x80 | ((sequence >>> 14) & 0x3f); + bytes[9] = (sequence >>> 6) & 0xff; + bytes[10] = ((sequence & 0x3f) << 2) | (random[10] & 0x03); + bytes[11] = random[11]; + bytes[12] = random[12]; + bytes[13] = random[13]; + bytes[14] = random[14]; + bytes[15] = random[15]; + + return formatUuid(bytes); +} + +function formatUuid(bytes: Uint8Array): string { + const hex = Array.from(bytes, (byte) => byte.toString(16).padStart(2, "0")); + return `${hex.slice(0, 4).join("")}-${hex.slice(4, 6).join("")}-${hex.slice(6, 8).join("")}-${hex.slice(8, 10).join("")}-${hex.slice(10, 16).join("")}`; +} diff --git a/packages/agent/src/vendor/pi-agent-core/harness/skills.ts b/packages/agent/src/vendor/pi-agent-core/harness/skills.ts new file mode 100644 index 00000000..db03d93c --- /dev/null +++ b/packages/agent/src/vendor/pi-agent-core/harness/skills.ts @@ -0,0 +1,303 @@ +import ignore from "ignore"; +import { parse } from "yaml"; +import type { ExecutionEnv, Skill } from "./types.js"; + +const MAX_NAME_LENGTH = 64; +const MAX_DESCRIPTION_LENGTH = 1024; +const IGNORE_FILE_NAMES = [".gitignore", ".ignore", ".fdignore"]; + +type IgnoreMatcher = ReturnType; + +/** Warning produced while loading skills. */ +export interface SkillDiagnostic { + /** Diagnostic severity. Currently only warnings are emitted. */ + type: "warning"; + /** Human-readable diagnostic message. */ + message: string; + /** Path associated with the diagnostic. */ + path: string; +} + +interface SkillFrontmatter { + name?: string; + description?: string; + "disable-model-invocation"?: boolean; + [key: string]: unknown; +} + +/** Format a skill invocation prompt, optionally appending additional user instructions. */ +export function formatSkillInvocation(skill: Skill, additionalInstructions?: string): string { + const skillBlock = `\nReferences are relative to ${dirnameEnvPath(skill.filePath)}.\n\n${skill.content}\n`; + return additionalInstructions ? `${skillBlock}\n\n${additionalInstructions}` : skillBlock; +} + +/** + * Load skills from one or more directories. + * + * Traverses directories recursively, loads `SKILL.md` files, loads direct root `.md` files as skills, honors ignore files, + * and returns diagnostics for invalid skill files. Missing input directories are skipped. + */ +export async function loadSkills( + env: ExecutionEnv, + dirs: string | string[], +): Promise<{ skills: Skill[]; diagnostics: SkillDiagnostic[] }> { + const skills: Skill[] = []; + const diagnostics: SkillDiagnostic[] = []; + for (const dir of Array.isArray(dirs) ? dirs : [dirs]) { + const rootInfo = await safeFileInfo(env, dir); + if (!rootInfo || (await resolveKind(env, rootInfo)) !== "directory") continue; + const result = await loadSkillsFromDirInternal(env, rootInfo.path, true, ignore(), rootInfo.path); + skills.push(...result.skills); + diagnostics.push(...result.diagnostics); + } + return { skills, diagnostics }; +} + +/** + * Load skills from source-tagged directories. + * + * Source values are preserved exactly and attached to every loaded skill and diagnostic. The agent package does not + * interpret source values; applications define their own provenance shape. + */ +export async function loadSourcedSkills( + env: ExecutionEnv, + inputs: Array<{ path: string; source: TSource }>, + mapSkill?: (skill: Skill, source: TSource) => TSkill, +): Promise<{ + skills: Array<{ skill: TSkill; source: TSource }>; + diagnostics: Array; +}> { + const skills: Array<{ skill: TSkill; source: TSource }> = []; + const diagnostics: Array = []; + for (const input of inputs) { + const result = await loadSkills(env, input.path); + for (const skill of result.skills) { + skills.push({ skill: mapSkill ? mapSkill(skill, input.source) : (skill as TSkill), source: input.source }); + } + for (const diagnostic of result.diagnostics) diagnostics.push({ ...diagnostic, source: input.source }); + } + return { skills, diagnostics }; +} + +async function loadSkillsFromDirInternal( + env: ExecutionEnv, + dir: string, + includeRootFiles: boolean, + ignoreMatcher: IgnoreMatcher, + rootDir: string, +): Promise<{ skills: Skill[]; diagnostics: SkillDiagnostic[] }> { + const skills: Skill[] = []; + const diagnostics: SkillDiagnostic[] = []; + + if (!(await env.exists(dir))) return { skills, diagnostics }; + const dirInfo = await safeFileInfo(env, dir); + if (!dirInfo || (await resolveKind(env, dirInfo)) !== "directory") return { skills, diagnostics }; + + await addIgnoreRules(env, ignoreMatcher, dir, rootDir); + + let entries: Awaited>; + try { + entries = await env.listDir(dir); + } catch { + return { skills, diagnostics }; + } + + for (const entry of entries) { + if (entry.name !== "SKILL.md") continue; + const fullPath = entry.path; + const kind = await resolveKind(env, entry); + if (kind !== "file") continue; + const relPath = relativeEnvPath(rootDir, fullPath); + if (ignoreMatcher.ignores(relPath)) continue; + + const result = await loadSkillFromFile(env, fullPath); + if (result.skill) skills.push(result.skill); + diagnostics.push(...result.diagnostics); + return { skills, diagnostics }; + } + + for (const entry of entries.sort((a, b) => a.name.localeCompare(b.name))) { + if (entry.name.startsWith(".") || entry.name === "node_modules") continue; + const fullPath = entry.path; + const kind = await resolveKind(env, entry); + if (!kind) continue; + + const relPath = relativeEnvPath(rootDir, fullPath); + const ignorePath = kind === "directory" ? `${relPath}/` : relPath; + if (ignoreMatcher.ignores(ignorePath)) continue; + + if (kind === "directory") { + const result = await loadSkillsFromDirInternal(env, fullPath, false, ignoreMatcher, rootDir); + skills.push(...result.skills); + diagnostics.push(...result.diagnostics); + continue; + } + + if (kind !== "file" || !includeRootFiles || !entry.name.endsWith(".md")) continue; + const result = await loadSkillFromFile(env, fullPath); + if (result.skill) skills.push(result.skill); + diagnostics.push(...result.diagnostics); + } + + return { skills, diagnostics }; +} + +async function addIgnoreRules(env: ExecutionEnv, ig: IgnoreMatcher, dir: string, rootDir: string): Promise { + const relativeDir = relativeEnvPath(rootDir, dir); + const prefix = relativeDir ? `${relativeDir}/` : ""; + + for (const filename of IGNORE_FILE_NAMES) { + const ignorePath = joinEnvPath(dir, filename); + const info = await safeFileInfo(env, ignorePath); + if (info?.kind !== "file") continue; + try { + const content = await env.readTextFile(ignorePath); + const patterns = content + .split(/\r?\n/) + .map((line) => prefixIgnorePattern(line, prefix)) + .filter((line): line is string => Boolean(line)); + if (patterns.length > 0) ig.add(patterns); + } catch {} + } +} + +function prefixIgnorePattern(line: string, prefix: string): string | null { + const trimmed = line.trim(); + if (!trimmed) return null; + if (trimmed.startsWith("#") && !trimmed.startsWith("\\#")) return null; + + let pattern = line; + let negated = false; + if (pattern.startsWith("!")) { + negated = true; + pattern = pattern.slice(1); + } else if (pattern.startsWith("\\!")) { + pattern = pattern.slice(1); + } + if (pattern.startsWith("/")) pattern = pattern.slice(1); + const prefixed = prefix ? `${prefix}${pattern}` : pattern; + return negated ? `!${prefixed}` : prefixed; +} + +async function loadSkillFromFile( + env: ExecutionEnv, + filePath: string, +): Promise<{ skill: Skill | null; diagnostics: SkillDiagnostic[] }> { + const diagnostics: SkillDiagnostic[] = []; + try { + const rawContent = await env.readTextFile(filePath); + const { frontmatter, body } = parseFrontmatter(rawContent); + const skillDir = dirnameEnvPath(filePath); + const parentDirName = basenameEnvPath(skillDir); + + for (const error of validateDescription(frontmatter.description)) { + diagnostics.push({ type: "warning", message: error, path: filePath }); + } + + const name = frontmatter.name || parentDirName; + for (const error of validateName(name, parentDirName)) { + diagnostics.push({ type: "warning", message: error, path: filePath }); + } + + if (!frontmatter.description || frontmatter.description.trim() === "") { + return { skill: null, diagnostics }; + } + + return { + skill: { + name, + description: frontmatter.description, + content: body, + filePath, + disableModelInvocation: frontmatter["disable-model-invocation"] === true, + }, + diagnostics, + }; + } catch (error) { + const message = error instanceof Error ? error.message : "failed to parse skill file"; + diagnostics.push({ type: "warning", message, path: filePath }); + return { skill: null, diagnostics }; + } +} + +function validateName(name: string, parentDirName: string): string[] { + const errors: string[] = []; + if (name !== parentDirName) errors.push(`name "${name}" does not match parent directory "${parentDirName}"`); + if (name.length > MAX_NAME_LENGTH) errors.push(`name exceeds ${MAX_NAME_LENGTH} characters (${name.length})`); + if (!/^[a-z0-9-]+$/.test(name)) { + errors.push("name contains invalid characters (must be lowercase a-z, 0-9, hyphens only)"); + } + if (name.startsWith("-") || name.endsWith("-")) errors.push("name must not start or end with a hyphen"); + if (name.includes("--")) errors.push("name must not contain consecutive hyphens"); + return errors; +} + +function validateDescription(description: string | undefined): string[] { + const errors: string[] = []; + if (!description || description.trim() === "") { + errors.push("description is required"); + } else if (description.length > MAX_DESCRIPTION_LENGTH) { + errors.push(`description exceeds ${MAX_DESCRIPTION_LENGTH} characters (${description.length})`); + } + return errors; +} + +function parseFrontmatter>(content: string): { frontmatter: T; body: string } { + const normalized = content.replace(/\r\n/g, "\n").replace(/\r/g, "\n"); + if (!normalized.startsWith("---")) return { frontmatter: {} as T, body: normalized }; + const endIndex = normalized.indexOf("\n---", 3); + if (endIndex === -1) return { frontmatter: {} as T, body: normalized }; + const yamlString = normalized.slice(4, endIndex); + const body = normalized.slice(endIndex + 4).trim(); + return { frontmatter: (parse(yamlString) ?? {}) as T, body }; +} + +async function safeFileInfo( + env: ExecutionEnv, + path: string, +): Promise> | undefined> { + try { + return await env.fileInfo(path); + } catch { + return undefined; + } +} + +async function resolveKind( + env: ExecutionEnv, + info: Awaited>, +): Promise<"file" | "directory" | undefined> { + if (info.kind === "file" || info.kind === "directory") return info.kind; + try { + const realPath = await env.realPath(info.path); + const target = await env.fileInfo(realPath); + return target.kind === "file" || target.kind === "directory" ? target.kind : undefined; + } catch { + return undefined; + } +} + +function joinEnvPath(base: string, child: string): string { + return `${base.replace(/\/+$/, "")}/${child.replace(/^\/+/, "")}`; +} + +function dirnameEnvPath(path: string): string { + const normalized = path.replace(/\/+$/, ""); + const slashIndex = normalized.lastIndexOf("/"); + return slashIndex <= 0 ? "/" : normalized.slice(0, slashIndex); +} + +function basenameEnvPath(path: string): string { + const normalized = path.replace(/\/+$/, ""); + const slashIndex = normalized.lastIndexOf("/"); + return slashIndex === -1 ? normalized : normalized.slice(slashIndex + 1); +} + +function relativeEnvPath(root: string, path: string): string { + const normalizedRoot = root.replace(/\/+$/, ""); + const normalizedPath = path.replace(/\/+$/, ""); + if (normalizedPath === normalizedRoot) return ""; + return normalizedPath.startsWith(`${normalizedRoot}/`) + ? normalizedPath.slice(normalizedRoot.length + 1) + : normalizedPath.replace(/^\/+/, ""); +} diff --git a/packages/agent/src/vendor/pi-agent-core/harness/system-prompt.ts b/packages/agent/src/vendor/pi-agent-core/harness/system-prompt.ts new file mode 100644 index 00000000..44b8f623 --- /dev/null +++ b/packages/agent/src/vendor/pi-agent-core/harness/system-prompt.ts @@ -0,0 +1,34 @@ +import type { Skill } from "./types.js"; + +export function formatSkillsForSystemPrompt(skills: Skill[]): string { + const visibleSkills = skills.filter((skill) => !skill.disableModelInvocation); + if (visibleSkills.length === 0) return ""; + + const lines = [ + "The following skills provide specialized instructions for specific tasks.", + "Read the full skill file when the task matches its description.", + "When a skill file references a relative path, resolve it against the skill directory (parent of SKILL.md / dirname of the path) and use that absolute path in tool commands.", + "", + "", + ]; + + for (const skill of visibleSkills) { + lines.push(" "); + lines.push(` ${escapeXml(skill.name)}`); + lines.push(` ${escapeXml(skill.description)}`); + lines.push(` ${escapeXml(skill.filePath)}`); + lines.push(" "); + } + + lines.push(""); + return lines.join("\n"); +} + +function escapeXml(value: string): string { + return value + .replace(/&/g, "&") + .replace(//g, ">") + .replace(/"/g, """) + .replace(/'/g, "'"); +} diff --git a/packages/agent/src/vendor/pi-agent-core/harness/types.ts b/packages/agent/src/vendor/pi-agent-core/harness/types.ts new file mode 100644 index 00000000..f50337f8 --- /dev/null +++ b/packages/agent/src/vendor/pi-agent-core/harness/types.ts @@ -0,0 +1,652 @@ +import type { ImageContent, Model, SimpleStreamOptions, TextContent, Transport } from "@earendil-works/pi-ai"; +import type { QueueMode } from "../agent.js"; +import type { AgentEvent, AgentMessage, AgentTool, ThinkingLevel } from "../index.js"; +import type { Session } from "./session/session.js"; + +/** + * Skill loaded from a `SKILL.md` file or provided by an application. + * + * `name`, `description`, and `filePath` are inserted into the system prompt in an XML-formatted block as suggested by agentskills.io. + * Use {@link formatSkillsForSystemPrompt} to generate the spec-compatible system prompt block. + */ +export interface Skill { + /** Stable skill name used for lookup and model-visible listings. */ + name: string; + /** Short model-visible description of when to use the skill. */ + description: string; + /** Full skill instructions. */ + content: string; + /** Absolute path to the skill file. Used for model-visible location and resolving relative references. */ + filePath: string; + /** Exclude this skill from model-visible skill lists while still allowing explicit application invocation. */ + disableModelInvocation?: boolean; +} + +/** Prompt template that can be formatted into a prompt for explicit invocation. */ +export interface PromptTemplate { + /** Stable template name used for lookup or application command routing. */ + name: string; + /** Optional description for command lists or autocomplete. */ + description?: string; + /** Template content. Argument placeholders are formatted by `formatPromptTemplateInvocation`. */ + content: string; +} + +/** Resources made available to explicit invocation methods and system-prompt callbacks. */ +export interface AgentHarnessResources< + TSkill extends Skill = Skill, + TPromptTemplate extends PromptTemplate = PromptTemplate, +> { + /** Prompt templates available for explicit invocation. */ + promptTemplates?: TPromptTemplate[]; + /** Skills available to the model and explicit skill invocation. */ + skills?: TSkill[]; +} + +/** Curated provider request options owned by the harness and snapshotted per turn. */ +export interface AgentHarnessStreamOptions { + /** Preferred transport forwarded to the stream function. */ + transport?: Transport; + /** Provider request timeout in milliseconds. */ + timeoutMs?: number; + /** Maximum provider retry attempts. */ + maxRetries?: number; + /** Optional cap for provider-requested retry delays. */ + maxRetryDelayMs?: number; + /** Additional request headers merged with auth and lifecycle headers. */ + headers?: Record; + /** Provider metadata forwarded with requests. */ + metadata?: SimpleStreamOptions["metadata"]; + /** Provider cache retention hint. */ + cacheRetention?: SimpleStreamOptions["cacheRetention"]; +} + +/** Per-request stream option patch returned by provider hooks. */ +export interface AgentHarnessStreamOptionsPatch + extends Omit, "headers" | "metadata"> { + /** Header patch. `undefined` values delete keys; explicit `headers: undefined` clears all headers. */ + headers?: Record; + /** Metadata patch. `undefined` values delete keys; explicit `metadata: undefined` clears all metadata. */ + metadata?: Record; +} + +/** Kind of filesystem object as addressed by an {@link ExecutionEnv}. Symlinks are not followed automatically. */ +export type FileKind = "file" | "directory" | "symlink"; + +/** Stable, backend-independent file error codes thrown by {@link ExecutionEnv} file operations. */ +export type FileErrorCode = + | "not_found" + | "permission_denied" + | "not_directory" + | "is_directory" + | "invalid" + | "not_supported" + | "unknown"; + +/** Error thrown by {@link ExecutionEnv} file operations. */ +export class FileError extends Error { + constructor( + /** Backend-independent error code. */ + public code: FileErrorCode, + message: string, + /** Absolute addressed path associated with the failure, when available. */ + public path?: string, + options?: ErrorOptions, + ) { + super(message, options); + this.name = "FileError"; + } +} + +/** Metadata for one filesystem object in an {@link ExecutionEnv}. */ +export interface FileInfo { + /** Basename of {@link path}. */ + name: string; + /** Absolute, syntactically normalized addressed path in the execution environment. Symlinks are not followed. */ + path: string; + /** Object kind. Symlink targets are not followed; use {@link ExecutionEnv.resolvePath} explicitly. */ + kind: FileKind; + /** Size in bytes for the addressed filesystem object. */ + size: number; + /** Modification time as milliseconds since Unix epoch. */ + mtimeMs: number; +} + +/** Options for {@link ExecutionEnv.exec}. */ +export interface ExecutionEnvExecOptions { + /** Working directory for the command. Relative paths are resolved against {@link ExecutionEnv.cwd}. */ + cwd?: string; + /** Additional environment variables for the command. Values override the environment defaults. */ + env?: Record; + /** Timeout in seconds. Implementations should reject when the command exceeds this duration. */ + timeout?: number; + /** Abort signal used to terminate the command. */ + signal?: AbortSignal; + /** Called with stdout chunks as they are produced. */ + onStdout?: (chunk: string) => void; + /** Called with stderr chunks as they are produced. */ + onStderr?: (chunk: string) => void; +} + +/** + * Filesystem and process execution environment used by the harness. + * + * Paths passed to methods may be absolute or relative to {@link cwd}. Paths returned by this interface are absolute + * addressed paths in the environment, but are not canonicalized through symlinks unless returned by {@link resolvePath}. + * + * File operations throw {@link FileError} for expected filesystem failures such as missing paths or permission errors. + */ +export interface ExecutionEnv { + /** Current working directory for relative paths and command execution. */ + cwd: string; + + /** Execute a shell command in {@link cwd} unless `options.cwd` is provided. */ + exec( + command: string, + options?: ExecutionEnvExecOptions, + ): Promise<{ stdout: string; stderr: string; exitCode: number }>; + + /** Read a UTF-8 text file. Throws {@link FileError}. */ + readTextFile(path: string): Promise; + /** Read a binary file. Throws {@link FileError}. */ + readBinaryFile(path: string): Promise; + /** Create or overwrite a file, creating parent directories when supported. Throws {@link FileError}. */ + writeFile(path: string, content: string | Uint8Array): Promise; + /** Return metadata for the addressed path without following symlinks. Throws {@link FileError}. */ + fileInfo(path: string): Promise; + /** List direct children of a directory without following symlinks. Throws {@link FileError}. */ + listDir(path: string): Promise; + /** Return the canonical path for a path, following symlinks. Throws {@link FileError}. */ + realPath(path: string): Promise; + /** Return false for missing paths. Other errors, such as permission failures, may throw {@link FileError}. */ + exists(path: string): Promise; + /** Create a directory. */ + createDir(path: string, options?: { recursive?: boolean }): Promise; + /** Remove a file or directory. */ + remove(path: string, options?: { recursive?: boolean; force?: boolean }): Promise; + /** Create a temporary directory and return its absolute path. */ + createTempDir(prefix?: string): Promise; + /** Create a temporary file and return its absolute path. */ + createTempFile(options?: { prefix?: string; suffix?: string }): Promise; + + /** Release resources owned by the environment. */ + 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"; // legacy name, kept for backwards compatibility + 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 SessionMetadata { + id: string; + createdAt: string; +} + +export interface JsonlSessionMetadata extends SessionMetadata { + cwd: string; + path: string; + parentSessionPath?: string; +} + +export interface SessionStorage { + getMetadata(): Promise; + getLeafId(): Promise; + setLeafId(leafId: string | null): Promise; + createEntryId(): Promise; + appendEntry(entry: SessionTreeEntry): Promise; + getEntry(id: string): Promise; + findEntries( + type: TType, + ): Promise>>; + getLabel(id: string): Promise; + getPathToRoot(leafId: string | null): Promise; + getEntries(): Promise; +} + +export type { Session } from "./session/session.js"; + +export interface SessionCreateOptions { + id?: string; +} + +export interface SessionForkOptions { + entryId?: string; + position?: "before" | "at"; + id?: string; +} + +export interface SessionRepo< + TMetadata extends SessionMetadata = SessionMetadata, + TCreateOptions extends SessionCreateOptions = SessionCreateOptions, + TListOptions = void, +> { + create(options: TCreateOptions): Promise>; + open(metadata: TMetadata): Promise>; + list(options?: TListOptions): Promise; + delete(metadata: TMetadata): Promise; + fork(source: TMetadata, options: SessionForkOptions & TCreateOptions): Promise>; +} + +export interface JsonlSessionCreateOptions extends SessionCreateOptions { + cwd: string; + parentSessionPath?: string; +} + +export interface JsonlSessionListOptions { + cwd?: string; +} + +export interface JsonlSessionRepoApi + extends SessionRepo {} + +export type AgentHarnessPhase = "idle" | "turn" | "compaction" | "branch_summary" | "retry"; + +export type PendingSessionWrite = SessionTreeEntry extends infer TEntry + ? TEntry extends SessionTreeEntry + ? Omit + : never + : never; + +export interface QueueUpdateEvent { + type: "queue_update"; + steer: AgentMessage[]; + followUp: AgentMessage[]; + nextTurn: AgentMessage[]; +} + +export interface SavePointEvent { + type: "save_point"; + hadPendingMutations: boolean; +} + +export interface AbortEvent { + type: "abort"; + clearedSteer: AgentMessage[]; + clearedFollowUp: AgentMessage[]; +} + +export interface SettledEvent { + type: "settled"; + nextTurnCount: number; +} + +export interface BeforeAgentStartEvent< + TSkill extends Skill = Skill, + TPromptTemplate extends PromptTemplate = PromptTemplate, +> { + type: "before_agent_start"; + prompt: string; + images?: ImageContent[]; + systemPrompt: string; + resources: AgentHarnessResources; +} + +export interface ContextEvent { + type: "context"; + messages: AgentMessage[]; +} + +export interface BeforeProviderRequestEvent { + type: "before_provider_request"; + model: Model; + sessionId: string; + streamOptions: AgentHarnessStreamOptions; +} + +export interface BeforeProviderPayloadEvent { + type: "before_provider_payload"; + model: Model; + 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 interface ResourcesUpdateEvent< + TSkill extends Skill = Skill, + TPromptTemplate extends PromptTemplate = PromptTemplate, +> { + type: "resources_update"; + resources: AgentHarnessResources; + previousResources: AgentHarnessResources; +} + +export type AgentHarnessOwnEvent< + TSkill extends Skill = Skill, + TPromptTemplate extends PromptTemplate = PromptTemplate, +> = + | QueueUpdateEvent + | SavePointEvent + | AbortEvent + | SettledEvent + | BeforeAgentStartEvent + | ContextEvent + | BeforeProviderRequestEvent + | BeforeProviderPayloadEvent + | AfterProviderResponseEvent + | ToolCallEvent + | ToolResultEvent + | SessionBeforeCompactEvent + | SessionCompactEvent + | SessionBeforeTreeEvent + | SessionTreeEvent + | ModelSelectEvent + | ThinkingLevelSelectEvent + | ResourcesUpdateEvent; + +export type AgentHarnessEvent = + | AgentEvent + | AgentHarnessOwnEvent; + +export interface BeforeAgentStartResult { + messages?: AgentMessage[]; + systemPrompt?: string; +} + +export interface ContextResult { + messages: AgentMessage[]; +} + +export interface BeforeProviderRequestResult { + streamOptions?: AgentHarnessStreamOptionsPatch; +} + +export interface BeforeProviderPayloadResult { + 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; + before_provider_payload: BeforeProviderPayloadResult | 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; + resources_update: 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< + TSkill extends Skill = Skill, + TPromptTemplate extends PromptTemplate = PromptTemplate, + TTool extends AgentTool = AgentTool, +> { + env: ExecutionEnv; + session: Session; + tools?: TTool[]; + /** + * Concrete resources available to explicit invocation methods and system-prompt callbacks. + * Applications own loading/reloading resources and should call `setResources()` with new values. + */ + resources?: AgentHarnessResources; + systemPrompt?: + | string + | ((context: { + env: ExecutionEnv; + session: Session; + model: Model; + thinkingLevel: ThinkingLevel; + activeTools: TTool[]; + resources: AgentHarnessResources; + }) => string | Promise); + getApiKeyAndHeaders?: ( + model: Model, + ) => Promise<{ apiKey: string; headers?: Record } | undefined>; + /** Curated stream/provider request options. Snapshotted at turn start. */ + streamOptions?: AgentHarnessStreamOptions; + model: Model; + thinkingLevel?: ThinkingLevel; + activeToolNames?: string[]; + steeringMode?: QueueMode; + followUpMode?: QueueMode; +} + +export type { AgentHarness } from "./agent-harness.js"; diff --git a/packages/agent/src/vendor/pi-agent-core/harness/utils/shell-output.ts b/packages/agent/src/vendor/pi-agent-core/harness/utils/shell-output.ts new file mode 100644 index 00000000..d31d7a06 --- /dev/null +++ b/packages/agent/src/vendor/pi-agent-core/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/vendor/pi-agent-core/harness/utils/truncate.ts b/packages/agent/src/vendor/pi-agent-core/harness/utils/truncate.ts new file mode 100644 index 00000000..18ac5d74 --- /dev/null +++ b/packages/agent/src/vendor/pi-agent-core/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/vendor/pi-agent-core/index.ts b/packages/agent/src/vendor/pi-agent-core/index.ts new file mode 100644 index 00000000..293ce196 --- /dev/null +++ b/packages/agent/src/vendor/pi-agent-core/index.ts @@ -0,0 +1,42 @@ +// Core Agent +export * from "./agent.js"; +// Loop functions +export * from "./agent-loop.js"; +export * from "./harness/agent-harness.js"; +export { + collectEntriesForBranchSummary, + generateBranchSummary, + prepareBranchEntries, +} from "./harness/compaction/branch-summarization.js"; +export { + calculateContextTokens, + compact, + DEFAULT_COMPACTION_SETTINGS, + estimateContextTokens, + estimateTokens, + findCutPoint, + findTurnStartIndex, + generateSummary, + getLastAssistantUsage, + prepareCompaction, + serializeConversation, + shouldCompact, +} from "./harness/compaction/compaction.js"; +export * from "./harness/execution-env.js"; +export * from "./harness/messages.js"; +export * from "./harness/prompt-templates.js"; +export * from "./harness/session/repo/jsonl.js"; +export * from "./harness/session/repo/memory.js"; +export * from "./harness/session/repo/shared.js"; +export * from "./harness/session/session.js"; +export { uuidv7 } from "./harness/session/uuid.js"; +export * from "./harness/skills.js"; +export * from "./harness/system-prompt.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 +export * from "./types.js"; diff --git a/packages/agent/src/vendor/pi-agent-core/proxy.ts b/packages/agent/src/vendor/pi-agent-core/proxy.ts new file mode 100644 index 00000000..5f0925c9 --- /dev/null +++ b/packages/agent/src/vendor/pi-agent-core/proxy.ts @@ -0,0 +1,367 @@ +/** + * Proxy stream function for apps that route LLM calls through a server. + * The server manages auth and proxies requests to LLM providers. + */ + +// Internal import for JSON parsing utility +import { + type AssistantMessage, + type AssistantMessageEvent, + type Context, + EventStream, + type Model, + parseStreamingJson, + type SimpleStreamOptions, + type StopReason, + type ToolCall, +} from "@earendil-works/pi-ai"; + +// Create stream class matching ProxyMessageEventStream +class ProxyMessageEventStream extends EventStream { + constructor() { + super( + (event) => event.type === "done" || event.type === "error", + (event) => { + if (event.type === "done") return event.message; + if (event.type === "error") return event.error; + throw new Error("Unexpected event type"); + }, + ); + } +} + +/** + * Proxy event types - server sends these with partial field stripped to reduce bandwidth. + */ +export type ProxyAssistantMessageEvent = + | { type: "start" } + | { type: "text_start"; contentIndex: number } + | { type: "text_delta"; contentIndex: number; delta: string } + | { type: "text_end"; contentIndex: number; contentSignature?: string } + | { type: "thinking_start"; contentIndex: number } + | { type: "thinking_delta"; contentIndex: number; delta: string } + | { type: "thinking_end"; contentIndex: number; contentSignature?: string } + | { type: "toolcall_start"; contentIndex: number; id: string; toolName: string } + | { type: "toolcall_delta"; contentIndex: number; delta: string } + | { type: "toolcall_end"; contentIndex: number } + | { + type: "done"; + reason: Extract; + usage: AssistantMessage["usage"]; + } + | { + type: "error"; + reason: Extract; + errorMessage?: string; + usage: AssistantMessage["usage"]; + }; + +type ProxySerializableStreamOptions = Pick< + SimpleStreamOptions, + | "temperature" + | "maxTokens" + | "reasoning" + | "cacheRetention" + | "sessionId" + | "headers" + | "metadata" + | "transport" + | "thinkingBudgets" + | "maxRetryDelayMs" +>; + +export interface ProxyStreamOptions extends ProxySerializableStreamOptions { + /** Local abort signal for the proxy request */ + signal?: AbortSignal; + /** Auth token for the proxy server */ + authToken: string; + /** Proxy server URL (e.g., "https://genai.example.com") */ + proxyUrl: string; +} + +/** + * Stream function that proxies through a server instead of calling LLM providers directly. + * The server strips the partial field from delta events to reduce bandwidth. + * We reconstruct the partial message client-side. + * + * Use this as the `streamFn` option when creating an Agent that needs to go through a proxy. + * + * @example + * ```typescript + * const agent = new Agent({ + * streamFn: (model, context, options) => + * streamProxy(model, context, { + * ...options, + * authToken: await getAuthToken(), + * proxyUrl: "https://genai.example.com", + * }), + * }); + * ``` + */ +function buildProxyRequestOptions(options: ProxyStreamOptions): ProxySerializableStreamOptions { + return { + temperature: options.temperature, + maxTokens: options.maxTokens, + reasoning: options.reasoning, + cacheRetention: options.cacheRetention, + sessionId: options.sessionId, + headers: options.headers, + metadata: options.metadata, + transport: options.transport, + thinkingBudgets: options.thinkingBudgets, + maxRetryDelayMs: options.maxRetryDelayMs, + }; +} + +export function streamProxy(model: Model, context: Context, options: ProxyStreamOptions): ProxyMessageEventStream { + const stream = new ProxyMessageEventStream(); + + (async () => { + // Initialize the partial message that we'll build up from events + const partial: AssistantMessage = { + role: "assistant", + stopReason: "stop", + content: [], + api: model.api, + provider: model.provider, + model: model.id, + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + timestamp: Date.now(), + }; + + let reader: ReadableStreamDefaultReader | undefined; + + const abortHandler = () => { + if (reader) { + reader.cancel("Request aborted by user").catch(() => {}); + } + }; + + if (options.signal) { + options.signal.addEventListener("abort", abortHandler); + } + + try { + const response = await fetch(`${options.proxyUrl}/api/stream`, { + method: "POST", + headers: { + Authorization: `Bearer ${options.authToken}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ + model, + context, + options: buildProxyRequestOptions(options), + }), + signal: options.signal, + }); + + if (!response.ok) { + let errorMessage = `Proxy error: ${response.status} ${response.statusText}`; + try { + const errorData = (await response.json()) as { error?: string }; + if (errorData.error) { + errorMessage = `Proxy error: ${errorData.error}`; + } + } catch { + // Couldn't parse error response + } + throw new Error(errorMessage); + } + + reader = response.body!.getReader(); + const decoder = new TextDecoder(); + let buffer = ""; + + while (true) { + const { done, value } = await reader.read(); + if (done) break; + + if (options.signal?.aborted) { + throw new Error("Request aborted by user"); + } + + buffer += decoder.decode(value, { stream: true }); + const lines = buffer.split("\n"); + buffer = lines.pop() || ""; + + for (const line of lines) { + if (line.startsWith("data: ")) { + const data = line.slice(6).trim(); + if (data) { + const proxyEvent = JSON.parse(data) as ProxyAssistantMessageEvent; + const event = processProxyEvent(proxyEvent, partial); + if (event) { + stream.push(event); + } + } + } + } + } + + if (options.signal?.aborted) { + throw new Error("Request aborted by user"); + } + + stream.end(); + } catch (error) { + const errorMessage = error instanceof Error ? error.message : String(error); + const reason = options.signal?.aborted ? "aborted" : "error"; + partial.stopReason = reason; + partial.errorMessage = errorMessage; + stream.push({ + type: "error", + reason, + error: partial, + }); + stream.end(); + } finally { + if (options.signal) { + options.signal.removeEventListener("abort", abortHandler); + } + } + })(); + + return stream; +} + +/** + * Process a proxy event and update the partial message. + */ +function processProxyEvent( + proxyEvent: ProxyAssistantMessageEvent, + partial: AssistantMessage, +): AssistantMessageEvent | undefined { + switch (proxyEvent.type) { + case "start": + return { type: "start", partial }; + + case "text_start": + partial.content[proxyEvent.contentIndex] = { type: "text", text: "" }; + return { type: "text_start", contentIndex: proxyEvent.contentIndex, partial }; + + case "text_delta": { + const content = partial.content[proxyEvent.contentIndex]; + if (content?.type === "text") { + content.text += proxyEvent.delta; + return { + type: "text_delta", + contentIndex: proxyEvent.contentIndex, + delta: proxyEvent.delta, + partial, + }; + } + throw new Error("Received text_delta for non-text content"); + } + + case "text_end": { + const content = partial.content[proxyEvent.contentIndex]; + if (content?.type === "text") { + content.textSignature = proxyEvent.contentSignature; + return { + type: "text_end", + contentIndex: proxyEvent.contentIndex, + content: content.text, + partial, + }; + } + throw new Error("Received text_end for non-text content"); + } + + case "thinking_start": + partial.content[proxyEvent.contentIndex] = { type: "thinking", thinking: "" }; + return { type: "thinking_start", contentIndex: proxyEvent.contentIndex, partial }; + + case "thinking_delta": { + const content = partial.content[proxyEvent.contentIndex]; + if (content?.type === "thinking") { + content.thinking += proxyEvent.delta; + return { + type: "thinking_delta", + contentIndex: proxyEvent.contentIndex, + delta: proxyEvent.delta, + partial, + }; + } + throw new Error("Received thinking_delta for non-thinking content"); + } + + case "thinking_end": { + const content = partial.content[proxyEvent.contentIndex]; + if (content?.type === "thinking") { + content.thinkingSignature = proxyEvent.contentSignature; + return { + type: "thinking_end", + contentIndex: proxyEvent.contentIndex, + content: content.thinking, + partial, + }; + } + throw new Error("Received thinking_end for non-thinking content"); + } + + case "toolcall_start": + partial.content[proxyEvent.contentIndex] = { + type: "toolCall", + id: proxyEvent.id, + name: proxyEvent.toolName, + arguments: {}, + partialJson: "", + } satisfies ToolCall & { partialJson: string } as ToolCall; + return { type: "toolcall_start", contentIndex: proxyEvent.contentIndex, partial }; + + case "toolcall_delta": { + const content = partial.content[proxyEvent.contentIndex]; + if (content?.type === "toolCall") { + (content as any).partialJson += proxyEvent.delta; + content.arguments = parseStreamingJson((content as any).partialJson) || {}; + partial.content[proxyEvent.contentIndex] = { ...content }; // Trigger reactivity + return { + type: "toolcall_delta", + contentIndex: proxyEvent.contentIndex, + delta: proxyEvent.delta, + partial, + }; + } + throw new Error("Received toolcall_delta for non-toolCall content"); + } + + case "toolcall_end": { + const content = partial.content[proxyEvent.contentIndex]; + if (content?.type === "toolCall") { + delete (content as any).partialJson; + return { + type: "toolcall_end", + contentIndex: proxyEvent.contentIndex, + toolCall: content, + partial, + }; + } + return undefined; + } + + case "done": + partial.stopReason = proxyEvent.reason; + partial.usage = proxyEvent.usage; + return { type: "done", reason: proxyEvent.reason, message: partial }; + + case "error": + partial.stopReason = proxyEvent.reason; + partial.errorMessage = proxyEvent.errorMessage; + partial.usage = proxyEvent.usage; + return { type: "error", reason: proxyEvent.reason, error: partial }; + + default: { + const _exhaustiveCheck: never = proxyEvent; + console.warn(`Unhandled proxy event type: ${(proxyEvent as any).type}`); + return undefined; + } + } +} diff --git a/packages/agent/src/vendor/pi-agent-core/types.ts b/packages/agent/src/vendor/pi-agent-core/types.ts new file mode 100644 index 00000000..285c1b02 --- /dev/null +++ b/packages/agent/src/vendor/pi-agent-core/types.ts @@ -0,0 +1,410 @@ +import type { + AssistantMessage, + AssistantMessageEvent, + ImageContent, + Message, + Model, + SimpleStreamOptions, + streamSimple, + TextContent, + Tool, + ToolResultMessage, +} from "@earendil-works/pi-ai"; +import type { Static, TSchema } from "typebox"; + +/** + * Stream function used by the agent loop. + * + * Contract: + * - Must not throw or return a rejected promise for request/model/runtime failures. + * - Must return an AssistantMessageEventStream. + * - Failures must be encoded in the returned stream via protocol events and a + * final AssistantMessage with stopReason "error" or "aborted" and errorMessage. + */ +export type StreamFn = ( + ...args: Parameters +) => ReturnType | Promise>; + +/** + * Configuration for how tool calls from a single assistant message are executed. + * + * - "sequential": each tool call is prepared, executed, and finalized before the next one starts. + * - "parallel": tool calls are prepared sequentially, then allowed tools execute concurrently. + * `tool_execution_end` is emitted in tool completion order after each tool is finalized, + * while tool-result message artifacts are emitted later in assistant source order. + */ +export type ToolExecutionMode = "sequential" | "parallel"; + +/** A single tool call content block emitted by an assistant message. */ +export type AgentToolCall = Extract; + +/** + * Result returned from `beforeToolCall`. + * + * Returning `{ block: true }` prevents the tool from executing. The loop emits an error tool result instead. + * `reason` becomes the text shown in that error result. If omitted, a default blocked message is used. + */ +export interface BeforeToolCallResult { + block?: boolean; + reason?: string; +} + +/** + * Partial override returned from `afterToolCall`. + * + * Merge semantics are field-by-field: + * - `content`: if provided, replaces the tool result content array in full + * - `details`: if provided, replaces the tool result details value in full + * - `isError`: if provided, replaces the tool result error flag + * - `terminate`: if provided, replaces the early-termination hint + * + * Omitted fields keep the original executed tool result values. + * There is no deep merge for `content` or `details`. + */ +export interface AfterToolCallResult { + content?: (TextContent | ImageContent)[]; + details?: unknown; + isError?: boolean; + /** + * Hint that the agent should stop after the current tool batch. + * Early termination only happens when every finalized tool result in the batch sets this to true. + */ + terminate?: boolean; +} + +/** Context passed to `beforeToolCall`. */ +export interface BeforeToolCallContext { + /** The assistant message that requested the tool call. */ + assistantMessage: AssistantMessage; + /** The raw tool call block from `assistantMessage.content`. */ + toolCall: AgentToolCall; + /** Validated tool arguments for the target tool schema. */ + args: unknown; + /** Current agent context at the time the tool call is prepared. */ + context: AgentContext; +} + +/** Context passed to `afterToolCall`. */ +export interface AfterToolCallContext { + /** The assistant message that requested the tool call. */ + assistantMessage: AssistantMessage; + /** The raw tool call block from `assistantMessage.content`. */ + toolCall: AgentToolCall; + /** Validated tool arguments for the target tool schema. */ + args: unknown; + /** The executed tool result before any `afterToolCall` overrides are applied. */ + result: AgentToolResult; + /** Whether the executed tool result is currently treated as an error. */ + isError: boolean; + /** Current agent context at the time the tool call is finalized. */ + context: AgentContext; +} + +/** Context passed to `shouldStopAfterTurn`. */ +export interface ShouldStopAfterTurnContext { + /** The assistant message that completed the turn. */ + message: AssistantMessage; + /** Tool result messages passed to the preceding `turn_end` event. */ + toolResults: ToolResultMessage[]; + /** Current agent context after the turn's assistant message and tool results have been appended. */ + context: AgentContext; + /** Messages that this loop invocation will return if it exits at this point. Prompt runs include the initial prompt messages; continuation runs do not include pre-existing context messages. */ + newMessages: AgentMessage[]; +} + +/** Replacement runtime state used by the agent loop before starting another provider request. */ +export interface AgentLoopTurnUpdate { + /** Context for the next provider request. */ + context?: AgentContext; + /** Model for the next provider request. */ + model?: Model; + /** Thinking level for the next provider request. */ + thinkingLevel?: ThinkingLevel; +} + +export interface PrepareNextTurnContext extends ShouldStopAfterTurnContext {} + +export interface AgentLoopConfig extends SimpleStreamOptions { + model: Model; + + /** + * Converts AgentMessage[] to LLM-compatible Message[] before each LLM call. + * + * Each AgentMessage must be converted to a UserMessage, AssistantMessage, or ToolResultMessage + * that the LLM can understand. AgentMessages that cannot be converted (e.g., UI-only notifications, + * status messages) should be filtered out. + * + * Contract: must not throw or reject. Return a safe fallback value instead. + * Throwing interrupts the low-level agent loop without producing a normal event sequence. + * + * @example + * ```typescript + * convertToLlm: (messages) => messages.flatMap(m => { + * if (m.role === "custom") { + * // Convert custom message to user message + * return [{ role: "user", content: m.content, timestamp: m.timestamp }]; + * } + * if (m.role === "notification") { + * // Filter out UI-only messages + * return []; + * } + * // Pass through standard LLM messages + * return [m]; + * }) + * ``` + */ + convertToLlm: (messages: AgentMessage[]) => Message[] | Promise; + + /** + * Optional transform applied to the context before `convertToLlm`. + * + * Use this for operations that work at the AgentMessage level: + * - Context window management (pruning old messages) + * - Injecting context from external sources + * + * Contract: must not throw or reject. Return the original messages or another + * safe fallback value instead. + * + * @example + * ```typescript + * transformContext: async (messages) => { + * if (estimateTokens(messages) > MAX_TOKENS) { + * return pruneOldMessages(messages); + * } + * return messages; + * } + * ``` + */ + transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise; + + /** + * Resolves an API key dynamically for each LLM call. + * + * Useful for short-lived OAuth tokens (e.g., GitHub Copilot) that may expire + * during long-running tool execution phases. + * + * Contract: must not throw or reject. Return undefined when no key is available. + */ + getApiKey?: (provider: string) => Promise | string | undefined; + + /** + * Called after each turn fully completes and `turn_end` has been emitted. + * + * If it returns true, the loop emits `agent_end` and exits before polling steering or follow-up queues, + * without starting another LLM call. The current assistant response and any tool executions finish normally. + * + * Use this to request a graceful stop after the current turn, e.g. before context gets too full. + * + * Contract: must not throw or reject. Throwing interrupts the low-level agent loop without producing a normal event sequence. + */ + shouldStopAfterTurn?: (context: ShouldStopAfterTurnContext) => boolean | Promise; + + /** + * Called after `turn_end` and before the loop decides whether another provider request should start. + * Return replacement context/model/thinking state to affect the next turn in this run. + * Return undefined to keep using the current context/config. + */ + prepareNextTurn?: ( + context: PrepareNextTurnContext, + ) => AgentLoopTurnUpdate | undefined | Promise; + + /** + * Returns steering messages to inject into the conversation mid-run. + * + * Called after the current assistant turn finishes executing its tool calls, unless `shouldStopAfterTurn` exits first. + * If messages are returned, they are added to the context before the next LLM call. + * Tool calls from the current assistant message are not skipped. + * + * Use this for "steering" the agent while it's working. + * + * Contract: must not throw or reject. Return [] when no steering messages are available. + */ + getSteeringMessages?: () => Promise; + + /** + * Returns follow-up messages to process after the agent would otherwise stop. + * + * Called when the agent has no more tool calls and no steering messages. + * If messages are returned, they're added to the context and the agent + * continues with another turn. + * + * Use this for follow-up messages that should wait until the agent finishes. + * + * Contract: must not throw or reject. Return [] when no follow-up messages are available. + */ + getFollowUpMessages?: () => Promise; + + /** + * Tool execution mode. + * - "sequential": execute tool calls one by one + * - "parallel": preflight tool calls sequentially, then execute allowed tools concurrently; + * emit `tool_execution_end` in tool completion order after each tool is finalized, + * then emit tool-result message artifacts later in assistant source order + * + * Default: "parallel" + */ + toolExecution?: ToolExecutionMode; + + /** + * Called before a tool is executed, after arguments have been validated. + * + * Return `{ block: true }` to prevent execution. The loop emits an error tool result instead. + * The hook receives the agent abort signal and is responsible for honoring it. + */ + beforeToolCall?: (context: BeforeToolCallContext, signal?: AbortSignal) => Promise; + + /** + * Called after a tool finishes executing, before `tool_execution_end` and tool-result message events are emitted. + * + * Return an `AfterToolCallResult` to override parts of the executed tool result: + * - `content` replaces the full content array + * - `details` replaces the full details payload + * - `isError` replaces the error flag + * - `terminate` replaces the early-termination hint + * + * Any omitted fields keep their original values. No deep merge is performed. + * The hook receives the agent abort signal and is responsible for honoring it. + */ + afterToolCall?: (context: AfterToolCallContext, signal?: AbortSignal) => Promise; +} + +/** + * Thinking/reasoning level for models that support it. + * Note: "xhigh" is only supported by selected model families. Use model thinking-level metadata + * from @earendil-works/pi-ai to detect support for a concrete model. + */ +export type ThinkingLevel = "off" | "minimal" | "low" | "medium" | "high" | "xhigh"; + +/** + * Extensible interface for custom app messages. + * Apps can extend via declaration merging: + * + * @example + * ```typescript + * declare module "@mariozechner/agent" { + * interface CustomAgentMessages { + * artifact: ArtifactMessage; + * notification: NotificationMessage; + * } + * } + * ``` + */ +export interface CustomAgentMessages { + // Empty by default - apps extend via declaration merging +} + +/** + * AgentMessage: Union of LLM messages + custom messages. + * This abstraction allows apps to add custom message types while maintaining + * type safety and compatibility with the base LLM messages. + */ +export type AgentMessage = Message | CustomAgentMessages[keyof CustomAgentMessages]; + +/** + * Public agent state. + * + * `tools` and `messages` use accessor properties so implementations can copy + * assigned arrays before storing them. + */ +export interface AgentState { + /** System prompt sent with each model request. */ + systemPrompt: string; + /** Active model used for future turns. */ + model: Model; + /** Requested reasoning level for future turns. */ + thinkingLevel: ThinkingLevel; + /** Available tools. Assigning a new array copies the top-level array. */ + set tools(tools: AgentTool[]); + get tools(): AgentTool[]; + /** Conversation transcript. Assigning a new array copies the top-level array. */ + set messages(messages: AgentMessage[]); + get messages(): AgentMessage[]; + /** + * True while the agent is processing a prompt or continuation. + * + * This remains true until awaited `agent_end` listeners settle. + */ + readonly isStreaming: boolean; + /** Partial assistant message for the current streamed response, if any. */ + readonly streamingMessage?: AgentMessage; + /** Tool call ids currently executing. */ + readonly pendingToolCalls: ReadonlySet; + /** Error message from the most recent failed or aborted assistant turn, if any. */ + readonly errorMessage?: string; +} + +/** Final or partial result produced by a tool. */ +export interface AgentToolResult { + /** Text or image content returned to the model. */ + content: (TextContent | ImageContent)[]; + /** Arbitrary structured details for logs or UI rendering. */ + details: T; + /** + * Hint that the agent should stop after the current tool batch. + * Early termination only happens when every finalized tool result in the batch sets this to true. + */ + terminate?: boolean; +} + +/** Callback used by tools to stream partial execution updates. */ +export type AgentToolUpdateCallback = (partialResult: AgentToolResult) => void; + +/** Tool definition used by the agent runtime. */ +export interface AgentTool extends Tool { + /** Human-readable label for UI display. */ + label: string; + /** + * Optional compatibility shim for raw tool-call arguments before schema validation. + * Must return an object that matches `TParameters`. + */ + prepareArguments?: (args: unknown) => Static; + /** Execute the tool call. Throw on failure instead of encoding errors in `content`. */ + execute: ( + toolCallId: string, + params: Static, + signal?: AbortSignal, + onUpdate?: AgentToolUpdateCallback, + ) => Promise>; + /** + * Per-tool execution mode override. + * - "sequential": this tool must execute one at a time with other tool calls. + * - "parallel": this tool can execute concurrently with other tool calls. + * + * If omitted, the default execution mode applies. + */ + executionMode?: ToolExecutionMode; +} + +/** Context snapshot passed into the low-level agent loop. */ +export interface AgentContext { + /** System prompt included with the request. */ + systemPrompt: string; + /** Transcript visible to the model. */ + messages: AgentMessage[]; + /** Tools available for this run. */ + tools?: AgentTool[]; +} + +/** + * Events emitted by the Agent for UI updates. + * + * `agent_end` is the last event emitted for a run, but awaited `Agent.subscribe()` + * listeners for that event are still part of run settlement. The agent becomes + * idle only after those listeners finish. + */ +export type AgentEvent = + // Agent lifecycle + | { type: "agent_start" } + | { type: "agent_end"; messages: AgentMessage[] } + // Turn lifecycle - a turn is one assistant response + any tool calls/results + | { type: "turn_start" } + | { type: "turn_end"; message: AgentMessage; toolResults: ToolResultMessage[] } + // Message lifecycle - emitted for user, assistant, and toolResult messages + | { type: "message_start"; message: AgentMessage } + // Only emitted for assistant messages during streaming + | { type: "message_update"; message: AgentMessage; assistantMessageEvent: AssistantMessageEvent } + | { type: "message_end"; message: AgentMessage } + // Tool execution lifecycle + | { type: "tool_execution_start"; toolCallId: string; toolName: string; args: any } + | { type: "tool_execution_update"; toolCallId: string; toolName: string; args: any; partialResult: any } + | { type: "tool_execution_end"; toolCallId: string; toolName: string; result: any; isError: boolean }; diff --git a/packages/agent/test/agent.test.ts b/packages/agent/test/agent.test.ts index e4460810..73683ca4 100644 --- a/packages/agent/test/agent.test.ts +++ b/packages/agent/test/agent.test.ts @@ -1,14 +1,54 @@ -import { Agent, type AgentTool } from "@earendil-works/pi-agent-core"; import { describe, expect, it } from "vitest"; +import { createAssistantMessageEventStream, type AssistantMessage } from "@earendil-works/pi-ai"; import { resolveCuaRuntimeSpec } from "@onkernel/cua-ai"; import type Kernel from "@onkernel/sdk"; -import { CuaAgent, CuaHarness, createCuaComputerTools, type KernelBrowser } from "../src/index"; +import { + Agent, + AgentHarness, + CuaAgent, + CuaAgentHarness, + InMemorySessionRepo, + NodeExecutionEnv, + createCuaComputerTools, + type AgentTool, + type KernelBrowser, + type StreamFn, +} from "../src/index"; const browser = { session_id: "browser_123" } as KernelBrowser; const client = {} as Kernel; +function createAssistantMessage(model: { api: string; provider: string; id: string }): AssistantMessage { + return { + role: "assistant", + content: [], + api: model.api, + provider: model.provider, + model: model.id, + 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(), + }; +} + +async function createHarnessServices() { + const sessionRepo = new InMemorySessionRepo(); + return { + env: new NodeExecutionEnv({ cwd: process.cwd() }), + session: await sessionRepo.create(), + }; +} + describe("CuaAgent", () => { it("extends pi Agent and resolves model refs in initialState", () => { + const runtime = resolveCuaRuntimeSpec("openai:gpt-5.5"); const agent = new CuaAgent({ browser, client, @@ -20,6 +60,7 @@ describe("CuaAgent", () => { expect(agent).toBeInstanceOf(Agent); expect(agent.state.model.id).toBe("gpt-5.5"); expect(agent.state.tools.length).toBeGreaterThan(0); + expect(agent.state.systemPrompt).toBe(runtime.defaultSystemPrompt); }); it("uses provided tools exactly", () => { @@ -57,7 +98,7 @@ describe("CuaAgent", () => { async execute() { return { content: [{ type: "text", text: "ok" }], details: {} }; }, - } satisfies AgentTool, + } satisfies AgentTool, ]; const agent = new CuaAgent({ @@ -73,29 +114,121 @@ describe("CuaAgent", () => { expect(agent.state.tools).toHaveLength(3); expect(agent.state.systemPrompt).toBe("Use the browser carefully."); }); + + it("refreshes CUA runtime state when state.model changes", () => { + const runtime = resolveCuaRuntimeSpec("google:gemini-3-pro-preview"); + const agent = new CuaAgent({ + browser, + client, + initialState: { + model: "openai:gpt-5.5", + }, + }); + + agent.state.model = "google:gemini-3-pro-preview"; + + expect(agent.state.model.id).toBe(runtime.model.id); + expect(agent.state.systemPrompt).toBe(runtime.defaultSystemPrompt); + expect(agent.state.tools).toHaveLength(runtime.toolDefinitions.length); + }); + + it("keeps caller-owned tools and system prompt when state.model changes", () => { + const tool: AgentTool = { + name: "custom", + label: "custom", + description: "custom tool", + parameters: { type: "object", properties: {}, additionalProperties: false } as never, + async execute() { + return { content: [{ type: "text", text: "ok" }], details: {} }; + }, + }; + const agent = new CuaAgent({ + browser, + client, + initialState: { + model: "openai:gpt-5.5", + tools: [tool], + systemPrompt: "custom prompt", + }, + }); + + agent.state.model = "google:gemini-3-pro-preview"; + + expect(agent.state.tools).toEqual([tool]); + expect(agent.state.systemPrompt).toBe("custom prompt"); + }); + + it("composes payload hooks for custom stream functions", async () => { + const payloads: unknown[] = []; + const streamFn: StreamFn = (model, _context, options) => { + const stream = createAssistantMessageEventStream(); + void (async () => { + payloads.push(await options?.onPayload?.({ provider: model.provider }, model)); + const message = createAssistantMessage(model); + stream.push({ type: "start", partial: message }); + stream.push({ type: "done", reason: "stop", message }); + stream.end(message); + })(); + return stream; + }; + + const agent = new CuaAgent({ + browser, + client, + streamFn, + onPayload: (payload) => ({ payload, userHook: true }), + initialState: { + model: "openai:gpt-5.5", + }, + }); + + await agent.prompt("hello"); + + expect(payloads).toEqual([{ payload: { provider: "openai", store: true }, userHook: true }]); + }); }); -describe("CuaHarness", () => { - it("wraps a pi Agent and resolves model refs", () => { - const harness = new CuaHarness({ +describe("CuaAgentHarness", () => { + it("extends pi AgentHarness and resolves model refs", async () => { + const harness = new CuaAgentHarness({ + ...(await createHarnessServices()), browser, client, model: "openai:gpt-5.5", - getApiKey: () => "test-key", + getApiKeyAndHeaders: async () => ({ apiKey: "test-key" }), }); + expect(harness).toBeInstanceOf(AgentHarness); expect(harness.agent).toBeInstanceOf(Agent); expect(harness.agent.state.model.id).toBe("gpt-5.5"); expect(harness.agent.state.tools.length).toBeGreaterThan(0); }); - it("exposes transcript snapshots", () => { - const harness = new CuaHarness({ + it("refreshes CUA runtime state through setModel", async () => { + const runtime = resolveCuaRuntimeSpec("google:gemini-3-pro-preview"); + const harness = new CuaAgentHarness({ + ...(await createHarnessServices()), + browser, + client, + model: "openai:gpt-5.5", + }); + + await harness.setModel("google:gemini-3-pro-preview"); + + expect(harness.agent.state.model.id).toBe(runtime.model.id); + expect(harness.agent.state.tools).toHaveLength(runtime.toolDefinitions.length); + }); + + it("preserves active tool selection when setModel refreshes tools", async () => { + const harness = new CuaAgentHarness({ + ...(await createHarnessServices()), browser, client, model: "openai:gpt-5.5", }); - const transcript = harness.getTranscript(); - expect(transcript).toEqual(harness.state.messages); - expect(transcript).not.toBe(harness.state.messages); + + await harness.setActiveTools([]); + await harness.setModel("google:gemini-3-pro-preview"); + + expect(harness.agent.state.tools).toEqual([]); }); }); diff --git a/packages/agent/test/e2e.live.test.ts b/packages/agent/test/e2e.live.test.ts index cba00632..06f97e0f 100644 --- a/packages/agent/test/e2e.live.test.ts +++ b/packages/agent/test/e2e.live.test.ts @@ -1,7 +1,14 @@ import Kernel from "@onkernel/sdk"; -import type { AgentMessage } from "@earendil-works/pi-agent-core"; import { describe, expect, it } from "vitest"; -import { CuaAgent, CuaHarness } from "../src/index"; +import { + CuaAgent, + CuaAgentHarness, + InMemorySessionRepo, + NodeExecutionEnv, + type AgentEvent, + type AgentHarnessEvent, + type AgentMessage, +} from "../src/index"; const LIVE = process.env.CUA_E2E_LIVE === "1"; const KERNEL_API_KEY = process.env.KERNEL_API_KEY; @@ -20,6 +27,13 @@ type ProviderCase = { timeoutMs: number; }; +type ModelSwitchCase = { + name: string; + from: ProviderCase; + to: ProviderCase; + timeoutMs: number; +}; + const cases: ProviderCase[] = [ { name: "openai", @@ -88,6 +102,15 @@ const cases: ProviderCase[] = [ }, ]; +const switchCases: ModelSwitchCase[] = [ + { + name: "openai-to-gemini", + from: cases[0]!, + to: cases[2]!, + timeoutMs: 420_000, + }, +]; + type RunStats = { toolCalls: number; toolResults: number; @@ -103,6 +126,14 @@ function shouldRunCase(c: ProviderCase): boolean { return Boolean(process.env[c.apiKeyEnvVar]); } +function shouldRunSwitchCase(c: ModelSwitchCase): boolean { + return shouldRunCase(c.from) && shouldRunCase(c.to); +} + +function createRunStats(): RunStats { + return { toolCalls: 0, toolResults: 0, hasReadArtifact: false, toolErrors: [], assistantErrors: [] }; +} + async function withBrowser(run: (client: Kernel, browser: Awaited>) => Promise): Promise { if (!KERNEL_API_KEY) { throw new Error("KERNEL_API_KEY is required"); @@ -116,6 +147,14 @@ async function withBrowser(run: (client: Kernel, browser: Awaited block.type === "image" || (block.type === "text" && /url\(\)|Current URL:/.test(block.text)), + ) + ) { + stats.hasReadArtifact = true; + } + } + if (event.type === "message_end" && event.message.role === "assistant") { + stats.finalAssistant = event.message; + if (event.message.errorMessage) { + stats.assistantErrors.push(event.message.errorMessage); + } + } +} + describe("Cua live e2e", () => { for (const c of cases) { const test = shouldRunCase(c) ? it : it.skip; @@ -139,7 +201,7 @@ describe("Cua live e2e", () => { `${c.name}: CuaAgent executes browser steps`, async () => { await withBrowser(async (client, browser) => { - const stats: RunStats = { toolCalls: 0, toolResults: 0, hasReadArtifact: false, toolErrors: [], assistantErrors: [] }; + const stats = createRunStats(); const agent = new CuaAgent({ browser, client, @@ -149,26 +211,7 @@ describe("Cua live e2e", () => { }, }); agent.subscribe((event) => { - if (event.type === "tool_execution_start") stats.toolCalls += 1; - if (event.type === "tool_execution_end" && event.isError) { - stats.toolErrors.push(`${event.toolName}: failed`); - } - if (event.type === "message_end" && event.message.role === "toolResult") { - stats.toolResults += 1; - if ( - event.message.content.some( - (block) => block.type === "image" || (block.type === "text" && /url\(\)|Current URL:/.test(block.text)), - ) - ) { - stats.hasReadArtifact = true; - } - } - if (event.type === "message_end" && event.message.role === "assistant") { - stats.finalAssistant = event.message; - if (event.message.errorMessage) { - stats.assistantErrors.push(event.message.errorMessage); - } - } + recordRunEvent(stats, event); }); await agent.prompt(c.prompt); @@ -179,38 +222,23 @@ describe("Cua live e2e", () => { ); test( - `${c.name}: CuaHarness executes browser steps`, + `${c.name}: CuaAgentHarness executes browser steps`, async () => { await withBrowser(async (client, browser) => { - const stats: RunStats = { toolCalls: 0, toolResults: 0, hasReadArtifact: false, toolErrors: [], assistantErrors: [] }; - const harness = new CuaHarness({ + const stats = createRunStats(); + const harness = new CuaAgentHarness({ + ...(await createHarnessServices(`${c.name}-harness`)), browser, client, model: c.modelRef, - getApiKey: () => process.env[c.apiKeyEnvVar], + getApiKeyAndHeaders: async () => { + const apiKey = process.env[c.apiKeyEnvVar]; + return apiKey ? { apiKey } : undefined; + }, }); harness.subscribe((event) => { - if (event.type === "tool_execution_start") stats.toolCalls += 1; - if (event.type === "tool_execution_end" && event.isError) { - stats.toolErrors.push(`${event.toolName}: failed`); - } - if (event.type === "message_end" && event.message.role === "toolResult") { - stats.toolResults += 1; - if ( - event.message.content.some( - (block) => block.type === "image" || (block.type === "text" && /url\(\)|Current URL:/.test(block.text)), - ) - ) { - stats.hasReadArtifact = true; - } - } - if (event.type === "message_end" && event.message.role === "assistant") { - stats.finalAssistant = event.message; - if (event.message.errorMessage) { - stats.assistantErrors.push(event.message.errorMessage); - } - } + recordRunEvent(stats, event); }); await harness.prompt(c.prompt); @@ -220,4 +248,79 @@ describe("Cua live e2e", () => { c.timeoutMs, ); } + + for (const c of switchCases) { + const test = shouldRunSwitchCase(c) ? it : it.skip; + + test( + `${c.name}: CuaAgent switches models after a turn`, + async () => { + await withBrowser(async (client, browser) => { + let stats = createRunStats(); + const agent = new CuaAgent({ + browser, + client, + getApiKey: (provider) => { + if (provider === c.from.modelRef.split(":")[0]) return process.env[c.from.apiKeyEnvVar]; + if (provider === c.to.modelRef.split(":")[0]) return process.env[c.to.apiKeyEnvVar]; + return undefined; + }, + initialState: { + model: c.from.modelRef, + }, + }); + agent.subscribe((event) => { + recordRunEvent(stats, event); + }); + + await agent.prompt(c.from.prompt); + assertStats(stats, c.from.expectToolCalls, c.from.name, "agent"); + + stats = createRunStats(); + agent.state.model = c.to.modelRef; + await agent.prompt(c.to.prompt); + assertStats(stats, c.to.expectToolCalls, c.to.name, "agent"); + }); + }, + c.timeoutMs, + ); + + test( + `${c.name}: CuaAgentHarness switches models after a turn`, + async () => { + await withBrowser(async (client, browser) => { + let stats = createRunStats(); + const harness = new CuaAgentHarness({ + ...(await createHarnessServices(`${c.name}-harness-switch`)), + browser, + client, + model: c.from.modelRef, + getApiKeyAndHeaders: async (model) => { + if (model.provider === c.from.modelRef.split(":")[0]) { + const apiKey = process.env[c.from.apiKeyEnvVar]; + return apiKey ? { apiKey } : undefined; + } + if (model.provider === c.to.modelRef.split(":")[0]) { + const apiKey = process.env[c.to.apiKeyEnvVar]; + return apiKey ? { apiKey } : undefined; + } + return undefined; + }, + }); + harness.subscribe((event) => { + recordRunEvent(stats, event); + }); + + await harness.prompt(c.from.prompt); + assertStats(stats, c.from.expectToolCalls, c.from.name, "harness"); + + stats = createRunStats(); + await harness.setModel(c.to.modelRef); + await harness.prompt(c.to.prompt); + assertStats(stats, c.to.expectToolCalls, c.to.name, "harness"); + }); + }, + c.timeoutMs, + ); + } }); diff --git a/packages/agent/test/translator.test.ts b/packages/agent/test/translator.test.ts new file mode 100644 index 00000000..77bf5efb --- /dev/null +++ b/packages/agent/test/translator.test.ts @@ -0,0 +1,59 @@ +import { describe, expect, it } from "vitest"; +import type Kernel from "@onkernel/sdk"; +import { InternalComputerTranslator, type KernelBrowser } from "../src/translator/translator"; + +const browser = { session_id: "browser_123" } as KernelBrowser; + +function createClient() { + const batches: unknown[] = []; + const client = { + browsers: { + computer: { + batch: async (_id: string, body: { actions: unknown[] }) => { + batches.push(body.actions); + }, + readClipboard: async () => ({ text: "https://example.com/" }), + }, + }, + } as unknown as Kernel; + return { batches, client }; +} + +describe("InternalComputerTranslator", () => { + it("holds modifiers instead of typing shortcut keys", async () => { + const { batches, client } = createClient(); + const translator = new InternalComputerTranslator({ browser, client }); + + await translator.executeBatch([ + { type: "goto", url: "https://example.com" }, + { type: "url" }, + { type: "keypress", keys: ["ctrl", "shift", "Tab"] }, + ]); + + expect(batches).toEqual([ + [ + { type: "press_key", press_key: { keys: ["l"], hold_keys: ["Control_L"] } }, + { type: "type_text", type_text: { text: "https://example.com" } }, + { type: "press_key", press_key: { keys: ["Return"] } }, + ], + [ + { type: "press_key", press_key: { keys: ["l"], hold_keys: ["Control_L"] } }, + { type: "press_key", press_key: { keys: ["c"], hold_keys: ["Control_L"] } }, + ], + [ + { type: "press_key", press_key: { keys: ["Tab"], hold_keys: ["Control_L", "Shift_L"] } }, + ], + ]); + }); + + it("accepts shortcut strings from provider adapters", async () => { + const { batches, client } = createClient(); + const translator = new InternalComputerTranslator({ browser, client }); + + await translator.executeBatch([{ type: "keypress", keys: ["Ctrl+L"] }]); + + expect(batches).toEqual([ + [{ type: "press_key", press_key: { keys: ["l"], hold_keys: ["Control_L"] } }], + ]); + }); +});