From f1c587dde39025c75d7397bc14532d8fa5c001d9 Mon Sep 17 00:00:00 2001 From: Armin Ronacher Date: Sun, 19 Jul 2026 20:18:49 +0200 Subject: [PATCH 01/62] fix(coding-agent): avoid duplicate session reads fixes #6793 --- packages/coding-agent/CHANGELOG.md | 1 + .../coding-agent/src/core/session-manager.ts | 112 +++++++++++++++--- .../session-manager/file-operations.test.ts | 61 ++++++++++ 3 files changed, 156 insertions(+), 18 deletions(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index e3ea47011..15bdbf80d 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -9,6 +9,7 @@ ### Fixed +- Fixed persisted sessions being read and parsed twice when opened, reducing startup latency for large sessions ([#6793](https://github.com/earendil-works/pi/issues/6793)). - Fixed prompt-template defaults for all arguments (`${@:-default}` and `${ARGUMENTS:-default}`) ([#6695](https://github.com/earendil-works/pi/issues/6695)). - Fixed obsolete custom UI, custom tool, and custom editor examples in the extension documentation ([#6735](https://github.com/earendil-works/pi/issues/6735)). - Fixed Kimi Coding sessions to show API-equivalent implied costs with the subscription indicator. diff --git a/packages/coding-agent/src/core/session-manager.ts b/packages/coding-agent/src/core/session-manager.ts index eadca30b1..f5b75c2cb 100644 --- a/packages/coding-agent/src/core/session-manager.ts +++ b/packages/coding-agent/src/core/session-manager.ts @@ -485,6 +485,16 @@ export function getDefaultSessionDir(cwd: string, agentDir: string = getDefaultA } const SESSION_READ_BUFFER_SIZE = 1024 * 1024; +const SESSION_HEADER_READ_BUFFER_SIZE = 4096; +/** Bound synchronous header discovery while allowing large cwd and custom metadata fields. */ +const MAX_SESSION_HEADER_SCAN_BYTES = 1024 * 1024; + +class SessionHeaderScanLimitError extends Error { + constructor(filePath: string) { + super(`Session header exceeds ${MAX_SESSION_HEADER_SCAN_BYTES}-byte scan limit: ${filePath}`); + this.name = "SessionHeaderScanLimitError"; + } +} function parseSessionEntryLine(line: string): FileEntry | null { if (!line.trim()) return null; @@ -541,20 +551,69 @@ export function loadEntriesFromFile(filePath: string): FileEntry[] { return entries; } +/** + * Inspect a physical line while searching for the first parsed session entry. + * Blank and malformed lines are skipped to match loadEntriesFromFile(). + * Returns undefined to keep scanning, null for a parsed non-header entry, or the header. + */ +function parseSessionHeaderCandidate(line: string): SessionHeader | null | undefined { + if (!line.trim()) return undefined; + const entry = parseSessionEntryLine(line); + if (!entry) return undefined; + if (entry.type !== "session" || typeof (entry as { id?: unknown }).id !== "string") return null; + return entry; +} + function readSessionHeader(filePath: string): SessionHeader | null { + const fd = openSync(filePath, "r"); try { - const fd = openSync(filePath, "r"); - const buffer = Buffer.alloc(512); - const bytesRead = readSync(fd, buffer, 0, 512, 0); - closeSync(fd); - const firstLine = buffer.toString("utf8", 0, bytesRead).split("\n")[0]; - if (!firstLine) return null; - const header = JSON.parse(firstLine) as Record; - if (header.type !== "session" || typeof header.id !== "string") { - return null; + const decoder = new StringDecoder("utf8"); + const buffer = Buffer.allocUnsafe(SESSION_HEADER_READ_BUFFER_SIZE); + const lineChunks: string[] = []; + let scannedBytes = 0; + + while (scannedBytes < MAX_SESSION_HEADER_SCAN_BYTES) { + const readLength = Math.min(buffer.length, MAX_SESSION_HEADER_SCAN_BYTES - scannedBytes); + const bytesRead = readSync(fd, buffer, 0, readLength, null); + if (bytesRead === 0) { + lineChunks.push(decoder.end()); + return parseSessionHeaderCandidate(lineChunks.join("")) ?? null; + } + scannedBytes += bytesRead; + + const chunk = decoder.write(buffer.subarray(0, bytesRead)); + let lineStart = 0; + let newlineIndex = chunk.indexOf("\n", lineStart); + while (newlineIndex !== -1) { + lineChunks.push(chunk.slice(lineStart, newlineIndex)); + const header = parseSessionHeaderCandidate(lineChunks.join("")); + if (header !== undefined) return header; + lineChunks.length = 0; + lineStart = newlineIndex + 1; + newlineIndex = chunk.indexOf("\n", lineStart); + } + lineChunks.push(chunk.slice(lineStart)); } - return header as unknown as SessionHeader; + + // Probe for EOF so a final header without a newline is allowed when it ends + // exactly at the scan limit. Any additional byte exceeds the bounded scan. + const probe = Buffer.allocUnsafe(1); + if (readSync(fd, probe, 0, probe.length, null) === 0) { + lineChunks.push(decoder.end()); + return parseSessionHeaderCandidate(lineChunks.join("")) ?? null; + } + throw new SessionHeaderScanLimitError(filePath); + } finally { + closeSync(fd); + } +} + +function readSessionHeaderForDiscovery(filePath: string): SessionHeader | null { + try { + return readSessionHeader(filePath); } catch { + // Discovery is best-effort: unreadable or oversized files are not sessions, + // and one corrupt file must not prevent other sessions from being found. return null; } } @@ -576,7 +635,7 @@ export function findMostRecentSession(sessionDir: string, cwd?: string): string const files = readdirSync(resolvedSessionDir) .filter((f) => f.endsWith(".jsonl")) .map((f) => join(resolvedSessionDir, f)) - .map((path) => ({ path, header: readSessionHeader(path) })) + .map((path) => ({ path, header: readSessionHeaderForDiscovery(path) })) .filter( (file): file is { path: string; header: SessionHeader } => file.header !== null && @@ -587,6 +646,7 @@ export function findMostRecentSession(sessionDir: string, cwd?: string): string return files[0]?.path || null; } catch { + // Directory access and stat races make recent-session discovery unavailable. return null; } } @@ -807,6 +867,7 @@ export class SessionManager { sessionFile: string | undefined, persist: boolean, newSessionOptions?: NewSessionOptions, + preloadedFileEntries?: FileEntry[], ) { this.cwd = resolvePath(cwd); this.sessionDir = normalizePath(sessionDir); @@ -816,7 +877,7 @@ export class SessionManager { } if (sessionFile) { - this.setSessionFile(sessionFile); + this._setSessionFile(sessionFile, preloadedFileEntries); } else { this.newSession(newSessionOptions); } @@ -824,9 +885,13 @@ export class SessionManager { /** Switch to a different session file (used for resume and branching) */ setSessionFile(sessionFile: string): void { + this._setSessionFile(sessionFile); + } + + private _setSessionFile(sessionFile: string, preloadedFileEntries?: FileEntry[]): void { this.sessionFile = resolvePath(sessionFile); if (existsSync(this.sessionFile)) { - this.fileEntries = loadEntriesFromFile(this.sessionFile); + this.fileEntries = preloadedFileEntries ?? loadEntriesFromFile(this.sessionFile); // If file was empty, initialize it with a valid session header. If it was // non-empty but did not parse as a pi session, fail without modifying it. @@ -1451,13 +1516,24 @@ export class SessionManager { */ static open(path: string, sessionDir?: string, cwdOverride?: string): SessionManager { const resolvedPath = resolvePath(path); - // Extract cwd from session header if possible, otherwise use process.cwd() - const entries = loadEntriesFromFile(resolvedPath); - const header = entries.find((e) => e.type === "session") as SessionHeader | undefined; - const cwd = cwdOverride ?? header?.cwd ?? process.cwd(); + let header: SessionHeader | null = null; + let preloadedFileEntries: FileEntry[] | undefined; + if (cwdOverride === undefined && existsSync(resolvedPath)) { + try { + header = readSessionHeader(resolvedPath); + } catch (error) { + if (!(error instanceof SessionHeaderScanLimitError)) throw error; + // The bounded scan is only a discovery optimization. A full load remains + // authoritative for legacy files with very large headers or prefixes. + preloadedFileEntries = loadEntriesFromFile(resolvedPath); + const firstEntry = preloadedFileEntries[0]; + header = firstEntry?.type === "session" ? firstEntry : null; + } + } + const cwd = cwdOverride ?? (header ? getSessionHeaderCwd(header) : undefined) ?? process.cwd(); // If no sessionDir provided, derive from file's parent directory const dir = sessionDir ? normalizePath(sessionDir) : resolve(resolvedPath, ".."); - return new SessionManager(cwd, dir, resolvedPath, true); + return new SessionManager(cwd, dir, resolvedPath, true, undefined, preloadedFileEntries); } /** diff --git a/packages/coding-agent/test/session-manager/file-operations.test.ts b/packages/coding-agent/test/session-manager/file-operations.test.ts index d4ec0d2f9..c978b3200 100644 --- a/packages/coding-agent/test/session-manager/file-operations.test.ts +++ b/packages/coding-agent/test/session-manager/file-operations.test.ts @@ -5,6 +5,8 @@ import { join } from "path"; import { afterEach, beforeEach, describe, expect, it } from "vitest"; import { findMostRecentSession, loadEntriesFromFile, SessionManager } from "../../src/core/session-manager.ts"; +const HEADER_SCAN_LIMIT_BYTES = 1024 * 1024; + describe("loadEntriesFromFile", () => { let tempDir: string; @@ -17,6 +19,19 @@ describe("loadEntriesFromFile", () => { rmSync(tempDir, { recursive: true, force: true }); }); + function writeSessionHeader(file: string, cwd: string, id: string, prefix = ""): void { + writeFileSync( + file, + `${prefix}${JSON.stringify({ + type: "session", + version: 3, + id, + timestamp: "2025-01-01T00:00:00Z", + cwd, + })}\n`, + ); + } + it("returns empty array for non-existent file", () => { const entries = loadEntriesFromFile(join(tempDir, "nonexistent.jsonl")); expect(entries).toEqual([]); @@ -65,6 +80,43 @@ describe("loadEntriesFromFile", () => { expect(entries).toHaveLength(2); }); + it.each([ + ["leading blank lines", "\n \n", "leading-blank"], + ["leading malformed lines", "not json\n{broken json\n", "leading-malformed"], + ["a multi-buffer header", "", "a".repeat(8192)], + ])("reads cwd from a session with %s", (_description, prefix, sessionId) => { + const file = join(tempDir, "header.jsonl"); + const storedCwd = join(tempDir, "stored-project"); + writeSessionHeader(file, storedCwd, sessionId, prefix); + + const sessionManager = SessionManager.open(file, tempDir); + expect(sessionManager.getSessionId()).toBe(sessionId); + expect(sessionManager.getCwd()).toBe(storedCwd); + }); + + it("opens compatible sessions beyond the discovery scan limit", () => { + const storedCwd = join(tempDir, "stored-project"); + const overrideCwd = join(tempDir, "override-project"); + const cases = [ + { name: "large-header", id: "a".repeat(HEADER_SCAN_LIMIT_BYTES + 1), prefix: "" }, + { + name: "large-prefix", + id: "large-prefix", + prefix: `${"x".repeat(HEADER_SCAN_LIMIT_BYTES + 1)}\n`, + }, + ]; + + for (const { name, id, prefix } of cases) { + const file = join(tempDir, `${name}.jsonl`); + writeSessionHeader(file, storedCwd, id, prefix); + for (const cwdOverride of [undefined, overrideCwd]) { + const sessionManager = SessionManager.open(file, tempDir, cwdOverride); + expect(sessionManager.getSessionId()).toBe(id); + expect(sessionManager.getCwd()).toBe(cwdOverride ?? storedCwd); + } + } + }); + it("opens session files larger than Node's max string length", () => { const file = join(tempDir, "large.jsonl"); writeFileSync( @@ -155,6 +207,15 @@ describe("findMostRecentSession", () => { expect(findMostRecentSession(tempDir)).toBe(valid); }); + it("skips oversized corrupt files and returns a valid session", () => { + const invalid = join(tempDir, "oversized.jsonl"); + const valid = join(tempDir, "valid.jsonl"); + writeFileSync(invalid, "x".repeat(HEADER_SCAN_LIMIT_BYTES + 1)); + writeFileSync(valid, '{"type":"session","id":"abc","timestamp":"2025-01-01T00:00:00Z","cwd":"/tmp"}\n'); + + expect(findMostRecentSession(tempDir)).toBe(valid); + }); + it("filters most recent session by cwd", async () => { const projectA = join(tempDir, "project-a"); const projectB = join(tempDir, "project-b"); From 956074697fc0f59e8a9787717e76d5bb0ee7bc72 Mon Sep 17 00:00:00 2001 From: Armin Ronacher Date: Sun, 19 Jul 2026 22:21:31 +0200 Subject: [PATCH 02/62] fix(ai): replace generic record checks --- packages/ai/src/api/pi-messages.ts | 9 +++---- packages/ai/src/providers/radius-config.ts | 31 +++++++++++----------- scripts/publish-model-catalog.mjs | 31 +++++++++++++++------- 3 files changed, 41 insertions(+), 30 deletions(-) diff --git a/packages/ai/src/api/pi-messages.ts b/packages/ai/src/api/pi-messages.ts index 0af13f88f..ee83f4ec6 100644 --- a/packages/ai/src/api/pi-messages.ts +++ b/packages/ai/src/api/pi-messages.ts @@ -103,14 +103,11 @@ export class PiMessagesResponseError extends Error { } } -function isRecord(value: unknown): value is Record { - return typeof value === "object" && value !== null && !Array.isArray(value); -} - function parsePiMessagesErrorBody(body: string): PiMessagesErrorBody | undefined { try { - const parsed = JSON.parse(body) as unknown; - return isRecord(parsed) && isRecord(parsed.error) ? (parsed as PiMessagesErrorBody) : undefined; + const parsed = JSON.parse(body) as PiMessagesErrorBody | null; + const error = parsed?.error; + return parsed && typeof error === "object" && error !== null && !Array.isArray(error) ? parsed : undefined; } catch { return undefined; } diff --git a/packages/ai/src/providers/radius-config.ts b/packages/ai/src/providers/radius-config.ts index dc96a300f..a2715aaec 100644 --- a/packages/ai/src/providers/radius-config.ts +++ b/packages/ai/src/providers/radius-config.ts @@ -23,28 +23,29 @@ export type RadiusOAuthCredential = OAuthCredential & { gatewayConfig?: RadiusGatewayConfig; }; -function isRecord(value: unknown): value is Record { - return typeof value === "object" && value !== null && !Array.isArray(value); -} - function isRadiusGatewayModel(value: unknown): value is RadiusGatewayModel { + if (typeof value !== "object" || value === null || Array.isArray(value)) return false; + const model = value as Partial; return ( - isRecord(value) && - typeof value.id === "string" && - typeof value.name === "string" && - typeof value.reasoning === "boolean" && - Array.isArray(value.input) && - isRecord(value.cost) && - typeof value.contextWindow === "number" && - typeof value.maxTokens === "number" + typeof model.id === "string" && + typeof model.name === "string" && + typeof model.reasoning === "boolean" && + Array.isArray(model.input) && + typeof model.cost === "object" && + model.cost !== null && + !Array.isArray(model.cost) && + typeof model.contextWindow === "number" && + typeof model.maxTokens === "number" ); } function sanitizeRadiusGatewayConfig(config: unknown): RadiusGatewayConfig | undefined { - if (!isRecord(config) || typeof config.baseUrl !== "string" || !Array.isArray(config.models)) return undefined; + if (typeof config !== "object" || config === null || Array.isArray(config)) return undefined; + const { baseUrl, models } = config as Partial; + if (typeof baseUrl !== "string" || !Array.isArray(models)) return undefined; return { - baseUrl: config.baseUrl, - models: config.models.filter(isRadiusGatewayModel).map((model) => ({ ...model })), + baseUrl, + models: models.filter(isRadiusGatewayModel).map((model) => ({ ...model })), }; } diff --git a/scripts/publish-model-catalog.mjs b/scripts/publish-model-catalog.mjs index a8a30d43f..922e19ebe 100644 --- a/scripts/publish-model-catalog.mjs +++ b/scripts/publish-model-catalog.mjs @@ -65,10 +65,6 @@ function readJson(path) { return JSON.parse(readFileSync(path, "utf8")); } -function isRecord(value) { - return typeof value === "object" && value !== null && !Array.isArray(value); -} - function validateBundle(inputDir) { const modelsPath = join(inputDir, "models.json"); const providerIndexPath = join(inputDir, "providers.json"); @@ -77,7 +73,9 @@ function validateBundle(inputDir) { const models = JSON.parse(modelsBytes.toString("utf8")); const providerIds = readJson(providerIndexPath); - if (!isRecord(models)) throw new Error("models.json must contain an object"); + if (typeof models !== "object" || models === null || Array.isArray(models)) { + throw new Error("models.json must contain an object"); + } if (!Array.isArray(providerIds) || !providerIds.every((value) => typeof value === "string")) { throw new Error("providers.json must contain an array of provider IDs"); } @@ -93,13 +91,21 @@ function validateBundle(inputDir) { let modelCount = 0; for (const providerId of providerIds) { const providerModels = models[providerId]; - if (!isRecord(providerModels)) throw new Error(`Provider catalog must be an object: ${providerId}`); + if (typeof providerModels !== "object" || providerModels === null || Array.isArray(providerModels)) { + throw new Error(`Provider catalog must be an object: ${providerId}`); + } const providerFile = readJson(join(providersDir, `${providerId}.json`)); if (!isDeepStrictEqual(providerFile, providerModels)) { throw new Error(`Provider shard does not match models.json: ${providerId}`); } for (const [modelId, model] of Object.entries(providerModels)) { - if (!isRecord(model) || model.id !== modelId || model.provider !== providerId) { + if ( + typeof model !== "object" || + model === null || + Array.isArray(model) || + model.id !== modelId || + model.provider !== providerId + ) { throw new Error(`Invalid model entry: ${providerId}/${modelId}`); } modelCount++; @@ -181,13 +187,20 @@ function uploadJson(bucket, endpoint, sourcePath, key, cacheControl) { } function validateIndex(index) { - if (!isRecord(index) || index.schemaVersion !== CATALOG_SCHEMA_VERSION) { + if ( + typeof index !== "object" || + index === null || + Array.isArray(index) || + index.schemaVersion !== CATALOG_SCHEMA_VERSION + ) { throw new Error(`Existing ${CATALOG_INDEX_KEY} has an unsupported schema`); } if (!Array.isArray(index.catalogs)) throw new Error(`Existing ${CATALOG_INDEX_KEY} has no catalogs array`); for (const catalog of index.catalogs) { if ( - !isRecord(catalog) || + typeof catalog !== "object" || + catalog === null || + Array.isArray(catalog) || typeof catalog.minimumPiVersion !== "string" || typeof catalog.revision !== "string" ) { From d2f8dafb0f07409758797c880fbc3d526fa7c5c6 Mon Sep 17 00:00:00 2001 From: Alexey Zaytsev Date: Sun, 19 Jul 2026 20:23:47 -0300 Subject: [PATCH 03/62] fix(ai,agent,coding-agent): share UUIDv7 and use for Codex (#6834) --- packages/agent/CHANGELOG.md | 4 ++++ .../agent/src/harness/session/jsonl-storage.ts | 2 +- .../agent/src/harness/session/memory-storage.ts | 2 +- packages/agent/src/harness/session/repo-utils.ts | 2 +- packages/agent/src/index.ts | 1 - packages/ai/CHANGELOG.md | 5 +++++ packages/ai/src/api/openai-codex-responses.ts | 10 ++-------- packages/ai/src/index.ts | 1 + .../src/harness/session => ai/src/utils}/uuid.ts | 16 +++++----------- .../test/uuid.test.ts} | 2 +- .../coding-agent/src/core/session-manager.ts | 4 ++-- 11 files changed, 23 insertions(+), 26 deletions(-) rename packages/{agent/src/harness/session => ai/src/utils}/uuid.ts (83%) rename packages/{agent/test/harness/session-uuid.test.ts => ai/test/uuid.test.ts} (96%) diff --git a/packages/agent/CHANGELOG.md b/packages/agent/CHANGELOG.md index b27954ca9..eba2a62a8 100644 --- a/packages/agent/CHANGELOG.md +++ b/packages/agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Breaking Changes + +- Moved the `uuidv7` export to `@earendil-works/pi-ai`. + ## [0.80.10] - 2026-07-16 ## [0.80.9] - 2026-07-16 diff --git a/packages/agent/src/harness/session/jsonl-storage.ts b/packages/agent/src/harness/session/jsonl-storage.ts index b41115726..d57398d53 100644 --- a/packages/agent/src/harness/session/jsonl-storage.ts +++ b/packages/agent/src/harness/session/jsonl-storage.ts @@ -1,7 +1,7 @@ +import { uuidv7 } from "@earendil-works/pi-ai"; import type { FileSystem, JsonlSessionMetadata, LeafEntry, SessionStorage, SessionTreeEntry } from "../types.ts"; import { SessionError, toError } from "../types.ts"; import { getFileSystemResultOrThrow } from "./repo-utils.ts"; -import { uuidv7 } from "./uuid.ts"; type JsonlSessionStorageFileSystem = Pick; diff --git a/packages/agent/src/harness/session/memory-storage.ts b/packages/agent/src/harness/session/memory-storage.ts index fba55dae5..68ec76231 100644 --- a/packages/agent/src/harness/session/memory-storage.ts +++ b/packages/agent/src/harness/session/memory-storage.ts @@ -1,3 +1,4 @@ +import { uuidv7 } from "@earendil-works/pi-ai"; import { type LeafEntry, SessionError, @@ -5,7 +6,6 @@ import { type SessionStorage, type SessionTreeEntry, } from "../types.ts"; -import { uuidv7 } from "./uuid.ts"; function updateLabelCache(labelsById: Map, entry: SessionTreeEntry): void { if (entry.type !== "label") return; diff --git a/packages/agent/src/harness/session/repo-utils.ts b/packages/agent/src/harness/session/repo-utils.ts index a25b62def..dd2e0f9aa 100644 --- a/packages/agent/src/harness/session/repo-utils.ts +++ b/packages/agent/src/harness/session/repo-utils.ts @@ -1,3 +1,4 @@ +import { uuidv7 } from "@earendil-works/pi-ai"; import { type FileError, type Result, @@ -7,7 +8,6 @@ import { type SessionTreeEntry, } from "../types.ts"; import { Session } from "./session.ts"; -import { uuidv7 } from "./uuid.ts"; export function createSessionId(): string { return uuidv7(); diff --git a/packages/agent/src/index.ts b/packages/agent/src/index.ts index 89ba269af..49baf081f 100644 --- a/packages/agent/src/index.ts +++ b/packages/agent/src/index.ts @@ -33,7 +33,6 @@ export * from "./harness/session/memory-repo.ts"; export * from "./harness/session/memory-storage.ts"; export * from "./harness/session/repo-utils.ts"; export * from "./harness/session/session.ts"; -export { uuidv7 } from "./harness/session/uuid.ts"; export * from "./harness/skills.ts"; export * from "./harness/system-prompt.ts"; // Harness diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 694c828b0..7f9388ac7 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -2,8 +2,13 @@ ## [Unreleased] +### Added + +- Added a shared `uuidv7` utility for time-ordered identifiers. + ### Fixed +- Fixed sessionless OpenAI Codex WebSocket requests to use UUIDv7 request IDs, enabling models that reject UUIDv4 IDs. - Fixed GitHub Copilot long-context pricing tiers in generated model metadata ([#6668](https://github.com/earendil-works/pi/issues/6668)). - Fixed Kimi Coding subscription models to report API-equivalent implied costs when models.dev reports zero pricing. - Fixed OpenAI Responses early stream endings to be classified as retryable provider errors ([#6727](https://github.com/earendil-works/pi/issues/6727)). diff --git a/packages/ai/src/api/openai-codex-responses.ts b/packages/ai/src/api/openai-codex-responses.ts index de0d17925..955964fec 100644 --- a/packages/ai/src/api/openai-codex-responses.ts +++ b/packages/ai/src/api/openai-codex-responses.ts @@ -46,6 +46,7 @@ import { formatProviderError, normalizeProviderError } from "../utils/error-body import { AssistantMessageEventStream } from "../utils/event-stream.ts"; import { headersToRecord } from "../utils/headers.ts"; import { resolveHttpProxyUrlForTarget } from "../utils/node-http-proxy.ts"; +import { uuidv7 } from "../utils/uuid.ts"; import { clampOpenAIPromptCacheKey } from "./openai-prompt-cache.ts"; import { convertResponsesMessages, convertResponsesTools, processResponsesStream } from "./openai-responses-shared.ts"; import { buildBaseOptions } from "./simple-options.ts"; @@ -259,7 +260,7 @@ export const stream: StreamFunction<"openai-codex-responses", OpenAICodexRespons body = nextBody as RequestBody; } const codexSessionId = clampOpenAIPromptCacheKey(options?.sessionId); - const websocketRequestId = codexSessionId || createCodexRequestId(); + const websocketRequestId = codexSessionId || uuidv7(); const sseHeaders = buildSSEHeaders(model.headers, options?.headers, accountId, apiKey, codexSessionId); const websocketHeaders = buildWebSocketHeaders( model.headers, @@ -1505,13 +1506,6 @@ function extractAccountId(token: string): string { } } -function createCodexRequestId(): string { - if (typeof globalThis.crypto?.randomUUID === "function") { - return globalThis.crypto.randomUUID(); - } - return `codex_${Date.now()}_${Math.random().toString(36).slice(2, 10)}`; -} - function buildBaseCodexHeaders( initHeaders: Record | undefined, additionalHeaders: ProviderHeaders | undefined, diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 1056dc5ae..1ba458113 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -42,4 +42,5 @@ export * from "./utils/json-parse.ts"; export * from "./utils/overflow.ts"; export * from "./utils/retry.ts"; export * from "./utils/typebox-helpers.ts"; +export { uuidv7 } from "./utils/uuid.ts"; export * from "./utils/validation.ts"; diff --git a/packages/agent/src/harness/session/uuid.ts b/packages/ai/src/utils/uuid.ts similarity index 83% rename from packages/agent/src/harness/session/uuid.ts rename to packages/ai/src/utils/uuid.ts index 0c4cef8e0..53d900ca3 100644 --- a/packages/agent/src/harness/session/uuid.ts +++ b/packages/ai/src/utils/uuid.ts @@ -1,10 +1,9 @@ let lastTimestamp = -Infinity; let sequence = 0; -function fillRandomBytes(bytes: Uint8Array): void { - const crypto = globalThis.crypto; - if (crypto?.getRandomValues) { - crypto.getRandomValues(bytes); +function fillRandomBytes(bytes: Uint8Array): void { + if (globalThis.crypto?.getRandomValues) { + globalThis.crypto.getRandomValues(bytes); return; } for (let i = 0; i < bytes.length; i++) { @@ -12,6 +11,7 @@ function fillRandomBytes(bytes: Uint8Array): void { } } +/** Generate a time-ordered UUIDv7. */ export function uuidv7(): string { const random = new Uint8Array(16); fillRandomBytes(random); @@ -22,9 +22,7 @@ export function uuidv7(): string { lastTimestamp = timestamp; } else { sequence = (sequence + 1) >>> 0; - if (sequence === 0) { - lastTimestamp++; - } + if (sequence === 0) lastTimestamp++; } const bytes = new Uint8Array(16); @@ -45,10 +43,6 @@ export function uuidv7(): string { 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/test/harness/session-uuid.test.ts b/packages/ai/test/uuid.test.ts similarity index 96% rename from packages/agent/test/harness/session-uuid.test.ts rename to packages/ai/test/uuid.test.ts index b8d571f4e..3e5426502 100644 --- a/packages/agent/test/harness/session-uuid.test.ts +++ b/packages/ai/test/uuid.test.ts @@ -1,5 +1,5 @@ import { afterEach, describe, expect, it, vi } from "vitest"; -import { uuidv7 } from "../../src/harness/session/uuid.ts"; +import { uuidv7 } from "../src/utils/uuid.ts"; const UUID_V7_RE = /^[0-9a-f]{8}-[0-9a-f]{4}-7[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$/; const TIMESTAMP = 0x0123456789ab; diff --git a/packages/coding-agent/src/core/session-manager.ts b/packages/coding-agent/src/core/session-manager.ts index f5b75c2cb..0ab665957 100644 --- a/packages/coding-agent/src/core/session-manager.ts +++ b/packages/coding-agent/src/core/session-manager.ts @@ -1,5 +1,5 @@ -import { type AgentMessage, uuidv7 } from "@earendil-works/pi-agent-core"; -import type { ImageContent, Message, TextContent } from "@earendil-works/pi-ai"; +import type { AgentMessage } from "@earendil-works/pi-agent-core"; +import { type ImageContent, type Message, type TextContent, uuidv7 } from "@earendil-works/pi-ai"; import { randomUUID } from "crypto"; import { appendFileSync, From 75cb0b873af17e4f0d95b21d82ab6996084b978e Mon Sep 17 00:00:00 2001 From: Armin Ronacher Date: Mon, 20 Jul 2026 01:28:36 +0200 Subject: [PATCH 04/62] fix(ai): support responses API for OpenCode Go --- packages/ai/src/providers/opencode-go.ts | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/packages/ai/src/providers/opencode-go.ts b/packages/ai/src/providers/opencode-go.ts index 608f579b7..e2f341eca 100644 --- a/packages/ai/src/providers/opencode-go.ts +++ b/packages/ai/src/providers/opencode-go.ts @@ -1,11 +1,12 @@ import { anthropicMessagesApi } from "../api/anthropic-messages.lazy.ts"; import { openAICompletionsApi } from "../api/openai-completions.lazy.ts"; +import { openAIResponsesApi } from "../api/openai-responses.lazy.ts"; import { envApiKeyAuth } from "../auth/helpers.ts"; import { createProvider, type Provider } from "../models.ts"; import { OPENCODE_GO_MODELS } from "./opencode-go.models.ts"; -export function opencodeGoProvider(): Provider<"anthropic-messages" | "openai-completions"> { - return createProvider({ +export function opencodeGoProvider(): Provider<"anthropic-messages" | "openai-completions" | "openai-responses"> { + return createProvider<"anthropic-messages" | "openai-completions" | "openai-responses">({ id: "opencode-go", name: "OpenCode Zen Go", auth: { apiKey: envApiKeyAuth("OpenCode API key", ["OPENCODE_API_KEY"]) }, @@ -13,6 +14,7 @@ export function opencodeGoProvider(): Provider<"anthropic-messages" | "openai-co api: { "anthropic-messages": anthropicMessagesApi(), "openai-completions": openAICompletionsApi(), + "openai-responses": openAIResponsesApi(), }, }); } From 94373d815d2b4a3a48864d5341afc824b8db45e3 Mon Sep 17 00:00:00 2001 From: Alexey Zaytsev Date: Sun, 19 Jul 2026 20:50:36 -0300 Subject: [PATCH 05/62] feat(ai): add shared contentText utility (#6840) Co-authored-by: Armin Ronacher --- packages/agent/src/harness/agent-harness.ts | 26 ++++++-------- .../compaction/branch-summarization.ts | 7 ++-- .../src/harness/compaction/compaction.ts | 22 ++++++------ .../agent/src/harness/compaction/utils.ts | 24 ++++--------- packages/ai/CHANGELOG.md | 1 + packages/ai/src/index.ts | 1 + packages/ai/src/utils/text.ts | 12 +++++++ packages/ai/test/text.test.ts | 33 +++++++++++++++++ .../coding-agent/src/core/agent-session.ts | 36 +++---------------- .../core/compaction/branch-summarization.ts | 6 ++-- .../src/core/compaction/compaction.ts | 11 ++---- .../coding-agent/src/core/compaction/utils.ts | 24 ++++--------- 12 files changed, 92 insertions(+), 111 deletions(-) create mode 100644 packages/ai/src/utils/text.ts create mode 100644 packages/ai/test/text.test.ts diff --git a/packages/agent/src/harness/agent-harness.ts b/packages/agent/src/harness/agent-harness.ts index 1d09b054e..1ce57b14d 100644 --- a/packages/agent/src/harness/agent-harness.ts +++ b/packages/agent/src/harness/agent-harness.ts @@ -1,4 +1,11 @@ -import type { AssistantMessage, ImageContent, Model, Models, UserMessage } from "@earendil-works/pi-ai"; +import { + type AssistantMessage, + contentText, + type ImageContent, + type Model, + type Models, + type UserMessage, +} from "@earendil-works/pi-ai"; import { runAgentLoop } from "../agent-loop.ts"; import type { AgentContext, @@ -781,23 +788,10 @@ export class AgentHarness< 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(""); + editorText = contentText(targetEntry.message.content, ""); } 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(""); + editorText = contentText(targetEntry.content, ""); } else { newLeafId = targetId; } diff --git a/packages/agent/src/harness/compaction/branch-summarization.ts b/packages/agent/src/harness/compaction/branch-summarization.ts index fdf1df498..a8d686739 100644 --- a/packages/agent/src/harness/compaction/branch-summarization.ts +++ b/packages/agent/src/harness/compaction/branch-summarization.ts @@ -1,4 +1,4 @@ -import type { Model, Models } from "@earendil-works/pi-ai"; +import { contentText, type Model, type Models } from "@earendil-works/pi-ai"; import type { AgentMessage } from "../../types.ts"; import { @@ -245,10 +245,7 @@ export async function generateBranchSummary( ); } - let summary = response.content - .filter((c): c is { type: "text"; text: string } => c.type === "text") - .map((c) => c.text) - .join("\n"); + let summary = contentText(response.content); summary = BRANCH_SUMMARY_PREAMBLE + summary; const { readFiles, modifiedFiles } = computeFileLists(fileOps); summary += formatFileOperations(readFiles, modifiedFiles); diff --git a/packages/agent/src/harness/compaction/compaction.ts b/packages/agent/src/harness/compaction/compaction.ts index 4132e3ddc..a7f8ebc06 100644 --- a/packages/agent/src/harness/compaction/compaction.ts +++ b/packages/agent/src/harness/compaction/compaction.ts @@ -1,4 +1,12 @@ -import type { AssistantMessage, ImageContent, Model, Models, TextContent, Usage } from "@earendil-works/pi-ai"; +import { + type AssistantMessage, + contentText, + type ImageContent, + type Model, + type Models, + type TextContent, + type Usage, +} from "@earendil-works/pi-ai"; import type { AgentMessage, ThinkingLevel } from "../../types.ts"; import { convertToLlm, @@ -513,10 +521,7 @@ export async function generateSummary( ); } - const textContent = response.content - .filter((c): c is { type: "text"; text: string } => c.type === "text") - .map((c) => c.text) - .join("\n"); + const textContent = contentText(response.content); return ok(textContent); } @@ -744,10 +749,5 @@ async function generateTurnPrefixSummary( ); } - return ok( - response.content - .filter((c): c is { type: "text"; text: string } => c.type === "text") - .map((c) => c.text) - .join("\n"), - ); + return ok(contentText(response.content)); } diff --git a/packages/agent/src/harness/compaction/utils.ts b/packages/agent/src/harness/compaction/utils.ts index 07535b79b..7fd3aac02 100644 --- a/packages/agent/src/harness/compaction/utils.ts +++ b/packages/agent/src/harness/compaction/utils.ts @@ -1,4 +1,4 @@ -import type { Message } from "@earendil-works/pi-ai"; +import { contentText, type Message } from "@earendil-works/pi-ai"; import type { AgentMessage } from "../../types.ts"; /** File paths touched by a session branch or compaction range. */ @@ -93,23 +93,14 @@ export function serializeConversation(messages: Message[]): 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(""); + const content = contentText(msg.content, ""); 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") { + if (block.type === "thinking") { thinkingParts.push(block.thinking); } else if (block.type === "toolCall") { const args = block.arguments as Record; @@ -123,17 +114,14 @@ export function serializeConversation(messages: Message[]): string { if (thinkingParts.length > 0) { parts.push(`[Assistant thinking]: ${thinkingParts.join("\n")}`); } - if (textParts.length > 0) { - parts.push(`[Assistant]: ${textParts.join("\n")}`); + if (msg.content.some((block) => block.type === "text")) { + parts.push(`[Assistant]: ${contentText(msg.content)}`); } 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(""); + const content = contentText(msg.content, ""); if (content) { parts.push(`[Tool result]: ${truncateForSummary(content, TOOL_RESULT_MAX_CHARS)}`); } diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 7f9388ac7..ed86ca720 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -4,6 +4,7 @@ ### Added +- Added `contentText` for extracting joined text from message content. - Added a shared `uuidv7` utility for time-ordered identifiers. ### Fixed diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 1ba458113..4ff678102 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -41,6 +41,7 @@ export * from "./utils/event-stream.ts"; export * from "./utils/json-parse.ts"; export * from "./utils/overflow.ts"; export * from "./utils/retry.ts"; +export { contentText } from "./utils/text.ts"; export * from "./utils/typebox-helpers.ts"; export { uuidv7 } from "./utils/uuid.ts"; export * from "./utils/validation.ts"; diff --git a/packages/ai/src/utils/text.ts b/packages/ai/src/utils/text.ts new file mode 100644 index 000000000..66f5b9c5b --- /dev/null +++ b/packages/ai/src/utils/text.ts @@ -0,0 +1,12 @@ +import type { ImageContent, TextContent, ThinkingContent, ToolCall } from "../types.ts"; + +type Content = TextContent | ImageContent | ThinkingContent | ToolCall; + +/** Extract and join text from message content. */ +export function contentText(content: string | readonly Content[], separator = "\n"): string { + if (typeof content === "string") return content; + return content + .filter((block) => block.type === "text") + .map((block) => block.text) + .join(separator); +} diff --git a/packages/ai/test/text.test.ts b/packages/ai/test/text.test.ts new file mode 100644 index 000000000..43dc21c84 --- /dev/null +++ b/packages/ai/test/text.test.ts @@ -0,0 +1,33 @@ +import { describe, expect, it } from "vitest"; +import { type AssistantMessage, contentText, type ToolResultMessage } from "../src/index.ts"; + +const content: AssistantMessage["content"] = [ + { type: "thinking", thinking: "reasoning" }, + { type: "text", text: "first" }, + { type: "toolCall", id: "1", name: "read", arguments: {} }, + { type: "text", text: "second" }, +]; + +describe("contentText", () => { + it("extracts assistant text blocks", () => { + expect(contentText(content)).toBe("first\nsecond"); + }); + + it("supports custom separators", () => { + expect(contentText(content, "")).toBe("firstsecond"); + }); + + it("passes string content through", () => { + expect(contentText("hello")).toBe("hello"); + }); + + it("extracts text from tool-result content", () => { + const toolResultContent: ToolResultMessage["content"] = [ + { type: "text", text: "first" }, + { type: "image", data: "...", mimeType: "image/png" }, + { type: "text", text: "second" }, + ]; + + expect(contentText(toolResultContent, "")).toBe("firstsecond"); + }); +}); diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 36e542e50..c4afb9b17 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -24,11 +24,11 @@ import type { PrepareNextTurnContext, ThinkingLevel, } from "@earendil-works/pi-agent-core"; +import { contentText } from "@earendil-works/pi-ai"; import type { AssistantMessage, AuthResult, ImageContent, - Message, Model, ProviderHeaders, TextContent, @@ -576,7 +576,7 @@ export class AgentSession { // This ensures the UI sees the updated queue state if (event.type === "message_start" && event.message.role === "user") { this._overflowRecoveryAttempted = false; - const messageText = this._getUserMessageText(event.message); + const messageText = contentText(event.message.content, ""); if (messageText) { // Check steering queue first const steeringIndex = this._steeringMessages.indexOf(messageText); @@ -659,15 +659,6 @@ export class AgentSession { return false; } - /** Extract text content from a message */ - private _getUserMessageText(message: Message): string { - if (message.role !== "user") return ""; - const content = message.content; - if (typeof content === "string") return content; - const textBlocks = content.filter((c) => c.type === "text"); - return textBlocks.map((c) => (c as TextContent).text).join(""); - } - /** Find the last assistant message in agent state (including aborted ones) */ private _findLastAssistantMessage(): AssistantMessage | undefined { const messages = this.agent.state.messages; @@ -2954,17 +2945,11 @@ export class AgentSession { if (targetEntry.type === "message" && targetEntry.message.role === "user") { // User message: leaf = parent (null if root), text goes to editor newLeafId = targetEntry.parentId; - editorText = this._extractUserMessageText(targetEntry.message.content); + editorText = contentText(targetEntry.message.content, ""); } else if (targetEntry.type === "custom_message") { // Custom message: leaf = parent (null if root), text goes to editor newLeafId = targetEntry.parentId; - editorText = - typeof targetEntry.content === "string" - ? targetEntry.content - : targetEntry.content - .filter((c): c is { type: "text"; text: string } => c.type === "text") - .map((c) => c.text) - .join(""); + editorText = contentText(targetEntry.content, ""); } else { // Non-user message: leaf = selected node newLeafId = targetId; @@ -3032,7 +3017,7 @@ export class AgentSession { if (entry.type !== "message") continue; if (entry.message.role !== "user") continue; - const text = this._extractUserMessageText(entry.message.content); + const text = contentText(entry.message.content, ""); if (text) { result.push({ entryId: entry.id, text }); } @@ -3041,17 +3026,6 @@ export class AgentSession { return result; } - private _extractUserMessageText(content: string | Array<{ type: string; text?: string }>): string { - if (typeof content === "string") return content; - if (Array.isArray(content)) { - return content - .filter((c): c is { type: "text"; text: string } => c.type === "text") - .map((c) => c.text) - .join(""); - } - return ""; - } - /** * Get session statistics. Aggregates over ALL session entries (including * history that was compacted away), so token/cost totals reflect what was diff --git a/packages/coding-agent/src/core/compaction/branch-summarization.ts b/packages/coding-agent/src/core/compaction/branch-summarization.ts index a96bb202d..f2e06876c 100644 --- a/packages/coding-agent/src/core/compaction/branch-summarization.ts +++ b/packages/coding-agent/src/core/compaction/branch-summarization.ts @@ -6,6 +6,7 @@ */ import type { AgentMessage, StreamFn } from "@earendil-works/pi-agent-core"; +import { contentText } from "@earendil-works/pi-ai"; import type { Model, SimpleStreamOptions } from "@earendil-works/pi-ai/compat"; import { completeSimple } from "@earendil-works/pi-ai/compat"; import { @@ -351,10 +352,7 @@ export async function generateBranchSummary( 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"); + let summary = contentText(response.content); // Prepend preamble to provide context about the branch summary summary = BRANCH_SUMMARY_PREAMBLE + summary; diff --git a/packages/coding-agent/src/core/compaction/compaction.ts b/packages/coding-agent/src/core/compaction/compaction.ts index 39b1042d3..5c71c74dc 100644 --- a/packages/coding-agent/src/core/compaction/compaction.ts +++ b/packages/coding-agent/src/core/compaction/compaction.ts @@ -6,6 +6,7 @@ */ import type { AgentMessage, StreamFn, ThinkingLevel } from "@earendil-works/pi-agent-core"; +import { contentText } from "@earendil-works/pi-ai"; import type { AssistantMessage, Context, Model, SimpleStreamOptions, Usage } from "@earendil-works/pi-ai/compat"; import { completeSimple } from "@earendil-works/pi-ai/compat"; import { convertToLlm } from "../messages.ts"; @@ -600,10 +601,7 @@ export async function generateSummary( 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"); + const textContent = contentText(response.content); return textContent; } @@ -865,8 +863,5 @@ async function generateTurnPrefixSummary( 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"); + return contentText(response.content); } diff --git a/packages/coding-agent/src/core/compaction/utils.ts b/packages/coding-agent/src/core/compaction/utils.ts index 6cfc16227..35d1a29bb 100644 --- a/packages/coding-agent/src/core/compaction/utils.ts +++ b/packages/coding-agent/src/core/compaction/utils.ts @@ -3,7 +3,7 @@ */ import type { AgentMessage } from "@earendil-works/pi-agent-core"; -import type { Message } from "@earendil-works/pi-ai"; +import { contentText, type Message } from "@earendil-works/pi-ai"; // ============================================================================ // File Operation Tracking @@ -111,23 +111,14 @@ export function serializeConversation(messages: Message[]): 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(""); + const content = contentText(msg.content, ""); 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") { + if (block.type === "thinking") { thinkingParts.push(block.thinking); } else if (block.type === "toolCall") { const args = block.arguments as Record; @@ -141,17 +132,14 @@ export function serializeConversation(messages: Message[]): string { if (thinkingParts.length > 0) { parts.push(`[Assistant thinking]: ${thinkingParts.join("\n")}`); } - if (textParts.length > 0) { - parts.push(`[Assistant]: ${textParts.join("\n")}`); + if (msg.content.some((block) => block.type === "text")) { + parts.push(`[Assistant]: ${contentText(msg.content)}`); } 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(""); + const content = contentText(msg.content, ""); if (content) { parts.push(`[Tool result]: ${truncateForSummary(content, TOOL_RESULT_MAX_CHARS)}`); } From 5f2f7d06793738605681dc69e17feea120cba3f9 Mon Sep 17 00:00:00 2001 From: Damjan 9000 Date: Mon, 20 Jul 2026 09:09:56 +0200 Subject: [PATCH 06/62] fix(tui): clear inverted cursor on exit to avoid dual cursor appearance (#6790) When pi exits, the reverse-video cursor character remains visible on the editor prompt line. Combined with the terminal's own cursor below, this creates the appearance of two cursors, which is confusing. Overwrite the inverted cursor with a normal space before exiting. --- packages/tui/src/tui.ts | 2 ++ 1 file changed, 2 insertions(+) diff --git a/packages/tui/src/tui.ts b/packages/tui/src/tui.ts index a7054888a..b7ebaea98 100644 --- a/packages/tui/src/tui.ts +++ b/packages/tui/src/tui.ts @@ -695,6 +695,8 @@ export class TUI extends Container { } // Move cursor to the end of the content to prevent overwriting/artifacts on exit if (this.previousLines.length > 0) { + // Overwrite the inverted cursor with a normal space to clear the artifact + this.terminal.write(" "); const targetRow = this.previousLines.length; // Line after the last content const lineDiff = targetRow - this.hardwareCursorRow; if (lineDiff > 0) { From 916de90dbf36658b83c9a54130afe768659c8332 Mon Sep 17 00:00:00 2001 From: Federico Jaramillo Martinez Date: Mon, 20 Jul 2026 09:18:14 +0200 Subject: [PATCH 07/62] fix(pi-ai): normalize bin path to avoid consumer lockfile churn (#6812) The npm registry normalizes bin paths in packument metadata ("dist/cli.js") while the published tarball keeps the authored "./dist/cli.js". npm consumers get perpetual one-line package-lock.json flips depending on whether resolution used registry metadata or the extracted package. Drop the "./" prefix to match the registry-normalized form and the sibling coding-agent package. Fixes #6811 --- packages/ai/package.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/ai/package.json b/packages/ai/package.json index a3ed55f04..58b9113f8 100644 --- a/packages/ai/package.json +++ b/packages/ai/package.json @@ -41,7 +41,7 @@ } }, "bin": { - "pi-ai": "./dist/cli.js" + "pi-ai": "dist/cli.js" }, "files": [ "dist", From 87845fc430c544e29e899cb3ef61545fd0e01751 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Mon, 20 Jul 2026 07:19:18 +0000 Subject: [PATCH 08/62] chore: approve contributor rsaryev --- .github/APPROVED_CONTRIBUTORS | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/APPROVED_CONTRIBUTORS b/.github/APPROVED_CONTRIBUTORS index dc3961ec9..d65a5452f 100644 --- a/.github/APPROVED_CONTRIBUTORS +++ b/.github/APPROVED_CONTRIBUTORS @@ -279,3 +279,5 @@ ananthakumaran pr andrebreijao pr anh-chu pr + +rsaryev pr From 85f89db9bcf4104b1e207ddb6f787bc5a4b631ce Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Mon, 20 Jul 2026 07:33:27 +0000 Subject: [PATCH 09/62] chore: approve contributor QuintinShaw --- .github/APPROVED_CONTRIBUTORS | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/APPROVED_CONTRIBUTORS b/.github/APPROVED_CONTRIBUTORS index d65a5452f..1c70bd0e9 100644 --- a/.github/APPROVED_CONTRIBUTORS +++ b/.github/APPROVED_CONTRIBUTORS @@ -281,3 +281,5 @@ andrebreijao pr anh-chu pr rsaryev pr + +QuintinShaw pr From 8b9373707845631ad286580b1d2c3f938e5c263f Mon Sep 17 00:00:00 2001 From: David Brailovsky Date: Mon, 20 Jul 2026 09:00:15 +0000 Subject: [PATCH 10/62] kimi: add low,high to k3 and remove k2p7 references --- packages/ai/scripts/generate-models.ts | 5 ++--- packages/ai/src/providers/kimi-coding.models.ts | 8 ++++---- packages/ai/test/abort.test.ts | 2 +- .../ai/test/anthropic-adaptive-thinking-models.test.ts | 2 +- .../ai/test/anthropic-force-adaptive-thinking.test.ts | 2 +- packages/ai/test/context-overflow.test.ts | 4 ++-- packages/ai/test/cross-provider-handoff.test.ts | 2 +- packages/ai/test/empty.test.ts | 2 +- packages/ai/test/image-tool-result.test.ts | 4 ++-- packages/ai/test/providers.test.ts | 1 - packages/ai/test/stream.test.ts | 4 ++-- packages/ai/test/supports-xhigh.test.ts | 4 ++-- packages/ai/test/tokens.test.ts | 2 +- packages/ai/test/tool-call-without-result.test.ts | 2 +- packages/ai/test/total-tokens.test.ts | 4 ++-- packages/ai/test/unicode-surrogate.test.ts | 2 +- 16 files changed, 24 insertions(+), 26 deletions(-) diff --git a/packages/ai/scripts/generate-models.ts b/packages/ai/scripts/generate-models.ts index 53d5a00a1..38c1464ea 100644 --- a/packages/ai/scripts/generate-models.ts +++ b/packages/ai/scripts/generate-models.ts @@ -245,9 +245,9 @@ const DEEPSEEK_V4_THINKING_LEVEL_MAP = { const KIMI_K3_THINKING_LEVEL_MAP = { off: null, minimal: null, - low: null, + low: "low", medium: null, - high: null, + high: "high", xhigh: null, max: "max", } as const; @@ -261,7 +261,6 @@ const KIMI_K3_COST = { // Kimi Coding is subscription-backed, so models.dev reports zero cost. Use the // equivalent Moonshot API rates to estimate the value of subscription usage. const KIMI_CODING_IMPLIED_COSTS: Record["cost"]> = { - k2p7: { input: 0.95, output: 4, cacheRead: 0.19, cacheWrite: 0 }, k3: KIMI_K3_COST, "kimi-for-coding": { input: 0.95, output: 4, cacheRead: 0.19, cacheWrite: 0 }, "kimi-for-coding-highspeed": { input: 1.9, output: 8, cacheRead: 0.38, cacheWrite: 0 }, diff --git a/packages/ai/src/providers/kimi-coding.models.ts b/packages/ai/src/providers/kimi-coding.models.ts index cfd4c68ec..ac810a751 100644 --- a/packages/ai/src/providers/kimi-coding.models.ts +++ b/packages/ai/src/providers/kimi-coding.models.ts @@ -5,14 +5,14 @@ import values from "./data/kimi-coding.json" with { type: "json" }; import type { Model } from "../types.ts"; export const KIMI_CODING_MODELS = values as { - "k2p7": Model<"anthropic-messages"> & { - id: "k2p7"; - provider: "kimi-coding"; - }; "k3": Model<"anthropic-messages"> & { id: "k3"; provider: "kimi-coding"; }; + "kimi-for-coding": Model<"anthropic-messages"> & { + id: "kimi-for-coding"; + provider: "kimi-coding"; + }; "kimi-for-coding-highspeed": Model<"anthropic-messages"> & { id: "kimi-for-coding-highspeed"; provider: "kimi-coding"; diff --git a/packages/ai/test/abort.test.ts b/packages/ai/test/abort.test.ts index ce81422cb..c79cf540b 100644 --- a/packages/ai/test/abort.test.ts +++ b/packages/ai/test/abort.test.ts @@ -250,7 +250,7 @@ describe("AI Providers Abort Tests", () => { }); describe.skipIf(!process.env.KIMI_API_KEY)("Kimi For Coding Provider Abort", () => { - const llm = getModel("kimi-coding", "k2p7"); + const llm = getModel("kimi-coding", "kimi-for-coding"); it("should abort mid-stream", { retry: 3 }, async () => { await testAbortSignal(llm); diff --git a/packages/ai/test/anthropic-adaptive-thinking-models.test.ts b/packages/ai/test/anthropic-adaptive-thinking-models.test.ts index 7807ffc17..329c8a725 100644 --- a/packages/ai/test/anthropic-adaptive-thinking-models.test.ts +++ b/packages/ai/test/anthropic-adaptive-thinking-models.test.ts @@ -7,7 +7,7 @@ const EXPECTED_CURRENT_ADAPTIVE_THINKING_MODELS = [ "anthropic/claude-opus-4-8", "anthropic/claude-sonnet-5", "cloudflare-ai-gateway/claude-fable-5", - "kimi-coding/k2p7", + "kimi-coding/kimi-for-coding", "kimi-coding/k3", "kimi-coding/kimi-for-coding-highspeed", "opencode/claude-opus-4-8", diff --git a/packages/ai/test/anthropic-force-adaptive-thinking.test.ts b/packages/ai/test/anthropic-force-adaptive-thinking.test.ts index 0a922db7c..ffaf48dbe 100644 --- a/packages/ai/test/anthropic-force-adaptive-thinking.test.ts +++ b/packages/ai/test/anthropic-force-adaptive-thinking.test.ts @@ -90,7 +90,7 @@ describe("Anthropic forceAdaptiveThinking compat override", () => { }); it.each([ - ["k2p7", "medium", "medium"], + ["kimi-for-coding", "medium", "medium"], ["k3", "max", "max"], ["kimi-for-coding-highspeed", "medium", "medium"], ] as const)( diff --git a/packages/ai/test/context-overflow.test.ts b/packages/ai/test/context-overflow.test.ts index e5acb6f39..2a8b62845 100644 --- a/packages/ai/test/context-overflow.test.ts +++ b/packages/ai/test/context-overflow.test.ts @@ -471,8 +471,8 @@ describe("Context overflow error handling", () => { // ============================================================================= describe.skipIf(!process.env.KIMI_API_KEY)("Kimi For Coding", () => { - it("k2p7 - should detect overflow via isContextOverflow", async () => { - const model = getModel("kimi-coding", "k2p7"); + it("kimi-for-coding - should detect overflow via isContextOverflow", async () => { + const model = getModel("kimi-coding", "kimi-for-coding"); const result = await testContextOverflow(model, process.env.KIMI_API_KEY!); logResult(result); diff --git a/packages/ai/test/cross-provider-handoff.test.ts b/packages/ai/test/cross-provider-handoff.test.ts index fe3b7adae..08da44de3 100644 --- a/packages/ai/test/cross-provider-handoff.test.ts +++ b/packages/ai/test/cross-provider-handoff.test.ts @@ -109,7 +109,7 @@ const PROVIDER_MODEL_PAIRS: ProviderModelPair[] = [ // Together AI { provider: "together", model: "moonshotai/Kimi-K2.6", label: "together-kimi-k2.6" }, // Kimi For Coding - { provider: "kimi-coding", model: "k2p7", label: "kimi-coding-k2p7" }, + { provider: "kimi-coding", model: "kimi-for-coding", label: "kimi-for-coding" }, // Mistral { provider: "mistral", model: "devstral-medium-latest", label: "mistral-devstral-medium" }, // MiniMax diff --git a/packages/ai/test/empty.test.ts b/packages/ai/test/empty.test.ts index 2f75f888b..502764c95 100644 --- a/packages/ai/test/empty.test.ts +++ b/packages/ai/test/empty.test.ts @@ -536,7 +536,7 @@ describe("AI Providers Empty Message Tests", () => { ); describe.skipIf(!process.env.KIMI_API_KEY)("Kimi For Coding Provider Empty Messages", () => { - const llm = getModel("kimi-coding", "k2p7"); + const llm = getModel("kimi-coding", "kimi-for-coding"); it("should handle empty content array", { retry: 3, timeout: 30000 }, async () => { await testEmptyMessage(llm); diff --git a/packages/ai/test/image-tool-result.test.ts b/packages/ai/test/image-tool-result.test.ts index 7c96d6cc9..946a443ad 100644 --- a/packages/ai/test/image-tool-result.test.ts +++ b/packages/ai/test/image-tool-result.test.ts @@ -381,8 +381,8 @@ describe("Tool Results with Images", () => { }, ); - describe.skipIf(!process.env.KIMI_API_KEY)("Kimi For Coding Provider (k2p7)", () => { - const llm = getModel("kimi-coding", "k2p7"); + describe.skipIf(!process.env.KIMI_API_KEY)("Kimi For Coding Provider (kimi-for-coding)", () => { + const llm = getModel("kimi-coding", "kimi-for-coding"); it("should handle tool result with only image", { retry: 3, timeout: 30000 }, async () => { await handleToolWithImageResult(llm); diff --git a/packages/ai/test/providers.test.ts b/packages/ai/test/providers.test.ts index fb2f67a34..e034ebecd 100644 --- a/packages/ai/test/providers.test.ts +++ b/packages/ai/test/providers.test.ts @@ -59,7 +59,6 @@ describe("builtin providers", () => { it("uses API-equivalent implied pricing for Kimi Coding subscription models", () => { const models = builtinModels(); const expectedCosts = { - k2p7: { input: 0.95, output: 4, cacheRead: 0.19, cacheWrite: 0 }, k3: { input: 3, output: 15, cacheRead: 0.3, cacheWrite: 0 }, "kimi-for-coding-highspeed": { input: 1.9, output: 8, cacheRead: 0.38, cacheWrite: 0 }, }; diff --git a/packages/ai/test/stream.test.ts b/packages/ai/test/stream.test.ts index 13a53c911..4209f94d1 100644 --- a/packages/ai/test/stream.test.ts +++ b/packages/ai/test/stream.test.ts @@ -1017,8 +1017,8 @@ describe("Generate E2E Tests", () => { }); }); - describe.skipIf(!process.env.KIMI_API_KEY)("Kimi For Coding Provider (k2p7 via Anthropic Messages)", () => { - const llm = getModel("kimi-coding", "k2p7"); + describe.skipIf(!process.env.KIMI_API_KEY)("Kimi For Coding Provider (kimi-for-coding via Anthropic Messages)", () => { + const llm = getModel("kimi-coding", "kimi-for-coding"); it("should complete basic text generation", { retry: 3 }, async () => { await basicTextGeneration(llm); diff --git a/packages/ai/test/supports-xhigh.test.ts b/packages/ai/test/supports-xhigh.test.ts index fa1eba91b..b6988ba06 100644 --- a/packages/ai/test/supports-xhigh.test.ts +++ b/packages/ai/test/supports-xhigh.test.ts @@ -110,10 +110,10 @@ describe("getSupportedThinkingLevels", () => { } }); - it("includes only max for Kimi Coding K3", () => { + it("includes only low, high, max for Kimi Coding K3", () => { const model = getModel("kimi-coding", "k3"); expect(model).toBeDefined(); - expect(getSupportedThinkingLevels(model!)).toEqual(["max"]); + expect(getSupportedThinkingLevels(model!)).toEqual(["low", "high", "max"]); }); it("includes only high for OpenCode Grok Build", () => { diff --git a/packages/ai/test/tokens.test.ts b/packages/ai/test/tokens.test.ts index 9422a86ec..42e6932d4 100644 --- a/packages/ai/test/tokens.test.ts +++ b/packages/ai/test/tokens.test.ts @@ -218,7 +218,7 @@ describe("Token Statistics on Abort", () => { }); describe.skipIf(!process.env.KIMI_API_KEY)("Kimi For Coding Provider", () => { - const llm = getModel("kimi-coding", "k2p7"); + const llm = getModel("kimi-coding", "kimi-for-coding"); it("should include token stats when aborted mid-stream", { retry: 3, timeout: 30000 }, async () => { await testTokensOnAbort(llm); diff --git a/packages/ai/test/tool-call-without-result.test.ts b/packages/ai/test/tool-call-without-result.test.ts index 22fb69492..439b97580 100644 --- a/packages/ai/test/tool-call-without-result.test.ts +++ b/packages/ai/test/tool-call-without-result.test.ts @@ -255,7 +255,7 @@ describe("Tool Call Without Result Tests", () => { }); describe.skipIf(!process.env.KIMI_API_KEY)("Kimi For Coding Provider", () => { - const model = getModel("kimi-coding", "k2p7"); + const model = getModel("kimi-coding", "kimi-for-coding"); it("should filter out tool calls without corresponding tool results", { retry: 3, timeout: 30000 }, async () => { await testToolCallWithoutResult(model); diff --git a/packages/ai/test/total-tokens.test.ts b/packages/ai/test/total-tokens.test.ts index 93f06c1f7..827dcdb19 100644 --- a/packages/ai/test/total-tokens.test.ts +++ b/packages/ai/test/total-tokens.test.ts @@ -567,8 +567,8 @@ describe("totalTokens field", () => { // ========================================================================= describe.skipIf(!process.env.KIMI_API_KEY)("Kimi For Coding", () => { - it("k2p7 - should return totalTokens equal to sum of components", { retry: 3, timeout: 60000 }, async () => { - const llm = getModel("kimi-coding", "k2p7"); + it("kimi-for-coding - should return totalTokens equal to sum of components", { retry: 3, timeout: 60000 }, async () => { + const llm = getModel("kimi-coding", "kimi-for-coding"); console.log(`\nKimi For Coding / ${llm.id}:`); const { first, second } = await testTotalTokensWithCache(llm, { apiKey: process.env.KIMI_API_KEY }); diff --git a/packages/ai/test/unicode-surrogate.test.ts b/packages/ai/test/unicode-surrogate.test.ts index baf4df100..8e059d1d0 100644 --- a/packages/ai/test/unicode-surrogate.test.ts +++ b/packages/ai/test/unicode-surrogate.test.ts @@ -696,7 +696,7 @@ describe("AI Providers Unicode Surrogate Pair Tests", () => { ); describe.skipIf(!process.env.KIMI_API_KEY)("Kimi For Coding Provider Unicode Handling", () => { - const llm = getModel("kimi-coding", "k2p7"); + const llm = getModel("kimi-coding", "kimi-for-coding"); it("should handle emoji in tool results", { retry: 3, timeout: 30000 }, async () => { await testEmojiInToolResults(llm); From da95649b957d6a3f4e1ab7f60baae7809cf65b03 Mon Sep 17 00:00:00 2001 From: David Brailovsky Date: Mon, 20 Jul 2026 10:19:01 +0000 Subject: [PATCH 11/62] lint fix --- packages/ai/test/stream.test.ts | 39 ++++++++++++++------------- packages/ai/test/total-tokens.test.ts | 22 ++++++++------- 2 files changed, 34 insertions(+), 27 deletions(-) diff --git a/packages/ai/test/stream.test.ts b/packages/ai/test/stream.test.ts index 4209f94d1..484b74518 100644 --- a/packages/ai/test/stream.test.ts +++ b/packages/ai/test/stream.test.ts @@ -1017,29 +1017,32 @@ describe("Generate E2E Tests", () => { }); }); - describe.skipIf(!process.env.KIMI_API_KEY)("Kimi For Coding Provider (kimi-for-coding via Anthropic Messages)", () => { - const llm = getModel("kimi-coding", "kimi-for-coding"); + describe.skipIf(!process.env.KIMI_API_KEY)( + "Kimi For Coding Provider (kimi-for-coding via Anthropic Messages)", + () => { + const llm = getModel("kimi-coding", "kimi-for-coding"); - it("should complete basic text generation", { retry: 3 }, async () => { - await basicTextGeneration(llm); - }); + it("should complete basic text generation", { retry: 3 }, async () => { + await basicTextGeneration(llm); + }); - it("should handle tool calling", { retry: 3 }, async () => { - await handleToolCall(llm); - }); + it("should handle tool calling", { retry: 3 }, async () => { + await handleToolCall(llm); + }); - it("should handle streaming", { retry: 3 }, async () => { - await handleStreaming(llm); - }); + it("should handle streaming", { retry: 3 }, async () => { + await handleStreaming(llm); + }); - it("should handle thinking mode", { retry: 3 }, async () => { - await handleThinking(llm, { thinkingEnabled: true, thinkingBudgetTokens: 2048 }); - }); + it("should handle thinking mode", { retry: 3 }, async () => { + await handleThinking(llm, { thinkingEnabled: true, thinkingBudgetTokens: 2048 }); + }); - it("should handle multi-turn with thinking and tools", { retry: 3 }, async () => { - await multiTurn(llm, { thinkingEnabled: true, thinkingBudgetTokens: 2048 }); - }); - }); + it("should handle multi-turn with thinking and tools", { retry: 3 }, async () => { + await multiTurn(llm, { thinkingEnabled: true, thinkingBudgetTokens: 2048 }); + }); + }, + ); describe.skipIf(!process.env.XIAOMI_API_KEY)( "Xiaomi MiMo (API billing) Provider (Xiaomi MiMo-V2.5-Pro via Anthropic Messages)", diff --git a/packages/ai/test/total-tokens.test.ts b/packages/ai/test/total-tokens.test.ts index 827dcdb19..ed5cae292 100644 --- a/packages/ai/test/total-tokens.test.ts +++ b/packages/ai/test/total-tokens.test.ts @@ -567,18 +567,22 @@ describe("totalTokens field", () => { // ========================================================================= describe.skipIf(!process.env.KIMI_API_KEY)("Kimi For Coding", () => { - it("kimi-for-coding - should return totalTokens equal to sum of components", { retry: 3, timeout: 60000 }, async () => { - const llm = getModel("kimi-coding", "kimi-for-coding"); + it( + "kimi-for-coding - should return totalTokens equal to sum of components", + { retry: 3, timeout: 60000 }, + async () => { + const llm = getModel("kimi-coding", "kimi-for-coding"); - console.log(`\nKimi For Coding / ${llm.id}:`); - const { first, second } = await testTotalTokensWithCache(llm, { apiKey: process.env.KIMI_API_KEY }); + console.log(`\nKimi For Coding / ${llm.id}:`); + const { first, second } = await testTotalTokensWithCache(llm, { apiKey: process.env.KIMI_API_KEY }); - logUsage("First request", first); - logUsage("Second request", second); + logUsage("First request", first); + logUsage("Second request", second); - assertTotalTokensEqualsComponents(first); - assertTotalTokensEqualsComponents(second); - }); + assertTotalTokensEqualsComponents(first); + assertTotalTokensEqualsComponents(second); + }, + ); }); // ========================================================================= From 21f92b319d984341264f49b6ece3c5a254d32659 Mon Sep 17 00:00:00 2001 From: David Brailovsky Date: Mon, 20 Jul 2026 10:21:34 +0000 Subject: [PATCH 12/62] run generate-models --- packages/ai/src/providers/nvidia.models.ts | 4 ---- .../ai/src/providers/opencode-go.models.ts | 2 +- packages/ai/src/providers/opencode.models.ts | 2 +- .../ai/src/providers/openrouter.models.ts | 20 ++++++++----------- packages/ai/src/providers/together.models.ts | 20 ++----------------- .../src/providers/vercel-ai-gateway.models.ts | 8 -------- 6 files changed, 12 insertions(+), 44 deletions(-) diff --git a/packages/ai/src/providers/nvidia.models.ts b/packages/ai/src/providers/nvidia.models.ts index beacbf3da..ec1369190 100644 --- a/packages/ai/src/providers/nvidia.models.ts +++ b/packages/ai/src/providers/nvidia.models.ts @@ -69,10 +69,6 @@ export const NVIDIA_MODELS = values as { id: "openai/gpt-oss-20b"; provider: "nvidia"; }; - "qwen/qwen3.5-122b-a10b": Model<"openai-completions"> & { - id: "qwen/qwen3.5-122b-a10b"; - provider: "nvidia"; - }; "stepfun-ai/step-3.5-flash": Model<"openai-completions"> & { id: "stepfun-ai/step-3.5-flash"; provider: "nvidia"; diff --git a/packages/ai/src/providers/opencode-go.models.ts b/packages/ai/src/providers/opencode-go.models.ts index 4f84e56da..6ae1443bd 100644 --- a/packages/ai/src/providers/opencode-go.models.ts +++ b/packages/ai/src/providers/opencode-go.models.ts @@ -21,7 +21,7 @@ export const OPENCODE_GO_MODELS = values as { id: "glm-5.2"; provider: "opencode-go"; }; - "grok-4.5": Model<"openai-completions"> & { + "grok-4.5": Model<"openai-responses"> & { id: "grok-4.5"; provider: "opencode-go"; }; diff --git a/packages/ai/src/providers/opencode.models.ts b/packages/ai/src/providers/opencode.models.ts index fb67b4ef7..4611ee2cb 100644 --- a/packages/ai/src/providers/opencode.models.ts +++ b/packages/ai/src/providers/opencode.models.ts @@ -165,7 +165,7 @@ export const OPENCODE_MODELS = values as { id: "gpt-5.6-terra"; provider: "opencode"; }; - "grok-4.5": Model<"openai-completions"> & { + "grok-4.5": Model<"openai-responses"> & { id: "grok-4.5"; provider: "opencode"; }; diff --git a/packages/ai/src/providers/openrouter.models.ts b/packages/ai/src/providers/openrouter.models.ts index 45dd6c03b..ad687c266 100644 --- a/packages/ai/src/providers/openrouter.models.ts +++ b/packages/ai/src/providers/openrouter.models.ts @@ -297,10 +297,6 @@ export const OPENROUTER_MODELS = values as { id: "meta-llama/llama-3.3-70b-instruct"; provider: "openrouter"; }; - "meta-llama/llama-3.3-70b-instruct:free": Model<"openai-completions"> & { - id: "meta-llama/llama-3.3-70b-instruct:free"; - provider: "openrouter"; - }; "meta-llama/llama-4-maverick": Model<"openai-completions"> & { id: "meta-llama/llama-4-maverick"; provider: "openrouter"; @@ -717,6 +713,10 @@ export const OPENROUTER_MODELS = values as { id: "openrouter/auto"; provider: "openrouter"; }; + "openrouter/auto-beta": Model<"openai-completions"> & { + id: "openrouter/auto-beta"; + provider: "openrouter"; + }; "openrouter/free": Model<"openai-completions"> & { id: "openrouter/free"; provider: "openrouter"; @@ -817,10 +817,6 @@ export const OPENROUTER_MODELS = values as { id: "qwen/qwen3-coder-plus"; provider: "openrouter"; }; - "qwen/qwen3-coder:free": Model<"openai-completions"> & { - id: "qwen/qwen3-coder:free"; - provider: "openrouter"; - }; "qwen/qwen3-max": Model<"openai-completions"> & { id: "qwen/qwen3-max"; provider: "openrouter"; @@ -833,10 +829,6 @@ export const OPENROUTER_MODELS = values as { id: "qwen/qwen3-next-80b-a3b-instruct"; provider: "openrouter"; }; - "qwen/qwen3-next-80b-a3b-instruct:free": Model<"openai-completions"> & { - id: "qwen/qwen3-next-80b-a3b-instruct:free"; - provider: "openrouter"; - }; "qwen/qwen3-next-80b-a3b-thinking": Model<"openai-completions"> & { id: "qwen/qwen3-next-80b-a3b-thinking"; provider: "openrouter"; @@ -969,6 +961,10 @@ export const OPENROUTER_MODELS = values as { id: "thedrummer/unslopnemo-12b"; provider: "openrouter"; }; + "thinkingmachines/inkling": Model<"openai-completions"> & { + id: "thinkingmachines/inkling"; + provider: "openrouter"; + }; "upstage/solar-pro-3": Model<"openai-completions"> & { id: "upstage/solar-pro-3"; provider: "openrouter"; diff --git a/packages/ai/src/providers/together.models.ts b/packages/ai/src/providers/together.models.ts index b2b494c8b..3ae1225cb 100644 --- a/packages/ai/src/providers/together.models.ts +++ b/packages/ai/src/providers/together.models.ts @@ -17,14 +17,6 @@ export const TOGETHER_MODELS = values as { id: "Qwen/Qwen2.5-7B-Instruct-Turbo"; provider: "together"; }; - "Qwen/Qwen3-235B-A22B-Instruct-2507-tput": Model<"openai-completions"> & { - id: "Qwen/Qwen3-235B-A22B-Instruct-2507-tput"; - provider: "together"; - }; - "Qwen/Qwen3.5-397B-A17B": Model<"openai-completions"> & { - id: "Qwen/Qwen3.5-397B-A17B"; - provider: "together"; - }; "Qwen/Qwen3.5-9B": Model<"openai-completions"> & { id: "Qwen/Qwen3.5-9B"; provider: "together"; @@ -41,10 +33,6 @@ export const TOGETHER_MODELS = values as { id: "deepseek-ai/DeepSeek-V4-Pro"; provider: "together"; }; - "essentialai/Rnj-1-Instruct": Model<"openai-completions"> & { - id: "essentialai/Rnj-1-Instruct"; - provider: "together"; - }; "google/gemma-4-31B-it": Model<"openai-completions"> & { id: "google/gemma-4-31B-it"; provider: "together"; @@ -73,12 +61,8 @@ export const TOGETHER_MODELS = values as { id: "openai/gpt-oss-20b"; provider: "together"; }; - "zai-org/GLM-5": Model<"openai-completions"> & { - id: "zai-org/GLM-5"; - provider: "together"; - }; - "zai-org/GLM-5.1": Model<"openai-completions"> & { - id: "zai-org/GLM-5.1"; + "thinkingmachines/Inkling": Model<"openai-completions"> & { + id: "thinkingmachines/Inkling"; provider: "together"; }; "zai-org/GLM-5.2": Model<"openai-completions"> & { diff --git a/packages/ai/src/providers/vercel-ai-gateway.models.ts b/packages/ai/src/providers/vercel-ai-gateway.models.ts index 11bfc6f3e..cabae8800 100644 --- a/packages/ai/src/providers/vercel-ai-gateway.models.ts +++ b/packages/ai/src/providers/vercel-ai-gateway.models.ts @@ -309,14 +309,6 @@ export const VERCEL_AI_GATEWAY_MODELS = values as { id: "meta/llama-3.1-8b"; provider: "vercel-ai-gateway"; }; - "meta/llama-3.2-11b": Model<"anthropic-messages"> & { - id: "meta/llama-3.2-11b"; - provider: "vercel-ai-gateway"; - }; - "meta/llama-3.2-90b": Model<"anthropic-messages"> & { - id: "meta/llama-3.2-90b"; - provider: "vercel-ai-gateway"; - }; "meta/llama-3.3-70b": Model<"anthropic-messages"> & { id: "meta/llama-3.3-70b"; provider: "vercel-ai-gateway"; From 84fb69c95868419ef15d8ea461773f40087af922 Mon Sep 17 00:00:00 2001 From: David Brailovsky Date: Mon, 20 Jul 2026 10:21:58 +0000 Subject: [PATCH 13/62] update npm shrinkwrap --- packages/coding-agent/install-lock/package-lock.json | 2 +- packages/coding-agent/npm-shrinkwrap.json | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/coding-agent/install-lock/package-lock.json b/packages/coding-agent/install-lock/package-lock.json index 8282baf66..31c29e590 100644 --- a/packages/coding-agent/install-lock/package-lock.json +++ b/packages/coding-agent/install-lock/package-lock.json @@ -481,7 +481,7 @@ "typebox": "1.1.38" }, "bin": { - "pi-ai": "./dist/cli.js" + "pi-ai": "dist/cli.js" }, "engines": { "node": ">=22.19.0" diff --git a/packages/coding-agent/npm-shrinkwrap.json b/packages/coding-agent/npm-shrinkwrap.json index 669d0feb9..fa901aeb5 100644 --- a/packages/coding-agent/npm-shrinkwrap.json +++ b/packages/coding-agent/npm-shrinkwrap.json @@ -505,7 +505,7 @@ "typebox": "1.1.38" }, "bin": { - "pi-ai": "./dist/cli.js" + "pi-ai": "dist/cli.js" }, "engines": { "node": ">=22.19.0" From 35f12c8c729010d8b77cbb161ba19f30018fd1e8 Mon Sep 17 00:00:00 2001 From: Aadish Verma Date: Mon, 20 Jul 2026 03:22:22 -0700 Subject: [PATCH 14/62] fix: gpt 5.6 context window (#6853) --- packages/ai/CHANGELOG.md | 1 + packages/ai/scripts/generate-models.ts | 4 ++-- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index ed86ca720..7739597b3 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -13,6 +13,7 @@ - Fixed GitHub Copilot long-context pricing tiers in generated model metadata ([#6668](https://github.com/earendil-works/pi/issues/6668)). - Fixed Kimi Coding subscription models to report API-equivalent implied costs when models.dev reports zero pricing. - Fixed OpenAI Responses early stream endings to be classified as retryable provider errors ([#6727](https://github.com/earendil-works/pi/issues/6727)). +- Fixed GPT-5.6 Codex models to default to the 272K context window, avoiding automatic long-context pricing ([#6838](https://github.com/earendil-works/pi/issues/6838)). ## [0.80.10] - 2026-07-16 diff --git a/packages/ai/scripts/generate-models.ts b/packages/ai/scripts/generate-models.ts index 38c1464ea..95d11332f 100644 --- a/packages/ai/scripts/generate-models.ts +++ b/packages/ai/scripts/generate-models.ts @@ -2112,10 +2112,10 @@ async function generateModels() { // OpenAI Codex (ChatGPT OAuth) models // NOTE: These are not fetched from models.dev; we keep a small, explicit list to avoid aliases. - // Older model limits are based on observed server behavior; GPT-5.6 follows Codex's 372k catalog limit. + // Older model limits are based on observed server behavior; GPT-5.6 follows Codex's 272k catalog limit (formerly 372k). const CODEX_BASE_URL = "https://chatgpt.com/backend-api"; const CODEX_CONTEXT = 272000; - const CODEX_GPT_56_CONTEXT = 372000; + const CODEX_GPT_56_CONTEXT = 272000; const CODEX_SPARK_CONTEXT = 128000; const CODEX_MAX_TOKENS = 128000; const codexModels: Model<"openai-codex-responses">[] = [ From 9e7fce70a3186e3a25074784b1db00facc816299 Mon Sep 17 00:00:00 2001 From: David Brailovsky Date: Mon, 20 Jul 2026 13:09:12 +0200 Subject: [PATCH 15/62] export missing message and tool execution event types (#6772) fixes #6687 --- packages/coding-agent/src/index.ts | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/packages/coding-agent/src/index.ts b/packages/coding-agent/src/index.ts index 26757fb83..196fc8ca5 100644 --- a/packages/coding-agent/src/index.ts +++ b/packages/coding-agent/src/index.ts @@ -99,8 +99,11 @@ export type { KeybindingsManager, LoadExtensionsResult, LsToolCallEvent, + MessageEndEvent, MessageRenderer, MessageRenderOptions, + MessageStartEvent, + MessageUpdateEvent, ProjectTrustContext, ProjectTrustEvent, ProjectTrustEventDecision, @@ -128,7 +131,10 @@ export type { ToolCallEvent, ToolCallEventResult, ToolDefinition, + ToolExecutionEndEvent, ToolExecutionMode, + ToolExecutionStartEvent, + ToolExecutionUpdateEvent, ToolInfo, ToolRenderResultOptions, ToolResultEvent, From d9f7f814730998f191ad6c32bbd178e0834cae18 Mon Sep 17 00:00:00 2001 From: Cristina Poncela Cubeiro <140309543+cristinaponcela@users.noreply.github.com> Date: Mon, 20 Jul 2026 13:51:16 +0200 Subject: [PATCH 16/62] fix: tool_call_id error when switching (#6854) --- packages/ai/src/api/openai-completions.ts | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/packages/ai/src/api/openai-completions.ts b/packages/ai/src/api/openai-completions.ts index 610017772..01eefcb05 100644 --- a/packages/ai/src/api/openai-completions.ts +++ b/packages/ai/src/api/openai-completions.ts @@ -34,6 +34,7 @@ import type { } from "../types.ts"; import { formatProviderError, normalizeProviderError } from "../utils/error-body.ts"; import { AssistantMessageEventStream } from "../utils/event-stream.ts"; +import { shortHash } from "../utils/hash.ts"; import { headersToRecord } from "../utils/headers.ts"; import { parseStreamingJson } from "../utils/json-parse.ts"; import { getProviderEnvValue } from "../utils/provider-env.ts"; @@ -895,10 +896,21 @@ export function convertMessages( // Format: {call_id}|{id} where {id} can be 400+ chars with special chars (+, /, =) // These come from providers like github-copilot, openai-codex, opencode // Extract just the call_id part and normalize it + // Multiple tool calls in the same turn can share call_id but differ by item_id. + // Preserve item-level uniqueness when replaying into Chat Completions, which + // requires distinct tool call ids. if (id.includes("|")) { - const [callId] = id.split("|"); // Sanitize to allowed chars and truncate to 40 chars (OpenAI limit) - return callId.replace(/[^a-zA-Z0-9_-]/g, "_").slice(0, 40); + const separatorIndex = id.indexOf("|"); + const callId = id.slice(0, separatorIndex).replace(/[^a-zA-Z0-9_-]/g, "_"); + const itemId = id.slice(separatorIndex + 1).replace(/[^a-zA-Z0-9_-]/g, "_"); + const combinedId = itemId.length > 0 ? `${callId}_${itemId}` : callId; + if (combinedId.length <= 40) { + return combinedId; + } + const hash = shortHash(id).slice(0, 8); + const prefix = callId.slice(0, Math.max(1, 40 - hash.length - 1)); + return `${prefix}_${hash}`; } if (model.provider === "openai") return id.length > 40 ? id.slice(0, 40) : id; From bbb91fa8ae56ecbf0e6ad4c3fe3bdc4a7f8bc696 Mon Sep 17 00:00:00 2001 From: QuintinShaw Date: Mon, 20 Jul 2026 19:53:30 +0800 Subject: [PATCH 17/62] feat(ai): add Qwen Token Plan as built-in provider (#6858) Add Alibaba Cloud Model Studio Token Plan subscription service as two built-in API-key providers: qwen-token-plan (international, Singapore) and qwen-token-plan-cn (China, Beijing). Each provider exposes 15 text-generation models (Qwen, DeepSeek, GLM, Kimi, MiniMax) via the OpenAI-compatible endpoint with DashScope enable_thinking support. Model metadata is sourced from models.dev; qwen3.8-max-preview is hardcoded until models.dev includes it. Also fixes kimi-coding test references (k2p7 -> kimi-for-coding) after models.dev catalog update picked up by generate-models. Closes #6850 --- packages/ai/README.md | 2 + packages/ai/scripts/generate-models.ts | 74 +++++++++++++++++++ packages/ai/src/env-api-keys.ts | 2 + packages/ai/src/models.generated.ts | 4 + packages/ai/src/providers/all.ts | 4 + .../providers/qwen-token-plan-cn.models.ts | 68 +++++++++++++++++ .../ai/src/providers/qwen-token-plan-cn.ts | 15 ++++ .../src/providers/qwen-token-plan.models.ts | 68 +++++++++++++++++ packages/ai/src/providers/qwen-token-plan.ts | 15 ++++ packages/ai/src/types.ts | 2 + packages/ai/src/utils/overflow.ts | 3 + packages/ai/test/abort.test.ts | 24 ++++++ packages/ai/test/context-overflow.test.ts | 24 ++++++ .../ai/test/cross-provider-handoff.test.ts | 3 + packages/ai/test/empty.test.ts | 40 ++++++++++ packages/ai/test/image-tool-result.test.ts | 24 ++++++ .../openai-completions-tool-choice.test.ts | 12 +++ .../ai/test/qwen-token-plan-models.test.ts | 38 ++++++++++ packages/ai/test/stream.test.ts | 59 +++++++++++++++ packages/ai/test/tokens.test.ts | 16 ++++ .../ai/test/tool-call-without-result.test.ts | 16 ++++ packages/ai/test/total-tokens.test.ts | 50 +++++++++++++ packages/ai/test/unicode-surrogate.test.ts | 32 ++++++++ packages/coding-agent/docs/providers.md | 4 + packages/coding-agent/src/cli/args.ts | 2 + .../coding-agent/src/core/model-resolver.ts | 2 + test.sh | 2 + 27 files changed, 605 insertions(+) create mode 100644 packages/ai/src/providers/qwen-token-plan-cn.models.ts create mode 100644 packages/ai/src/providers/qwen-token-plan-cn.ts create mode 100644 packages/ai/src/providers/qwen-token-plan.models.ts create mode 100644 packages/ai/src/providers/qwen-token-plan.ts create mode 100644 packages/ai/test/qwen-token-plan-models.test.ts diff --git a/packages/ai/README.md b/packages/ai/README.md index ee047aeb4..3c791621d 100644 --- a/packages/ai/README.md +++ b/packages/ai/README.md @@ -434,6 +434,8 @@ Built-in providers resolve these env vars (Node.js; in browsers pass `apiKey` ex | Hugging Face | `HF_TOKEN` | | OpenCode Zen / OpenCode Go | `OPENCODE_API_KEY` | | Kimi For Coding | `KIMI_API_KEY` | +| Qwen Token Plan | `QWEN_TOKEN_PLAN_API_KEY` | +| Qwen Token Plan (China) | `QWEN_TOKEN_PLAN_CN_API_KEY` | | Xiaomi MiMo (API billing) | `XIAOMI_API_KEY` | | Xiaomi MiMo Token Plan (China) | `XIAOMI_TOKEN_PLAN_CN_API_KEY` | | Xiaomi MiMo Token Plan (Amsterdam) | `XIAOMI_TOKEN_PLAN_AMS_API_KEY` | diff --git a/packages/ai/scripts/generate-models.ts b/packages/ai/scripts/generate-models.ts index 95d11332f..40e7275c7 100644 --- a/packages/ai/scripts/generate-models.ts +++ b/packages/ai/scripts/generate-models.ts @@ -1825,6 +1825,57 @@ async function loadModelsDevData(): Promise[]> { } } + // Process Alibaba Cloud Model Studio Token Plan models + // Two regions (international / cn) with identical catalogs, separate + // endpoints and API keys (sk-sp- prefix). models.dev keys are + // "alibaba-token-plan[-cn]"; pi exposes them as "qwen-token-plan[-cn]". + const qwenTokenPlanCompat: OpenAICompletionsCompat = { + thinkingFormat: "qwen", + supportsDeveloperRole: false, + supportsStore: false, + }; + const qwenTokenPlanVariants = [ + { + source: "alibaba-token-plan", + provider: "qwen-token-plan", + baseUrl: "https://token-plan.ap-southeast-1.maas.aliyuncs.com/compatible-mode/v1", + }, + { + source: "alibaba-token-plan-cn", + provider: "qwen-token-plan-cn", + baseUrl: "https://token-plan.cn-beijing.maas.aliyuncs.com/compatible-mode/v1", + }, + ] as const; + + for (const { source, provider, baseUrl } of qwenTokenPlanVariants) { + const providerModels = data[source]?.models; + if (!providerModels) continue; + + for (const [modelId, model] of Object.entries(providerModels)) { + const m = model as ModelsDevModel; + if (m.tool_call !== true) continue; + + models.push({ + id: modelId, + name: m.name || modelId, + api: "openai-completions", + provider, + baseUrl, + compat: qwenTokenPlanCompat, + reasoning: m.reasoning === true, + input: m.modalities?.input?.includes("image") ? ["text", "image"] : ["text"], + cost: { + input: m.cost?.input || 0, + output: m.cost?.output || 0, + cacheRead: m.cost?.cache_read || 0, + cacheWrite: m.cost?.cache_write || 0, + }, + contextWindow: m.limit?.context || 4096, + maxTokens: m.limit?.output || 4096, + }); + } + } + console.log(`Loaded ${models.length} tool-capable models from models.dev`); return models; } catch (error) { @@ -2227,6 +2278,29 @@ async function generateModels() { }); } + // Add qwen3.8-max-preview to Qwen Token Plan providers until models.dev includes it + for (const qwenTpProvider of ["qwen-token-plan", "qwen-token-plan-cn"] as const) { + if (!allModels.some((m) => m.provider === qwenTpProvider && m.id === "qwen3.8-max-preview")) { + const baseUrl = + qwenTpProvider === "qwen-token-plan" + ? "https://token-plan.ap-southeast-1.maas.aliyuncs.com/compatible-mode/v1" + : "https://token-plan.cn-beijing.maas.aliyuncs.com/compatible-mode/v1"; + allModels.push({ + id: "qwen3.8-max-preview", + name: "Qwen3.8 Max Preview", + api: "openai-completions", + provider: qwenTpProvider, + baseUrl, + compat: { thinkingFormat: "qwen", supportsDeveloperRole: false, supportsStore: false } satisfies OpenAICompletionsCompat, + reasoning: true, + input: ["text", "image"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 1000000, + maxTokens: 65536, + }); + } + } + // Add "auto" alias for openrouter/auto if (!allModels.some(m => m.provider === "openrouter" && m.id === "auto")) { allModels.push({ diff --git a/packages/ai/src/env-api-keys.ts b/packages/ai/src/env-api-keys.ts index 7a4f8f08f..e6fa23dca 100644 --- a/packages/ai/src/env-api-keys.ts +++ b/packages/ai/src/env-api-keys.ts @@ -73,6 +73,8 @@ function getApiKeyEnvVars(provider: string): readonly string[] | undefined { const envMap: Record = { "ant-ling": "ANT_LING_API_KEY", + "qwen-token-plan": "QWEN_TOKEN_PLAN_API_KEY", + "qwen-token-plan-cn": "QWEN_TOKEN_PLAN_CN_API_KEY", openai: "OPENAI_API_KEY", "azure-openai-responses": "AZURE_OPENAI_API_KEY", nvidia: "NVIDIA_API_KEY", diff --git a/packages/ai/src/models.generated.ts b/packages/ai/src/models.generated.ts index 0129ddeee..035c50b03 100644 --- a/packages/ai/src/models.generated.ts +++ b/packages/ai/src/models.generated.ts @@ -27,6 +27,8 @@ import { OPENAI_CODEX_MODELS } from "./providers/openai-codex.models.ts"; import { OPENCODE_MODELS } from "./providers/opencode.models.ts"; import { OPENCODE_GO_MODELS } from "./providers/opencode-go.models.ts"; import { OPENROUTER_MODELS } from "./providers/openrouter.models.ts"; +import { QWEN_TOKEN_PLAN_MODELS } from "./providers/qwen-token-plan.models.ts"; +import { QWEN_TOKEN_PLAN_CN_MODELS } from "./providers/qwen-token-plan-cn.models.ts"; import { TOGETHER_MODELS } from "./providers/together.models.ts"; import { VERCEL_AI_GATEWAY_MODELS } from "./providers/vercel-ai-gateway.models.ts"; import { XAI_MODELS } from "./providers/xai.models.ts"; @@ -64,6 +66,8 @@ export const MODELS = { "opencode": OPENCODE_MODELS, "opencode-go": OPENCODE_GO_MODELS, "openrouter": OPENROUTER_MODELS, + "qwen-token-plan": QWEN_TOKEN_PLAN_MODELS, + "qwen-token-plan-cn": QWEN_TOKEN_PLAN_CN_MODELS, "together": TOGETHER_MODELS, "vercel-ai-gateway": VERCEL_AI_GATEWAY_MODELS, "xai": XAI_MODELS, diff --git a/packages/ai/src/providers/all.ts b/packages/ai/src/providers/all.ts index 1261ed099..fd7bea58d 100644 --- a/packages/ai/src/providers/all.ts +++ b/packages/ai/src/providers/all.ts @@ -29,6 +29,8 @@ import { opencodeProvider } from "./opencode.ts"; import { opencodeGoProvider } from "./opencode-go.ts"; import { openrouterProvider } from "./openrouter.ts"; import { openrouterImagesProvider } from "./openrouter-images.ts"; +import { qwenTokenPlanProvider } from "./qwen-token-plan.ts"; +import { qwenTokenPlanCnProvider } from "./qwen-token-plan-cn.ts"; import { radiusProvider } from "./radius.ts"; import { togetherProvider } from "./together.ts"; import { vercelAIGatewayProvider } from "./vercel-ai-gateway.ts"; @@ -103,6 +105,8 @@ export function builtinProviders(): Provider[] { opencodeProvider(), opencodeGoProvider(), openrouterProvider(), + qwenTokenPlanProvider(), + qwenTokenPlanCnProvider(), radiusProvider(), togetherProvider(), vercelAIGatewayProvider(), diff --git a/packages/ai/src/providers/qwen-token-plan-cn.models.ts b/packages/ai/src/providers/qwen-token-plan-cn.models.ts new file mode 100644 index 000000000..84976ef94 --- /dev/null +++ b/packages/ai/src/providers/qwen-token-plan-cn.models.ts @@ -0,0 +1,68 @@ +// This file is auto-generated by scripts/generate-models.ts +// Do not edit manually - run 'npm run generate-models' to update + +import values from "./data/qwen-token-plan-cn.json" with { type: "json" }; +import type { Model } from "../types.ts"; + +export const QWEN_TOKEN_PLAN_CN_MODELS = values as { + "MiniMax-M2.5": Model<"openai-completions"> & { + id: "MiniMax-M2.5"; + provider: "qwen-token-plan-cn"; + }; + "deepseek-v3.2": Model<"openai-completions"> & { + id: "deepseek-v3.2"; + provider: "qwen-token-plan-cn"; + }; + "deepseek-v4-flash": Model<"openai-completions"> & { + id: "deepseek-v4-flash"; + provider: "qwen-token-plan-cn"; + }; + "deepseek-v4-pro": Model<"openai-completions"> & { + id: "deepseek-v4-pro"; + provider: "qwen-token-plan-cn"; + }; + "glm-5": Model<"openai-completions"> & { + id: "glm-5"; + provider: "qwen-token-plan-cn"; + }; + "glm-5.1": Model<"openai-completions"> & { + id: "glm-5.1"; + provider: "qwen-token-plan-cn"; + }; + "glm-5.2": Model<"openai-completions"> & { + id: "glm-5.2"; + provider: "qwen-token-plan-cn"; + }; + "kimi-k2.5": Model<"openai-completions"> & { + id: "kimi-k2.5"; + provider: "qwen-token-plan-cn"; + }; + "kimi-k2.6": Model<"openai-completions"> & { + id: "kimi-k2.6"; + provider: "qwen-token-plan-cn"; + }; + "kimi-k2.7-code": Model<"openai-completions"> & { + id: "kimi-k2.7-code"; + provider: "qwen-token-plan-cn"; + }; + "qwen3.6-flash": Model<"openai-completions"> & { + id: "qwen3.6-flash"; + provider: "qwen-token-plan-cn"; + }; + "qwen3.6-plus": Model<"openai-completions"> & { + id: "qwen3.6-plus"; + provider: "qwen-token-plan-cn"; + }; + "qwen3.7-max": Model<"openai-completions"> & { + id: "qwen3.7-max"; + provider: "qwen-token-plan-cn"; + }; + "qwen3.7-plus": Model<"openai-completions"> & { + id: "qwen3.7-plus"; + provider: "qwen-token-plan-cn"; + }; + "qwen3.8-max-preview": Model<"openai-completions"> & { + id: "qwen3.8-max-preview"; + provider: "qwen-token-plan-cn"; + }; +}; diff --git a/packages/ai/src/providers/qwen-token-plan-cn.ts b/packages/ai/src/providers/qwen-token-plan-cn.ts new file mode 100644 index 000000000..259b5a7f0 --- /dev/null +++ b/packages/ai/src/providers/qwen-token-plan-cn.ts @@ -0,0 +1,15 @@ +import { openAICompletionsApi } from "../api/openai-completions.lazy.ts"; +import { envApiKeyAuth } from "../auth/helpers.ts"; +import { createProvider, type Provider } from "../models.ts"; +import { QWEN_TOKEN_PLAN_CN_MODELS } from "./qwen-token-plan-cn.models.ts"; + +export function qwenTokenPlanCnProvider(): Provider<"openai-completions"> { + return createProvider({ + id: "qwen-token-plan-cn", + name: "Qwen Token Plan CN", + baseUrl: "https://token-plan.cn-beijing.maas.aliyuncs.com/compatible-mode/v1", + auth: { apiKey: envApiKeyAuth("Qwen Token Plan CN API key", ["QWEN_TOKEN_PLAN_CN_API_KEY"]) }, + models: Object.values(QWEN_TOKEN_PLAN_CN_MODELS), + api: openAICompletionsApi(), + }); +} diff --git a/packages/ai/src/providers/qwen-token-plan.models.ts b/packages/ai/src/providers/qwen-token-plan.models.ts new file mode 100644 index 000000000..e8f91f012 --- /dev/null +++ b/packages/ai/src/providers/qwen-token-plan.models.ts @@ -0,0 +1,68 @@ +// This file is auto-generated by scripts/generate-models.ts +// Do not edit manually - run 'npm run generate-models' to update + +import values from "./data/qwen-token-plan.json" with { type: "json" }; +import type { Model } from "../types.ts"; + +export const QWEN_TOKEN_PLAN_MODELS = values as { + "MiniMax-M2.5": Model<"openai-completions"> & { + id: "MiniMax-M2.5"; + provider: "qwen-token-plan"; + }; + "deepseek-v3.2": Model<"openai-completions"> & { + id: "deepseek-v3.2"; + provider: "qwen-token-plan"; + }; + "deepseek-v4-flash": Model<"openai-completions"> & { + id: "deepseek-v4-flash"; + provider: "qwen-token-plan"; + }; + "deepseek-v4-pro": Model<"openai-completions"> & { + id: "deepseek-v4-pro"; + provider: "qwen-token-plan"; + }; + "glm-5": Model<"openai-completions"> & { + id: "glm-5"; + provider: "qwen-token-plan"; + }; + "glm-5.1": Model<"openai-completions"> & { + id: "glm-5.1"; + provider: "qwen-token-plan"; + }; + "glm-5.2": Model<"openai-completions"> & { + id: "glm-5.2"; + provider: "qwen-token-plan"; + }; + "kimi-k2.5": Model<"openai-completions"> & { + id: "kimi-k2.5"; + provider: "qwen-token-plan"; + }; + "kimi-k2.6": Model<"openai-completions"> & { + id: "kimi-k2.6"; + provider: "qwen-token-plan"; + }; + "kimi-k2.7-code": Model<"openai-completions"> & { + id: "kimi-k2.7-code"; + provider: "qwen-token-plan"; + }; + "qwen3.6-flash": Model<"openai-completions"> & { + id: "qwen3.6-flash"; + provider: "qwen-token-plan"; + }; + "qwen3.6-plus": Model<"openai-completions"> & { + id: "qwen3.6-plus"; + provider: "qwen-token-plan"; + }; + "qwen3.7-max": Model<"openai-completions"> & { + id: "qwen3.7-max"; + provider: "qwen-token-plan"; + }; + "qwen3.7-plus": Model<"openai-completions"> & { + id: "qwen3.7-plus"; + provider: "qwen-token-plan"; + }; + "qwen3.8-max-preview": Model<"openai-completions"> & { + id: "qwen3.8-max-preview"; + provider: "qwen-token-plan"; + }; +}; diff --git a/packages/ai/src/providers/qwen-token-plan.ts b/packages/ai/src/providers/qwen-token-plan.ts new file mode 100644 index 000000000..295560ba2 --- /dev/null +++ b/packages/ai/src/providers/qwen-token-plan.ts @@ -0,0 +1,15 @@ +import { openAICompletionsApi } from "../api/openai-completions.lazy.ts"; +import { envApiKeyAuth } from "../auth/helpers.ts"; +import { createProvider, type Provider } from "../models.ts"; +import { QWEN_TOKEN_PLAN_MODELS } from "./qwen-token-plan.models.ts"; + +export function qwenTokenPlanProvider(): Provider<"openai-completions"> { + return createProvider({ + id: "qwen-token-plan", + name: "Qwen Token Plan", + baseUrl: "https://token-plan.ap-southeast-1.maas.aliyuncs.com/compatible-mode/v1", + auth: { apiKey: envApiKeyAuth("Qwen Token Plan API key", ["QWEN_TOKEN_PLAN_API_KEY"]) }, + models: Object.values(QWEN_TOKEN_PLAN_MODELS), + api: openAICompletionsApi(), + }); +} diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index c123fd897..9fa1ab20f 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -64,6 +64,8 @@ export type KnownProvider = | "kimi-coding" | "cloudflare-workers-ai" | "cloudflare-ai-gateway" + | "qwen-token-plan" + | "qwen-token-plan-cn" | "xiaomi" | "xiaomi-token-plan-cn" | "xiaomi-token-plan-ams" diff --git a/packages/ai/src/utils/overflow.ts b/packages/ai/src/utils/overflow.ts index c10ee80d3..6554b0b03 100644 --- a/packages/ai/src/utils/overflow.ts +++ b/packages/ai/src/utils/overflow.ts @@ -31,6 +31,7 @@ import type { AssistantMessage } from "../types.ts"; * - Xiaomi MiMo: Truncates input to fill contextWindow exactly, then returns finish_reason "length" * with output=0 (no room left to generate). Detected via stopReason "length" + zero output + * input filling the context window. + * - DashScope/Qwen: "Range of input length should be [1, X]" (HTTP 400 invalid_parameter_error) * - Ollama: Some deployments truncate silently, others return errors like "prompt too long; exceeded max context length by X tokens" */ const OVERFLOW_PATTERNS = [ @@ -54,6 +55,7 @@ const OVERFLOW_PATTERNS = [ /prompt has [\d,]+ tokens?, but the configured context size is [\d,]+ tokens?/i, // DS4 server /model_context_window_exceeded/i, // z.ai non-standard finish_reason surfaced as error text /prompt too long; exceeded (?:max )?context length/i, // Ollama explicit overflow error + /range of input length should be/i, // DashScope / Qwen Token Plan /context[_ ]length[_ ]exceeded/i, // Generic fallback /too many tokens/i, // Generic fallback /token limit exceeded/i, // Generic fallback @@ -101,6 +103,7 @@ const NON_OVERFLOW_PATTERNS = [ * - LM Studio: "greater than the context length" * - Kimi For Coding: "exceeded model token limit: X (requested: Y)" * - DS4: "Prompt has X tokens, but the configured context size is Y tokens" + * - DashScope/Qwen: "Range of input length should be [1, X]" * * **Unreliable detection:** * - z.ai: Sometimes accepts overflow silently (detectable via usage.input > contextWindow), diff --git a/packages/ai/test/abort.test.ts b/packages/ai/test/abort.test.ts index c79cf540b..df1904db9 100644 --- a/packages/ai/test/abort.test.ts +++ b/packages/ai/test/abort.test.ts @@ -249,6 +249,30 @@ describe("AI Providers Abort Tests", () => { }); }); + describe.skipIf(!process.env.QWEN_TOKEN_PLAN_API_KEY)("Qwen Token Plan Provider Abort", () => { + const llm = getModel("qwen-token-plan", "qwen3.7-max"); + + it("should abort mid-stream", { retry: 3 }, async () => { + await testAbortSignal(llm); + }); + + it("should handle immediate abort", { retry: 3 }, async () => { + await testImmediateAbort(llm); + }); + }); + + describe.skipIf(!process.env.QWEN_TOKEN_PLAN_CN_API_KEY)("Qwen Token Plan (CN) Provider Abort", () => { + const llm = getModel("qwen-token-plan-cn", "qwen3.7-max"); + + it("should abort mid-stream", { retry: 3 }, async () => { + await testAbortSignal(llm); + }); + + it("should handle immediate abort", { retry: 3 }, async () => { + await testImmediateAbort(llm); + }); + }); + describe.skipIf(!process.env.KIMI_API_KEY)("Kimi For Coding Provider Abort", () => { const llm = getModel("kimi-coding", "kimi-for-coding"); diff --git a/packages/ai/test/context-overflow.test.ts b/packages/ai/test/context-overflow.test.ts index 2a8b62845..ab5fd9274 100644 --- a/packages/ai/test/context-overflow.test.ts +++ b/packages/ai/test/context-overflow.test.ts @@ -466,6 +466,30 @@ describe("Context overflow error handling", () => { }, 120000); }); + describe.skipIf(!process.env.QWEN_TOKEN_PLAN_API_KEY)("Qwen Token Plan", () => { + it("qwen3.7-max - should detect overflow via isContextOverflow", async () => { + const model = getModel("qwen-token-plan", "qwen3.7-max"); + const result = await testContextOverflow(model, process.env.QWEN_TOKEN_PLAN_API_KEY!); + logResult(result); + + expect(result.stopReason).toBe("error"); + expect(result.errorMessage).toMatch(/input length/i); + expect(isContextOverflow(result.response, model.contextWindow)).toBe(true); + }, 120000); + }); + + describe.skipIf(!process.env.QWEN_TOKEN_PLAN_CN_API_KEY)("Qwen Token Plan (CN)", () => { + it("qwen3.7-max - should detect overflow via isContextOverflow", async () => { + const model = getModel("qwen-token-plan-cn", "qwen3.7-max"); + const result = await testContextOverflow(model, process.env.QWEN_TOKEN_PLAN_CN_API_KEY!); + logResult(result); + + expect(result.stopReason).toBe("error"); + expect(result.errorMessage).toMatch(/input length/i); + expect(isContextOverflow(result.response, model.contextWindow)).toBe(true); + }, 120000); + }); + // ============================================================================= // Kimi For Coding // ============================================================================= diff --git a/packages/ai/test/cross-provider-handoff.test.ts b/packages/ai/test/cross-provider-handoff.test.ts index 08da44de3..f81513860 100644 --- a/packages/ai/test/cross-provider-handoff.test.ts +++ b/packages/ai/test/cross-provider-handoff.test.ts @@ -130,6 +130,9 @@ const PROVIDER_MODEL_PAIRS: ProviderModelPair[] = [ { provider: "xiaomi-token-plan-cn", model: "mimo-v2.5-pro", label: "xiaomi-token-plan-cn-mimo-v2.5-pro" }, { provider: "xiaomi-token-plan-ams", model: "mimo-v2.5-pro", label: "xiaomi-token-plan-ams-mimo-v2.5-pro" }, { provider: "xiaomi-token-plan-sgp", model: "mimo-v2.5-pro", label: "xiaomi-token-plan-sgp-mimo-v2.5-pro" }, + // Qwen Token Plan + { provider: "qwen-token-plan", model: "qwen3.7-max", label: "qwen-token-plan-qwen3.7-max" }, + { provider: "qwen-token-plan-cn", model: "qwen3.7-max", label: "qwen-token-plan-cn-qwen3.7-max" }, ]; // Cached context structure diff --git a/packages/ai/test/empty.test.ts b/packages/ai/test/empty.test.ts index 502764c95..0a0705f51 100644 --- a/packages/ai/test/empty.test.ts +++ b/packages/ai/test/empty.test.ts @@ -535,6 +535,46 @@ describe("AI Providers Empty Message Tests", () => { }, ); + describe.skipIf(!process.env.QWEN_TOKEN_PLAN_API_KEY)("Qwen Token Plan Provider Empty Messages", () => { + const llm = getModel("qwen-token-plan", "qwen3.7-max"); + + it("should handle empty content array", { retry: 3, timeout: 30000 }, async () => { + await testEmptyMessage(llm); + }); + + it("should handle empty string content", { retry: 3, timeout: 30000 }, async () => { + await testEmptyStringMessage(llm); + }); + + it("should handle whitespace-only content", { retry: 3, timeout: 30000 }, async () => { + await testWhitespaceOnlyMessage(llm); + }); + + it("should handle empty assistant message in conversation", { retry: 3, timeout: 30000 }, async () => { + await testEmptyAssistantMessage(llm); + }); + }); + + describe.skipIf(!process.env.QWEN_TOKEN_PLAN_CN_API_KEY)("Qwen Token Plan (CN) Provider Empty Messages", () => { + const llm = getModel("qwen-token-plan-cn", "qwen3.7-max"); + + it("should handle empty content array", { retry: 3, timeout: 30000 }, async () => { + await testEmptyMessage(llm); + }); + + it("should handle empty string content", { retry: 3, timeout: 30000 }, async () => { + await testEmptyStringMessage(llm); + }); + + it("should handle whitespace-only content", { retry: 3, timeout: 30000 }, async () => { + await testWhitespaceOnlyMessage(llm); + }); + + it("should handle empty assistant message in conversation", { retry: 3, timeout: 30000 }, async () => { + await testEmptyAssistantMessage(llm); + }); + }); + describe.skipIf(!process.env.KIMI_API_KEY)("Kimi For Coding Provider Empty Messages", () => { const llm = getModel("kimi-coding", "kimi-for-coding"); diff --git a/packages/ai/test/image-tool-result.test.ts b/packages/ai/test/image-tool-result.test.ts index 946a443ad..150f3d8b7 100644 --- a/packages/ai/test/image-tool-result.test.ts +++ b/packages/ai/test/image-tool-result.test.ts @@ -381,6 +381,30 @@ describe("Tool Results with Images", () => { }, ); + describe.skipIf(!process.env.QWEN_TOKEN_PLAN_API_KEY)("Qwen Token Plan Provider (qwen3.7-max)", () => { + const llm = getModel("qwen-token-plan", "qwen3.7-max"); + + it("should handle tool result with only image", { retry: 3, timeout: 30000 }, async () => { + await handleToolWithImageResult(llm); + }); + + it("should handle tool result with text and image", { retry: 3, timeout: 30000 }, async () => { + await handleToolWithTextAndImageResult(llm); + }); + }); + + describe.skipIf(!process.env.QWEN_TOKEN_PLAN_CN_API_KEY)("Qwen Token Plan (CN) Provider (qwen3.7-max)", () => { + const llm = getModel("qwen-token-plan-cn", "qwen3.7-max"); + + it("should handle tool result with only image", { retry: 3, timeout: 30000 }, async () => { + await handleToolWithImageResult(llm); + }); + + it("should handle tool result with text and image", { retry: 3, timeout: 30000 }, async () => { + await handleToolWithTextAndImageResult(llm); + }); + }); + describe.skipIf(!process.env.KIMI_API_KEY)("Kimi For Coding Provider (kimi-for-coding)", () => { const llm = getModel("kimi-coding", "kimi-for-coding"); diff --git a/packages/ai/test/openai-completions-tool-choice.test.ts b/packages/ai/test/openai-completions-tool-choice.test.ts index b1f077737..a74849d8a 100644 --- a/packages/ai/test/openai-completions-tool-choice.test.ts +++ b/packages/ai/test/openai-completions-tool-choice.test.ts @@ -1088,6 +1088,18 @@ describe("openai-completions tool_choice", () => { } }); + it("stores Qwen Token Plan reasoning replay compat in built-in metadata", () => { + const providers = ["qwen-token-plan", "qwen-token-plan-cn"] as const; + + for (const provider of providers) { + const model = getModel(provider, "qwen3.7-max")!; + expect(model.compat?.thinkingFormat).toBe("qwen"); + expect(model.compat?.requiresReasoningContentOnAssistantMessages).toBeUndefined(); + expect(model.compat?.supportsDeveloperRole).toBe(false); + expect(model.compat?.supportsStore).toBe(false); + } + }); + it("replays Xiaomi MiMo assistant tool calls with empty reasoning_content when thinking is missing", async () => { const model = getModel("xiaomi", "mimo-v2.5-pro")!; const assistantMessage: AssistantMessage = { diff --git a/packages/ai/test/qwen-token-plan-models.test.ts b/packages/ai/test/qwen-token-plan-models.test.ts new file mode 100644 index 000000000..ac5bbc6a9 --- /dev/null +++ b/packages/ai/test/qwen-token-plan-models.test.ts @@ -0,0 +1,38 @@ +import { describe, expect, it } from "vitest"; +import { getModels } from "../src/compat.ts"; + +const TEXT_MODELS = [ + "MiniMax-M2.5", + "deepseek-v3.2", + "deepseek-v4-flash", + "deepseek-v4-pro", + "glm-5", + "glm-5.1", + "glm-5.2", + "kimi-k2.5", + "kimi-k2.6", + "kimi-k2.7-code", + "qwen3.6-flash", + "qwen3.6-plus", + "qwen3.7-max", + "qwen3.7-plus", + "qwen3.8-max-preview", +]; + +const IMAGE_MODELS = ["qwen-image-2.0", "qwen-image-2.0-pro", "wan2.7-image", "wan2.7-image-pro"]; + +describe("Qwen Token Plan models", () => { + it.each(["qwen-token-plan", "qwen-token-plan-cn"] as const)("exposes all text models on %s", (provider) => { + const modelIds = getModels(provider).map((model) => model.id); + for (const expected of TEXT_MODELS) { + expect(modelIds, `${provider} should include ${expected}`).toContain(expected); + } + }); + + it.each(["qwen-token-plan", "qwen-token-plan-cn"] as const)("omits image models from %s", (provider) => { + const modelIds = getModels(provider).map((model) => model.id); + for (const excluded of IMAGE_MODELS) { + expect(modelIds, `${provider} should not include ${excluded}`).not.toContain(excluded); + } + }); +}); diff --git a/packages/ai/test/stream.test.ts b/packages/ai/test/stream.test.ts index 484b74518..232df7a08 100644 --- a/packages/ai/test/stream.test.ts +++ b/packages/ai/test/stream.test.ts @@ -1168,6 +1168,65 @@ describe("Generate E2E Tests", () => { }, ); + describe.skipIf(!process.env.QWEN_TOKEN_PLAN_API_KEY)( + "Qwen Token Plan Provider (Qwen3.7-Max, international)", + () => { + const llm = getModel("qwen-token-plan", "qwen3.7-max"); + const thinkingOptions = { + thinkingEnabled: true, + reasoningEffort: "high", + } satisfies StreamOptionsWithExtras; + + it("should complete basic text generation", { retry: 3 }, async () => { + await basicTextGeneration(llm); + }); + + it("should handle tool calling", { retry: 3 }, async () => { + await handleToolCall(llm); + }); + + it("should handle streaming", { retry: 3 }, async () => { + await handleStreaming(llm); + }); + + it("should handle thinking mode", { retry: 3 }, async () => { + await handleThinking(llm, thinkingOptions); + }); + + it("should handle multi-turn with thinking and tools", { retry: 3 }, async () => { + await multiTurn(llm, thinkingOptions); + }); + }, + ); + + describe.skipIf(!process.env.QWEN_TOKEN_PLAN_CN_API_KEY)("Qwen Token Plan Provider (Qwen3.7-Max, CN region)", () => { + const llm = getModel("qwen-token-plan-cn", "qwen3.7-max"); + const thinkingOptions = { + thinkingEnabled: true, + reasoningEffort: "high", + } satisfies StreamOptionsWithExtras; + + it("should complete basic text generation", { retry: 3 }, async () => { + await basicTextGeneration(llm); + }); + + it("should handle tool calling", { retry: 3 }, async () => { + await handleToolCall(llm); + }); + + it("should handle streaming", { retry: 3 }, async () => { + await handleStreaming(llm); + }); + + it("should handle thinking mode", { retry: 3 }, async () => { + await handleThinking(llm, thinkingOptions); + }); + + it("should handle multi-turn with thinking and tools", { retry: 3 }, async () => { + await multiTurn(llm, thinkingOptions); + }); + }); + describe.skipIf(!process.env.ANT_LING_API_KEY)("Ant Ling Provider (Ling 2.6 Flash via OpenAI Completions)", () => { const llm = getModel("ant-ling", "Ling-2.6-flash"); diff --git a/packages/ai/test/tokens.test.ts b/packages/ai/test/tokens.test.ts index 42e6932d4..665e90099 100644 --- a/packages/ai/test/tokens.test.ts +++ b/packages/ai/test/tokens.test.ts @@ -276,6 +276,22 @@ describe("Token Statistics on Abort", () => { }); }); + describe.skipIf(!process.env.QWEN_TOKEN_PLAN_API_KEY)("Qwen Token Plan Provider", () => { + const llm = getModel("qwen-token-plan", "qwen3.7-max"); + + it("should include token stats when aborted mid-stream", { retry: 3, timeout: 30000 }, async () => { + await testTokensOnAbort(llm); + }); + }); + + describe.skipIf(!process.env.QWEN_TOKEN_PLAN_CN_API_KEY)("Qwen Token Plan (CN) Provider", () => { + const llm = getModel("qwen-token-plan-cn", "qwen3.7-max"); + + it("should include token stats when aborted mid-stream", { retry: 3, timeout: 30000 }, async () => { + await testTokensOnAbort(llm); + }); + }); + // ========================================================================= // OAuth-based providers (credentials from ~/.pi/agent/oauth.json) // ========================================================================= diff --git a/packages/ai/test/tool-call-without-result.test.ts b/packages/ai/test/tool-call-without-result.test.ts index 439b97580..f089e7a89 100644 --- a/packages/ai/test/tool-call-without-result.test.ts +++ b/packages/ai/test/tool-call-without-result.test.ts @@ -254,6 +254,22 @@ describe("Tool Call Without Result Tests", () => { }); }); + describe.skipIf(!process.env.QWEN_TOKEN_PLAN_API_KEY)("Qwen Token Plan Provider", () => { + const model = getModel("qwen-token-plan", "qwen3.7-max"); + + it("should filter out tool calls without corresponding tool results", { retry: 3, timeout: 30000 }, async () => { + await testToolCallWithoutResult(model); + }); + }); + + describe.skipIf(!process.env.QWEN_TOKEN_PLAN_CN_API_KEY)("Qwen Token Plan (CN) Provider", () => { + const model = getModel("qwen-token-plan-cn", "qwen3.7-max"); + + it("should filter out tool calls without corresponding tool results", { retry: 3, timeout: 30000 }, async () => { + await testToolCallWithoutResult(model); + }); + }); + describe.skipIf(!process.env.KIMI_API_KEY)("Kimi For Coding Provider", () => { const model = getModel("kimi-coding", "kimi-for-coding"); diff --git a/packages/ai/test/total-tokens.test.ts b/packages/ai/test/total-tokens.test.ts index ed5cae292..00ba4ab2f 100644 --- a/packages/ai/test/total-tokens.test.ts +++ b/packages/ai/test/total-tokens.test.ts @@ -562,6 +562,56 @@ describe("totalTokens field", () => { ); }); + // ========================================================================= + // Qwen Token Plan + // ========================================================================= + + describe.skipIf(!process.env.QWEN_TOKEN_PLAN_API_KEY)("Qwen Token Plan", () => { + it( + "qwen3.7-max - should return totalTokens equal to sum of components", + { retry: 3, timeout: 60000 }, + async () => { + const llm = getModel("qwen-token-plan", "qwen3.7-max"); + + console.log(`\nQwen Token Plan / ${llm.id}:`); + const { first, second } = await testTotalTokensWithCache(llm, { + apiKey: process.env.QWEN_TOKEN_PLAN_API_KEY, + }); + + logUsage("First request", first); + logUsage("Second request", second); + + assertTotalTokensEqualsComponents(first); + assertTotalTokensEqualsComponents(second); + }, + ); + }); + + // ========================================================================= + // Qwen Token Plan CN + // ========================================================================= + + describe.skipIf(!process.env.QWEN_TOKEN_PLAN_CN_API_KEY)("Qwen Token Plan (CN)", () => { + it( + "qwen3.7-max - should return totalTokens equal to sum of components", + { retry: 3, timeout: 60000 }, + async () => { + const llm = getModel("qwen-token-plan-cn", "qwen3.7-max"); + + console.log(`\nQwen Token Plan CN / ${llm.id}:`); + const { first, second } = await testTotalTokensWithCache(llm, { + apiKey: process.env.QWEN_TOKEN_PLAN_CN_API_KEY, + }); + + logUsage("First request", first); + logUsage("Second request", second); + + assertTotalTokensEqualsComponents(first); + assertTotalTokensEqualsComponents(second); + }, + ); + }); + // ========================================================================= // Kimi For Coding // ========================================================================= diff --git a/packages/ai/test/unicode-surrogate.test.ts b/packages/ai/test/unicode-surrogate.test.ts index 8e059d1d0..89f9602e0 100644 --- a/packages/ai/test/unicode-surrogate.test.ts +++ b/packages/ai/test/unicode-surrogate.test.ts @@ -695,6 +695,38 @@ describe("AI Providers Unicode Surrogate Pair Tests", () => { }, ); + describe.skipIf(!process.env.QWEN_TOKEN_PLAN_API_KEY)("Qwen Token Plan Provider Unicode Handling", () => { + const llm = getModel("qwen-token-plan", "qwen3.7-max"); + + it("should handle emoji in tool results", { retry: 3, timeout: 30000 }, async () => { + await testEmojiInToolResults(llm); + }); + + it("should handle real-world LinkedIn comment data with emoji", { retry: 3, timeout: 30000 }, async () => { + await testRealWorldLinkedInData(llm); + }); + + it("should handle unpaired high surrogate (0xD83D) in tool results", { retry: 3, timeout: 30000 }, async () => { + await testUnpairedHighSurrogate(llm); + }); + }); + + describe.skipIf(!process.env.QWEN_TOKEN_PLAN_CN_API_KEY)("Qwen Token Plan (CN) Provider Unicode Handling", () => { + const llm = getModel("qwen-token-plan-cn", "qwen3.7-max"); + + it("should handle emoji in tool results", { retry: 3, timeout: 30000 }, async () => { + await testEmojiInToolResults(llm); + }); + + it("should handle real-world LinkedIn comment data with emoji", { retry: 3, timeout: 30000 }, async () => { + await testRealWorldLinkedInData(llm); + }); + + it("should handle unpaired high surrogate (0xD83D) in tool results", { retry: 3, timeout: 30000 }, async () => { + await testUnpairedHighSurrogate(llm); + }); + }); + describe.skipIf(!process.env.KIMI_API_KEY)("Kimi For Coding Provider Unicode Handling", () => { const llm = getModel("kimi-coding", "kimi-for-coding"); diff --git a/packages/coding-agent/docs/providers.md b/packages/coding-agent/docs/providers.md index 9517dd0cb..e8758a3d8 100644 --- a/packages/coding-agent/docs/providers.md +++ b/packages/coding-agent/docs/providers.md @@ -87,6 +87,8 @@ pi | Kimi For Coding | `KIMI_API_KEY` | `kimi-coding` | | MiniMax | `MINIMAX_API_KEY` | `minimax` | | MiniMax (China) | `MINIMAX_CN_API_KEY` | `minimax-cn` | +| Qwen Token Plan | `QWEN_TOKEN_PLAN_API_KEY` | `qwen-token-plan` | +| Qwen Token Plan (China) | `QWEN_TOKEN_PLAN_CN_API_KEY` | `qwen-token-plan-cn` | | Xiaomi MiMo | `XIAOMI_API_KEY` | `xiaomi` | | Xiaomi MiMo Token Plan (China) | `XIAOMI_TOKEN_PLAN_CN_API_KEY` | `xiaomi-token-plan-cn` | | Xiaomi MiMo Token Plan (Amsterdam) | `XIAOMI_TOKEN_PLAN_AMS_API_KEY` | `xiaomi-token-plan-ams` | @@ -109,6 +111,8 @@ Store credentials in `~/.pi/agent/auth.json`: "opencode": { "type": "api_key", "key": "..." }, "opencode-go": { "type": "api_key", "key": "..." }, "together": { "type": "api_key", "key": "..." }, + "qwen-token-plan": { "type": "api_key", "key": "sk-sp-..." }, + "qwen-token-plan-cn": { "type": "api_key", "key": "sk-sp-..." }, "xiaomi": { "type": "api_key", "key": "..." }, "xiaomi-token-plan-cn": { "type": "api_key", "key": "..." }, "xiaomi-token-plan-ams": { "type": "api_key", "key": "..." }, diff --git a/packages/coding-agent/src/cli/args.ts b/packages/coding-agent/src/cli/args.ts index 0ca339694..c194579c5 100644 --- a/packages/coding-agent/src/cli/args.ts +++ b/packages/coding-agent/src/cli/args.ts @@ -362,6 +362,8 @@ ${chalk.bold("Environment Variables:")} CLOUDFLARE_API_KEY - Cloudflare API token (Workers AI and AI Gateway) CLOUDFLARE_ACCOUNT_ID - Cloudflare account id (required for both) CLOUDFLARE_GATEWAY_ID - Cloudflare AI Gateway slug (required for AI Gateway) + QWEN_TOKEN_PLAN_API_KEY - Qwen Token Plan API key (international region) + QWEN_TOKEN_PLAN_CN_API_KEY - Qwen Token Plan API key (China region) XIAOMI_API_KEY - Xiaomi MiMo API key (api.xiaomimimo.com billing) XIAOMI_TOKEN_PLAN_CN_API_KEY - Xiaomi MiMo Token Plan API key (China region) XIAOMI_TOKEN_PLAN_AMS_API_KEY - Xiaomi MiMo Token Plan API key (Amsterdam region) diff --git a/packages/coding-agent/src/core/model-resolver.ts b/packages/coding-agent/src/core/model-resolver.ts index 1f4dceaf3..bcf668e56 100644 --- a/packages/coding-agent/src/core/model-resolver.ts +++ b/packages/coding-agent/src/core/model-resolver.ts @@ -44,6 +44,8 @@ export const defaultModelPerProvider: Record = { "kimi-coding": "kimi-for-coding", "cloudflare-workers-ai": "@cf/moonshotai/kimi-k2.6", "cloudflare-ai-gateway": "workers-ai/@cf/moonshotai/kimi-k2.6", + "qwen-token-plan": "qwen3.7-max", + "qwen-token-plan-cn": "qwen3.7-max", xiaomi: "mimo-v2.5-pro", "xiaomi-token-plan-cn": "mimo-v2.5-pro", "xiaomi-token-plan-ams": "mimo-v2.5-pro", diff --git a/test.sh b/test.sh index 9b563b795..e6315a0e4 100755 --- a/test.sh +++ b/test.sh @@ -55,6 +55,8 @@ unset XIAOMI_API_KEY unset XIAOMI_TOKEN_PLAN_CN_API_KEY unset XIAOMI_TOKEN_PLAN_AMS_API_KEY unset XIAOMI_TOKEN_PLAN_SGP_API_KEY +unset QWEN_TOKEN_PLAN_API_KEY +unset QWEN_TOKEN_PLAN_CN_API_KEY unset RADIUS_API_KEY unset PI_GATEWAY unset PI_EXPERIMENTAL From 2e810c2402a8f5a9f96043b47aedfc6ec8d2d8e7 Mon Sep 17 00:00:00 2001 From: Cristina Poncela Cubeiro <140309543+cristinaponcela@users.noreply.github.com> Date: Mon, 20 Jul 2026 13:59:18 +0200 Subject: [PATCH 18/62] fix: issue template links --- .github/ISSUE_TEMPLATE/bug.yml | 4 ++-- .github/ISSUE_TEMPLATE/contribution.yml | 4 ++-- .github/ISSUE_TEMPLATE/package-report.yml | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/.github/ISSUE_TEMPLATE/bug.yml b/.github/ISSUE_TEMPLATE/bug.yml index 43cb73c57..b486675bd 100644 --- a/.github/ISSUE_TEMPLATE/bug.yml +++ b/.github/ISSUE_TEMPLATE/bug.yml @@ -5,9 +5,9 @@ body: - type: markdown attributes: value: | - **Before you start:** Read [CONTRIBUTING.md](https://github.com/earendil-works/pi-mono/blob/main/CONTRIBUTING.md). + **Before you start:** Read [CONTRIBUTING.md](https://github.com/earendil-works/pi/blob/main/CONTRIBUTING.md). - New issues from new contributors are auto-closed by default. Maintainers review auto-closed issues daily. Issues that do not meet the quality bar in [CONTRIBUTING.md](https://github.com/earendil-works/pi-mono/blob/main/CONTRIBUTING.md) will not be reopened or receive a reply. + New issues from new contributors are auto-closed by default. Maintainers review auto-closed issues daily. Issues that do not meet the quality bar in [CONTRIBUTING.md](https://github.com/earendil-works/pi/blob/main/CONTRIBUTING.md) will not be reopened or receive a reply. Keep this short. If it doesn't fit on one screen, it's too long. Write in your own voice. diff --git a/.github/ISSUE_TEMPLATE/contribution.yml b/.github/ISSUE_TEMPLATE/contribution.yml index 0cf6c5be4..1a36a1a35 100644 --- a/.github/ISSUE_TEMPLATE/contribution.yml +++ b/.github/ISSUE_TEMPLATE/contribution.yml @@ -5,9 +5,9 @@ body: - type: markdown attributes: value: | - **Before you start:** Read [CONTRIBUTING.md](https://github.com/earendil-works/pi-mono/blob/main/CONTRIBUTING.md). + **Before you start:** Read [CONTRIBUTING.md](https://github.com/earendil-works/pi/blob/main/CONTRIBUTING.md). - New issues from new contributors are auto-closed by default. Maintainers review auto-closed issues daily. Issues that do not meet the quality bar in [CONTRIBUTING.md](https://github.com/earendil-works/pi-mono/blob/main/CONTRIBUTING.md) will not be reopened or receive a reply. + New issues from new contributors are auto-closed by default. Maintainers review auto-closed issues daily. Issues that do not meet the quality bar in [CONTRIBUTING.md](https://github.com/earendil-works/pi/blob/main/CONTRIBUTING.md) will not be reopened or receive a reply. Keep this short. If it doesn't fit on one screen, it's too long. Write in your own voice. diff --git a/.github/ISSUE_TEMPLATE/package-report.yml b/.github/ISSUE_TEMPLATE/package-report.yml index 846e25ee5..dfd5c80d4 100644 --- a/.github/ISSUE_TEMPLATE/package-report.yml +++ b/.github/ISSUE_TEMPLATE/package-report.yml @@ -7,7 +7,7 @@ body: value: | Use this form to report a package listed on pi.dev. For Pi core bugs, use the bug report template instead. - New issues from new contributors are auto-closed by default. Maintainers review auto-closed issues daily. Issues that do not meet the quality bar in [CONTRIBUTING.md](https://github.com/earendil-works/pi-mono/blob/main/CONTRIBUTING.md) will not be reopened or receive a reply. + New issues from new contributors are auto-closed by default. Maintainers review auto-closed issues daily. Issues that do not meet the quality bar in [CONTRIBUTING.md](https://github.com/earendil-works/pi/blob/main/CONTRIBUTING.md) will not be reopened or receive a reply. Keep this short. If it doesn't fit on one screen, it's too long. Write in your own voice. From 3595e080cb013232c4141cab9976eebcdddcc417 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Mon, 20 Jul 2026 14:03:15 +0200 Subject: [PATCH 19/62] fix(tui): keep paste registry in sync when deleting paste markers Undo snapshots now restore paste content and counters alongside editor text. Paste marker renumbering shifts registry entries in ascending ID order before rewriting markers, preventing literal or incorrect paste content on submit.\n\nCloses #6844 --- packages/tui/CHANGELOG.md | 4 ++ packages/tui/src/components/editor.ts | 37 ++++++++----- packages/tui/test/editor.test.ts | 75 +++++++++++++++++++++++++++ 3 files changed, 104 insertions(+), 12 deletions(-) diff --git a/packages/tui/CHANGELOG.md b/packages/tui/CHANGELOG.md index 2255c1032..7c18f034d 100644 --- a/packages/tui/CHANGELOG.md +++ b/packages/tui/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Fixed editor paste registry corruption when deleting paste markers: undo now restores the paste registry together with the text, and marker renumbering shifts registry entries in ascending id order, so submitted prompts no longer contain literal `[paste #N ...]` markers or the wrong paste's content ([#6844](https://github.com/earendil-works/pi/issues/6844)). + ## [0.80.10] - 2026-07-16 ## [0.80.9] - 2026-07-16 diff --git a/packages/tui/src/components/editor.ts b/packages/tui/src/components/editor.ts index 3b55dc552..706324c77 100644 --- a/packages/tui/src/components/editor.ts +++ b/packages/tui/src/components/editor.ts @@ -212,6 +212,13 @@ interface EditorState { cursorCol: number; } +/** Undo snapshot: editor text state plus the paste registry. */ +interface EditorSnapshot { + state: EditorState; + pastes: Map; + pasteCounter: number; +} + interface LayoutLine { text: string; hasCursor: boolean; @@ -318,7 +325,7 @@ export class Editor implements Component, Focusable { private snappedFromCursorCol: number | null = null; // Undo support - private undoStack = new UndoStack(); + private undoStack = new UndoStack(); public onSubmit?: (text: string) => void; public onChange?: (text: string) => void; @@ -999,13 +1006,13 @@ export class Editor implements Component, Focusable { this.cancelAutocomplete(); this.lastAction = null; this.exitHistoryBrowsing(); - this.pastes.clear(); - this.pasteCounter = 0; const normalized = this.normalizeText(text); // Push undo snapshot if content differs (makes programmatic changes undoable) if (this.getText() !== normalized) { this.pushUndoSnapshot(); } + this.pastes.clear(); + this.pasteCounter = 0; this.setTextInternal(normalized); } @@ -1284,17 +1291,21 @@ export class Editor implements Component, Focusable { this.pastes.delete(targetId); this.pasteCounter--; - // We got to update id of markers which are greater than the removed one + // Shift registry entries down in ascending id order, independent + // of marker order in the text ([paste #3] becomes [paste #2] when + // [paste #1] is removed). + const higherIds = [...this.pastes.keys()].filter((id) => id > targetId).sort((a, b) => a - b); + for (const id of higherIds) { + this.pastes.set(id - 1, this.pastes.get(id)!); + this.pastes.delete(id); + } + + // Renumber markers with ids greater than the removed one. this.state.lines = this.state.lines.map((line) => line.replace(PASTE_MARKER_REGEX, (fullMatch, idGroup, suffixGroup) => { const x = Number(idGroup); if (x <= targetId) return fullMatch; - - // [paste #3] become [paste #2] if we remove [paste #1] - const newText = `[paste #${x - 1}${suffixGroup}]`; - this.pastes.set(x - 1, this.pastes.get(x) ?? newText); - this.pastes.delete(x); - return newText; + return `[paste #${x - 1}${suffixGroup}]`; }), ); } @@ -1994,14 +2005,16 @@ export class Editor implements Component, Focusable { } private pushUndoSnapshot(): void { - this.undoStack.push(this.state); + this.undoStack.push({ state: this.state, pastes: this.pastes, pasteCounter: this.pasteCounter }); } private undo(): void { this.exitHistoryBrowsing(); const snapshot = this.undoStack.pop(); if (!snapshot) return; - Object.assign(this.state, snapshot); + Object.assign(this.state, snapshot.state); + this.pastes = snapshot.pastes; + this.pasteCounter = snapshot.pasteCounter; this.lastAction = null; this.preferredVisualCol = null; if (this.onChange) { diff --git a/packages/tui/test/editor.test.ts b/packages/tui/test/editor.test.ts index 0f33370e1..d2c278d76 100644 --- a/packages/tui/test/editor.test.ts +++ b/packages/tui/test/editor.test.ts @@ -3553,6 +3553,11 @@ describe("Editor component", () => { return editor.getText(); } + /** Helper: 12-line paste content with a distinguishing tag */ + function bigPaste(tag: string): string { + return Array.from({ length: 12 }, (_, i) => `${tag}${i}`).join("\n"); + } + it("creates a paste marker for large pastes", () => { const editor = new Editor(createTestTUI(), defaultEditorTheme); const text = pasteWithMarker(editor); @@ -3690,6 +3695,76 @@ describe("Editor component", () => { assert.strictEqual(editor.getText(), textBefore); }); + it("undo after paste marker deletion restores the paste registry", () => { + const editor = new Editor(createTestTUI(), defaultEditorTheme); + let submitted = ""; + editor.onSubmit = (t) => { + submitted = t; + }; + + const paste = bigPaste("alpha"); + editor.handleInput(`\x1b[200~${paste}\x1b[201~`); + editor.handleInput("\x7f"); // delete the marker + editor.handleInput("\x1b[45;5u"); // undo: restores marker text and registry + editor.handleInput("\r"); + assert.strictEqual(submitted, paste); + }); + + it("undo after deleting the first of two paste markers restores both registry entries", () => { + const editor = new Editor(createTestTUI(), defaultEditorTheme); + let submitted = ""; + editor.onSubmit = (t) => { + submitted = t; + }; + + const pasteA = bigPaste("alpha"); + const pasteB = bigPaste("beta"); + editor.handleInput(`\x1b[200~${pasteA}\x1b[201~`); // #1 = A + editor.handleInput(`\x1b[200~${pasteB}\x1b[201~`); // #2 = B, cursor at end + editor.handleInput("\x01"); // Ctrl+A + editor.handleInput("\x1b[C"); // right over marker #1 + editor.handleInput("\x7f"); // delete marker #1, renumbers #2 -> #1 + editor.handleInput("\x1b[45;5u"); // undo + editor.handleInput("\r"); + assert.strictEqual(submitted, pasteA + pasteB); + }); + + it("renumbers the paste registry in ascending id order when markers are out of order in text", () => { + const editor = new Editor(createTestTUI(), defaultEditorTheme); + let submitted = ""; + editor.onSubmit = (t) => { + submitted = t; + }; + + const pasteA = bigPaste("alpha"); + const pasteB = bigPaste("beta"); + const pasteC = bigPaste("gamma"); + editor.handleInput(`\x1b[200~${pasteA}\x1b[201~`); // #1 = A + editor.handleInput("\x01"); // Ctrl+A + editor.handleInput(`\x1b[200~${pasteB}\x1b[201~`); // #2 = B, text: [#2][#1] + editor.handleInput("\x01"); // Ctrl+A + editor.handleInput(`\x1b[200~${pasteC}\x1b[201~`); // #3 = C, text: [#3][#2][#1] + editor.handleInput("\x05"); // Ctrl+E + editor.handleInput("\x7f"); // delete marker #1, renumber #3 -> #2 and #2 -> #1 + editor.handleInput("\r"); + assert.strictEqual(submitted, pasteC + pasteB); + }); + + it("undo after setText restores paste markers and registry", () => { + const editor = new Editor(createTestTUI(), defaultEditorTheme); + let submitted = ""; + editor.onSubmit = (t) => { + submitted = t; + }; + + const paste = bigPaste("alpha"); + editor.handleInput(`\x1b[200~${paste}\x1b[201~`); + editor.setText("replacement"); + editor.handleInput("\x1b[45;5u"); // undo + editor.handleInput("\r"); + assert.strictEqual(submitted, paste); + }); + it("handles multiple paste markers in same line", () => { const editor = new Editor(createTestTUI(), defaultEditorTheme); pasteWithMarker(editor); From 13437ca828894f43f973c630d208b488637d8fa9 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Mon, 20 Jul 2026 14:03:33 +0200 Subject: [PATCH 20/62] fix(ai): normalize Kimi K2.7 to the canonical coding model Treat the models.dev k2p7 entry as an alias for kimi-for-coding and regenerate provider catalogs. This restores consistency between generated model types and values and unblocks repository type checks. --- packages/ai/scripts/generate-models.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/ai/scripts/generate-models.ts b/packages/ai/scripts/generate-models.ts index 40e7275c7..82700ac99 100644 --- a/packages/ai/scripts/generate-models.ts +++ b/packages/ai/scripts/generate-models.ts @@ -1673,12 +1673,12 @@ async function loadModelsDevData(): Promise[]> { const kimiModels = data["kimi-for-coding"].models as Record; const hasCanonicalModel = Object.prototype.hasOwnProperty.call(kimiModels, "kimi-for-coding"); - const kimiAliases = new Set(["k2p5", "k2p6"]); + const kimiAliases = new Set(["k2p5", "k2p6", "k2p7"]); for (const [modelId, model] of Object.entries(kimiModels)) { const m = model as ModelsDevModel; if (m.tool_call !== true) continue; - // models.dev may expose versioned aliases (e.g. k2p5/k2p6). + // models.dev may expose versioned aliases (e.g. k2p5/k2p6/k2p7). // Normalize aliases to the canonical model id and drop duplicates when canonical exists. if (kimiAliases.has(modelId) && hasCanonicalModel) continue; From 1942b2600f2bbb7b5bc86c30379e14c0766c058c Mon Sep 17 00:00:00 2001 From: Cristina Poncela Cubeiro <140309543+cristinaponcela@users.noreply.github.com> Date: Mon, 20 Jul 2026 16:20:49 +0200 Subject: [PATCH 21/62] fix: env section ignored (#6864) --- packages/ai/src/auth/helpers.ts | 4 +++- packages/ai/src/providers/amazon-bedrock.ts | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/packages/ai/src/auth/helpers.ts b/packages/ai/src/auth/helpers.ts index 5c64535f4..251c60da1 100644 --- a/packages/ai/src/auth/helpers.ts +++ b/packages/ai/src/auth/helpers.ts @@ -14,7 +14,9 @@ export function envApiKeyAuth(name: string, envVars: readonly string[]): ApiKeyA return { type: "api_key", key }; }, resolve: async ({ ctx, credential }) => { - if (credential?.key) return { auth: { apiKey: credential.key }, source: "stored credential" }; + if (credential?.key) { + return { auth: { apiKey: credential.key }, env: credential.env, source: "stored credential" }; + } for (const envVar of envVars) { const value = await ctx.env(envVar); if (value) return { auth: { apiKey: value }, source: envVar }; diff --git a/packages/ai/src/providers/amazon-bedrock.ts b/packages/ai/src/providers/amazon-bedrock.ts index 83a7050fd..61deef5cb 100644 --- a/packages/ai/src/providers/amazon-bedrock.ts +++ b/packages/ai/src/providers/amazon-bedrock.ts @@ -50,7 +50,9 @@ const bedrockAuth: ApiKeyAuth = { return { type: "api_key" }; }, resolve: async ({ ctx, credential }) => { - if (credential?.key) return { auth: { apiKey: credential.key }, source: "stored credential" }; + if (credential?.key) { + return { auth: { apiKey: credential.key }, env: credential.env, source: "stored credential" }; + } if (await ctx.env("AWS_BEARER_TOKEN_BEDROCK")) return { auth: {}, source: "AWS_BEARER_TOKEN_BEDROCK" }; if (credential?.env?.AWS_PROFILE ?? (await ctx.env("AWS_PROFILE"))) { return { From c179395218416b0352d2d11a4025fac0668a2676 Mon Sep 17 00:00:00 2001 From: Cristina Poncela Cubeiro <140309543+cristinaponcela@users.noreply.github.com> Date: Mon, 20 Jul 2026 16:29:10 +0200 Subject: [PATCH 22/62] feat: get_available_thinking_levels rpc (#6865) --- packages/coding-agent/docs/rpc.md | 20 +++++++++++++++++ .../coding-agent/src/modes/rpc/rpc-client.ts | 8 +++++++ .../coding-agent/src/modes/rpc/rpc-mode.ts | 5 +++++ .../coding-agent/src/modes/rpc/rpc-types.ts | 8 +++++++ packages/coding-agent/test/rpc.test.ts | 22 +++++++++++++++++++ 5 files changed, 63 insertions(+) diff --git a/packages/coding-agent/docs/rpc.md b/packages/coding-agent/docs/rpc.md index 1493f2d55..8074ceb95 100644 --- a/packages/coding-agent/docs/rpc.md +++ b/packages/coding-agent/docs/rpc.md @@ -313,6 +313,26 @@ Response: } ``` +#### get_available_thinking_levels + +List the thinking levels supported by the current model. Returns `["off"]` for a model without reasoning support. + +```json +{"type": "get_available_thinking_levels"} +``` + +Response: +```json +{ + "type": "response", + "command": "get_available_thinking_levels", + "success": true, + "data": { + "levels": ["off", "minimal", "low", "medium", "high"] + } +} +``` + ### Queue Modes #### set_steering_mode diff --git a/packages/coding-agent/src/modes/rpc/rpc-client.ts b/packages/coding-agent/src/modes/rpc/rpc-client.ts index 934c6f1fc..79cd7f2f9 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-client.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-client.ts @@ -280,6 +280,14 @@ export class RpcClient { return this.getData(response); } + /** + * Get list of available thinking levels for the current model. + */ + async getAvailableThinkingLevels(): Promise { + const response = await this.send({ type: "get_available_thinking_levels" }); + return this.getData<{ levels: ThinkingLevel[] }>(response).levels; + } + /** * Set steering mode. */ diff --git a/packages/coding-agent/src/modes/rpc/rpc-mode.ts b/packages/coding-agent/src/modes/rpc/rpc-mode.ts index 58b645313..06eadccb9 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-mode.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-mode.ts @@ -504,6 +504,11 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise { + await client.start(); + + const levels = await client.getAvailableThinkingLevels(); + expect(levels.length).toBeGreaterThan(0); + + // The current level reported by get_state must be in the available list + const state = await client.getState(); + expect(levels).toContain(state.thinkingLevel); + + // cycle_thinking_level must only ever land on levels from get_available_thinking_levels + const initialLevel = state.thinkingLevel; + const cycled = await client.cycleThinkingLevel(); + if (cycled) { + expect(levels).toContain(cycled.level); + // distinct cycle step (unless only one level) + if (levels.length > 1) { + expect(cycled.level).not.toBe(initialLevel); + } + } + }, 30000); + test("should get available models", async () => { await client.start(); From 2fd3868401f0fe79496f3df731b20348ba57538a Mon Sep 17 00:00:00 2001 From: David Brailovsky Date: Mon, 20 Jul 2026 16:41:43 +0200 Subject: [PATCH 23/62] add usage info to branch summary, compaction and tool result entries (#6671) * add usage info to branch summary entries * add usage to compaction entries * allow custom tools to report llm usage in tool results * allow observing and patching usage in tool_result hooks * agent-harness: save usage in entries for compaction, branch summaries and tool results --- packages/agent/src/agent-loop.ts | 2 + packages/agent/src/harness/agent-harness.ts | 24 +++- .../compaction/branch-summarization.ts | 1 + .../src/harness/compaction/compaction.ts | 77 +++++++--- packages/agent/src/harness/session/session.ts | 7 +- packages/agent/src/harness/types.ts | 24 +++- packages/agent/src/types.ts | 9 +- packages/agent/test/agent-loop.test.ts | 26 ++++ .../agent/test/harness/agent-harness.test.ts | 132 +++++++++++++++++- .../agent/test/harness/compaction.test.ts | 44 +++++- packages/agent/test/harness/session.test.ts | 43 ++++++ packages/agent/test/utils/calculate.ts | 8 ++ packages/ai/src/types.ts | 2 + .../coding-agent/src/core/agent-session.ts | 53 ++++--- .../core/compaction/branch-summarization.ts | 4 +- .../src/core/compaction/compaction.ts | 80 ++++++++--- .../src/core/extensions/runner.ts | 5 + .../coding-agent/src/core/extensions/types.ts | 5 + .../coding-agent/src/core/session-manager.ts | 17 ++- .../coding-agent/src/core/usage-totals.ts | 27 ++++ .../modes/interactive/components/footer.ts | 32 ++--- .../test/agent-session-stats.test.ts | 83 ++++++++++- .../test/branch-summary-extensions.test.ts | 55 ++++++++ .../test/compaction-summary-reasoning.test.ts | 14 +- .../coding-agent/test/footer-width.test.ts | 88 ++++++++++-- .../session-manager/tree-traversal.test.ts | 22 ++- .../suite/agent-session-compaction.test.ts | 37 +++++ .../agent-session-model-extension.test.ts | 42 ++++-- .../6324-branch-summary-ambient-auth.test.ts | 3 +- 29 files changed, 844 insertions(+), 122 deletions(-) create mode 100644 packages/coding-agent/src/core/usage-totals.ts create mode 100644 packages/coding-agent/test/branch-summary-extensions.test.ts diff --git a/packages/agent/src/agent-loop.ts b/packages/agent/src/agent-loop.ts index 14eb9f6e1..c9e731cd1 100644 --- a/packages/agent/src/agent-loop.ts +++ b/packages/agent/src/agent-loop.ts @@ -737,6 +737,7 @@ async function finalizeExecutedToolCall( ...result, content: afterResult.content ?? result.content, details: afterResult.details ?? result.details, + usage: afterResult.usage ?? result.usage, terminate: afterResult.terminate ?? result.terminate, }; isError = afterResult.isError ?? isError; @@ -780,6 +781,7 @@ function createToolResultMessage(finalized: FinalizedToolCallOutcome): ToolResul // so the null never enters session history or provider payloads. content: finalized.result.content ?? [], details: finalized.result.details, + usage: finalized.result.usage, ...(finalized.result.addedToolNames?.length ? { addedToolNames: finalized.result.addedToolNames } : {}), isError: finalized.isError, timestamp: Date.now(), diff --git a/packages/agent/src/harness/agent-harness.ts b/packages/agent/src/harness/agent-harness.ts index 1ce57b14d..85af6b28c 100644 --- a/packages/agent/src/harness/agent-harness.ts +++ b/packages/agent/src/harness/agent-harness.ts @@ -32,6 +32,7 @@ import type { AgentHarnessResources, AgentHarnessStreamOptions, AgentHarnessStreamOptionsPatch, + CompactResult, ExecutionEnv, NavigateTreeResult, PendingSessionWrite, @@ -434,9 +435,16 @@ export class AgentHarness< content: result.content, details: result.details, isError, + usage: result.usage, }); return patch - ? { content: patch.content, details: patch.details, isError: patch.isError, terminate: patch.terminate } + ? { + content: patch.content, + details: patch.details, + isError: patch.isError, + usage: patch.usage, + terminate: patch.terminate, + } : undefined; }, prepareNextTurn: async () => { @@ -690,9 +698,7 @@ export class AgentHarness< } } - async compact( - customInstructions?: string, - ): Promise<{ summary: string; firstKeptEntryId: string; tokensBefore: number; details?: unknown }> { + async compact(customInstructions?: string): Promise { if (this.phase !== "idle") throw new AgentHarnessError("busy", "compact() requires idle harness"); this.phase = "compaction"; try { @@ -723,6 +729,7 @@ export class AgentHarness< result.tokensBefore, result.details, provided !== undefined, + result.usage, ); const entry = await this.session.getEntry(entryId); if (entry?.type === "compaction") { @@ -764,6 +771,7 @@ export class AgentHarness< let summaryEntry: NavigateTreeResult["summaryEntry"]; let summaryText: string | undefined = hookResult?.summary?.summary; let summaryDetails: unknown = hookResult?.summary?.details; + let summaryUsage = hookResult?.summary?.usage; if (!summaryText && options?.summarize && entries.length > 0) { const model = this.model; if (!model) throw new AgentHarnessError("invalid_state", "No model set for branch summary"); @@ -779,6 +787,7 @@ export class AgentHarness< throw new AgentHarnessError("branch_summary", branchSummary.error.message, branchSummary.error); } summaryText = branchSummary.value.summary; + summaryUsage = branchSummary.value.usage; summaryDetails = { readFiles: branchSummary.value.readFiles, modifiedFiles: branchSummary.value.modifiedFiles, @@ -798,7 +807,12 @@ export class AgentHarness< const summaryId = await this.session.moveTo( newLeafId, summaryText - ? { summary: summaryText, details: summaryDetails, fromHook: hookResult?.summary !== undefined } + ? { + summary: summaryText, + details: summaryDetails, + usage: summaryUsage, + fromHook: hookResult?.summary !== undefined, + } : undefined, ); if (summaryId) { diff --git a/packages/agent/src/harness/compaction/branch-summarization.ts b/packages/agent/src/harness/compaction/branch-summarization.ts index a8d686739..ab56f3ac7 100644 --- a/packages/agent/src/harness/compaction/branch-summarization.ts +++ b/packages/agent/src/harness/compaction/branch-summarization.ts @@ -252,6 +252,7 @@ export async function generateBranchSummary( return ok({ summary: summary || "No summary generated", + usage: response.usage, readFiles, modifiedFiles, }); diff --git a/packages/agent/src/harness/compaction/compaction.ts b/packages/agent/src/harness/compaction/compaction.ts index a7f8ebc06..58938c13c 100644 --- a/packages/agent/src/harness/compaction/compaction.ts +++ b/packages/agent/src/harness/compaction/compaction.ts @@ -101,10 +101,35 @@ export interface CompactionResult { firstKeptEntryId: string; /** Estimated context tokens before compaction. */ tokensBefore: number; + /** Usage from the LLM call(s) that generated this summary, if available. */ + usage?: Usage; /** Optional implementation-specific details stored with the compaction entry. */ details?: T; } +function combineUsage(first: Usage, second: Usage): Usage { + return { + input: first.input + second.input, + output: first.output + second.output, + cacheRead: first.cacheRead + second.cacheRead, + cacheWrite: first.cacheWrite + second.cacheWrite, + ...(first.cacheWrite1h !== undefined || second.cacheWrite1h !== undefined + ? { cacheWrite1h: (first.cacheWrite1h ?? 0) + (second.cacheWrite1h ?? 0) } + : {}), + ...(first.reasoning !== undefined || second.reasoning !== undefined + ? { reasoning: (first.reasoning ?? 0) + (second.reasoning ?? 0) } + : {}), + totalTokens: first.totalTokens + second.totalTokens, + cost: { + input: first.cost.input + second.cost.input, + output: first.cost.output + second.cost.output, + cacheRead: first.cost.cacheRead + second.cost.cacheRead, + cacheWrite: first.cost.cacheWrite + second.cost.cacheWrite, + total: first.cost.total + second.cost.total, + }, + }; +} + /** Compaction thresholds and retention settings. */ export interface CompactionSettings { /** Enable automatic compaction decisions. */ @@ -474,7 +499,7 @@ export async function generateSummary( customInstructions?: string, previousSummary?: string, thinkingLevel?: ThinkingLevel, -): Promise> { +): Promise> { const maxTokens = Math.min( Math.floor(0.8 * reserveTokens), model.maxTokens > 0 ? model.maxTokens : Number.POSITIVE_INFINITY, @@ -523,7 +548,7 @@ export async function generateSummary( const textContent = contentText(response.content); - return ok(textContent); + return ok({ text: textContent, usage: response.usage }); } /** Prepared inputs for a compaction run. */ @@ -656,22 +681,26 @@ export async function compact( } let summary: string; + let summaryUsage: Usage; if (isSplitTurn && turnPrefixMessages.length > 0) { - const historyResult = - messagesToSummarize.length > 0 - ? await generateSummary( - messagesToSummarize, - models, - model, - settings.reserveTokens, - signal, - customInstructions, - previousSummary, - thinkingLevel, - ) - : ok("No prior history."); - if (!historyResult.ok) return err(historyResult.error); + let historyText = "No prior history."; + let historyUsage: Usage | undefined; + if (messagesToSummarize.length > 0) { + const historyResult = await generateSummary( + messagesToSummarize, + models, + model, + settings.reserveTokens, + signal, + customInstructions, + previousSummary, + thinkingLevel, + ); + if (!historyResult.ok) return err(historyResult.error); + historyText = historyResult.value.text; + historyUsage = historyResult.value.usage; + } const turnPrefixResult = await generateTurnPrefixSummary( turnPrefixMessages, models, @@ -681,7 +710,10 @@ export async function compact( thinkingLevel, ); if (!turnPrefixResult.ok) return err(turnPrefixResult.error); - summary = `${historyResult.value}\n\n---\n\n**Turn Context (split turn):**\n\n${turnPrefixResult.value}`; + summary = `${historyText}\n\n---\n\n**Turn Context (split turn):**\n\n${turnPrefixResult.value.text}`; + summaryUsage = historyUsage + ? combineUsage(historyUsage, turnPrefixResult.value.usage) + : turnPrefixResult.value.usage; } else { const summaryResult = await generateSummary( messagesToSummarize, @@ -694,7 +726,8 @@ export async function compact( thinkingLevel, ); if (!summaryResult.ok) return err(summaryResult.error); - summary = summaryResult.value; + summary = summaryResult.value.text; + summaryUsage = summaryResult.value.usage; } const { readFiles, modifiedFiles } = computeFileLists(fileOps); @@ -704,6 +737,7 @@ export async function compact( summary, firstKeptEntryId, tokensBefore, + usage: summaryUsage, details: { readFiles, modifiedFiles } as CompactionDetails, }); } @@ -714,7 +748,7 @@ async function generateTurnPrefixSummary( reserveTokens: number, signal?: AbortSignal, thinkingLevel?: ThinkingLevel, -): Promise> { +): Promise> { const maxTokens = Math.min( Math.floor(0.5 * reserveTokens), model.maxTokens > 0 ? model.maxTokens : Number.POSITIVE_INFINITY, @@ -749,5 +783,8 @@ async function generateTurnPrefixSummary( ); } - return ok(contentText(response.content)); + return ok({ + text: contentText(response.content), + usage: response.usage, + }); } diff --git a/packages/agent/src/harness/session/session.ts b/packages/agent/src/harness/session/session.ts index f5224c922..bff41fdd7 100644 --- a/packages/agent/src/harness/session/session.ts +++ b/packages/agent/src/harness/session/session.ts @@ -1,4 +1,4 @@ -import type { ImageContent, TextContent } from "@earendil-works/pi-ai"; +import type { ImageContent, TextContent, Usage } from "@earendil-works/pi-ai"; import type { AgentMessage } from "../../types.ts"; import { createBranchSummaryMessage, createCompactionSummaryMessage, createCustomMessage } from "../messages.ts"; import type { @@ -247,6 +247,7 @@ export class Session { tokensBefore: number, details?: T, fromHook?: boolean, + usage?: Usage, ): Promise { return this.appendTypedEntry({ type: "compaction", @@ -257,6 +258,7 @@ export class Session { firstKeptEntryId, tokensBefore, details, + usage, fromHook, } satisfies CompactionEntry); } @@ -317,7 +319,7 @@ export class Session { async moveTo( entryId: string | null, - summary?: { summary: string; details?: unknown; fromHook?: boolean }, + summary?: { summary: string; details?: unknown; usage?: Usage; fromHook?: boolean }, ): Promise { if (entryId !== null && !(await this.storage.getEntry(entryId))) { throw new SessionError("not_found", `Entry ${entryId} not found`); @@ -332,6 +334,7 @@ export class Session { fromId: entryId ?? "root", summary: summary.summary, details: summary.details, + usage: summary.usage, fromHook: summary.fromHook, } satisfies BranchSummaryEntry); } diff --git a/packages/agent/src/harness/types.ts b/packages/agent/src/harness/types.ts index 3f9655504..cdc7a5426 100644 --- a/packages/agent/src/harness/types.ts +++ b/packages/agent/src/harness/types.ts @@ -1,4 +1,12 @@ -import type { ImageContent, Model, Models, SimpleStreamOptions, TextContent, Transport } from "@earendil-works/pi-ai"; +import type { + ImageContent, + Model, + Models, + SimpleStreamOptions, + TextContent, + Transport, + Usage, +} from "@earendil-works/pi-ai"; import type { AgentEvent, AgentMessage, AgentTool, QueueMode, ThinkingLevel } from "../index.ts"; import type { Session } from "./session/session.ts"; @@ -365,6 +373,7 @@ export interface CompactionEntry extends SessionTreeEntryBase { firstKeptEntryId: string; tokensBefore: number; details?: T; + usage?: Usage; fromHook?: boolean; } @@ -373,6 +382,7 @@ export interface BranchSummaryEntry extends SessionTreeEntryBase { fromId: string; summary: string; details?: T; + usage?: Usage; fromHook?: boolean; } @@ -572,6 +582,7 @@ export interface ToolResultEvent { content: Array; details: unknown; isError: boolean; + usage?: Usage; } export interface SessionBeforeCompactEvent { @@ -687,6 +698,7 @@ export interface ToolResultPatch { content?: Array; details?: unknown; isError?: boolean; + usage?: Usage; terminate?: boolean; } @@ -697,7 +709,12 @@ export interface SessionBeforeCompactResult { export interface SessionBeforeTreeResult { cancel?: boolean; - summary?: { summary: string; details?: unknown }; + summary?: { + summary: string; + details?: unknown; + /** Usage from the LLM call that generated this summary, if available. */ + usage?: Usage; + }; customInstructions?: string; replaceInstructions?: boolean; label?: string; @@ -738,6 +755,8 @@ export interface CompactResult { summary: string; firstKeptEntryId: string; tokensBefore: number; + /** Usage from the LLM call(s) that generated this summary, if available. */ + usage?: Usage; details?: unknown; } @@ -793,6 +812,7 @@ export interface GenerateBranchSummaryOptions { export interface BranchSummaryResult { summary: string; + usage?: Usage; readFiles: string[]; modifiedFiles: string[]; } diff --git a/packages/agent/src/types.ts b/packages/agent/src/types.ts index 2de50144b..52f34891a 100644 --- a/packages/agent/src/types.ts +++ b/packages/agent/src/types.ts @@ -11,6 +11,7 @@ import type { TextContent, Tool, ToolResultMessage, + Usage, } from "@earendil-works/pi-ai"; import type { Static, TSchema } from "typebox"; @@ -69,15 +70,18 @@ export interface BeforeToolCallResult { * - `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 + * - `usage`: if provided, replaces the tool result usage * - `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`. + * There is no deep merge for `content`, `details`, or `usage`. */ export interface AfterToolCallResult { content?: (TextContent | ImageContent)[]; details?: unknown; isError?: boolean; + /** Usage from the final tool execution itself, if available. Not used for main LLM context accounting. */ + usage?: Usage; /** * 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. @@ -273,6 +277,7 @@ export interface AgentLoopConfig extends SimpleStreamOptions { * - `content` replaces the full content array * - `details` replaces the full details payload * - `isError` replaces the error flag + * - `usage` replaces the tool result usage * - `terminate` replaces the early-termination hint * * Any omitted fields keep their original values. No deep merge is performed. @@ -352,6 +357,8 @@ export interface AgentToolResult { content: (TextContent | ImageContent)[]; /** Arbitrary structured details for logs or UI rendering. */ details: T; + /** Usage from the final tool execution itself, if available. Not used for main LLM context accounting. */ + usage?: Usage; /** Names of tools introduced by this result and available from this transcript point onward. */ addedToolNames?: string[]; /** diff --git a/packages/agent/test/agent-loop.test.ts b/packages/agent/test/agent-loop.test.ts index 038680478..f9641ce74 100644 --- a/packages/agent/test/agent-loop.test.ts +++ b/packages/agent/test/agent-loop.test.ts @@ -239,6 +239,23 @@ describe("agentLoop with AgentMessage", () => { it("should handle tool calls and results", async () => { const toolSchema = Type.Object({ value: Type.String() }); const executed: string[] = []; + const toolUsage = { + input: 1, + output: 2, + cacheRead: 3, + cacheWrite: 4, + totalTokens: 10, + cost: { input: 0.1, output: 0.2, cacheRead: 0.3, cacheWrite: 0.4, total: 1 }, + }; + const patchedToolUsage = { + input: 5, + output: 6, + cacheRead: 7, + cacheWrite: 8, + totalTokens: 26, + cost: { input: 0.5, output: 0.6, cacheRead: 0.7, cacheWrite: 0.8, total: 2.6 }, + }; + let observedToolUsage: typeof toolUsage | undefined; const tool: AgentTool = { name: "echo", label: "Echo", @@ -249,6 +266,7 @@ describe("agentLoop with AgentMessage", () => { return { content: [{ type: "text", text: `echoed: ${params.value}` }], details: { value: params.value }, + usage: toolUsage, }; }, }; @@ -264,6 +282,10 @@ describe("agentLoop with AgentMessage", () => { const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, + afterToolCall: async ({ result }) => { + observedToolUsage = result.usage; + return { usage: patchedToolUsage }; + }, }; let callIndex = 0; @@ -305,6 +327,10 @@ describe("agentLoop with AgentMessage", () => { if (toolEnd?.type === "tool_execution_end") { expect(toolEnd.isError).toBe(false); } + expect(observedToolUsage).toEqual(toolUsage); + const messages = await stream.result(); + const toolResult = messages.find((message) => message.role === "toolResult"); + expect(toolResult?.role === "toolResult" ? toolResult.usage : undefined).toEqual(patchedToolUsage); }); it("should not execute tool calls from a length-truncated assistant message", async () => { diff --git a/packages/agent/test/harness/agent-harness.test.ts b/packages/agent/test/harness/agent-harness.test.ts index d13eca84a..48f18e249 100644 --- a/packages/agent/test/harness/agent-harness.test.ts +++ b/packages/agent/test/harness/agent-harness.test.ts @@ -5,6 +5,7 @@ import { fauxProvider, fauxToolCall, type RegisterFauxProviderOptions, + type Usage, } from "@earendil-works/pi-ai"; import { getModel } from "@earendil-works/pi-ai/compat"; import { describe, expect, it } from "vitest"; @@ -14,7 +15,7 @@ import { InMemorySessionStorage } from "../../src/harness/session/memory-storage import { Session } from "../../src/harness/session/session.ts"; import type { PromptTemplate, Skill } from "../../src/harness/types.ts"; import type { AgentMessage, AgentTool } from "../../src/types.ts"; -import { calculateTool } from "../utils/calculate.ts"; +import { calculateTool, createCalculateToolWithUsage } from "../utils/calculate.ts"; import { getCurrentTimeTool } from "../utils/get-current-time.ts"; interface AppSkill extends Skill { @@ -60,6 +61,34 @@ function getReasoning(options: unknown): unknown { return options.reasoning; } +function createUsage(input: number, output: number, cacheRead = 0, cacheWrite = 0): Usage { + return { + input, + output, + cacheRead, + cacheWrite, + totalTokens: input + output + cacheRead + cacheWrite, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }; +} + +function createUserMessage(text: string): AgentMessage { + return { role: "user", content: [{ type: "text", text }], timestamp: Date.now() }; +} + +function createAssistantMessage(text: string): AgentMessage { + return { + role: "assistant", + content: [{ type: "text", text }], + api: "faux", + provider: "faux", + model: "faux-1", + usage: createUsage(100, 50), + stopReason: "stop", + timestamp: Date.now(), + }; +} + describe("AgentHarness", () => { it("constructs directly and exposes queue modes", () => { const session = new Session(new InMemorySessionStorage()); @@ -427,14 +456,18 @@ describe("AgentHarness", () => { }), ]); const session = new Session(new InMemorySessionStorage()); + const toolUsage = createUsage(1, 2, 3, 4); + const patchedToolUsage = createUsage(5, 6, 7, 8); + const calculateToolWithUsage = createCalculateToolWithUsage(toolUsage); const harness = new AgentHarness({ models, env: new NodeExecutionEnv({ cwd: process.cwd() }), session, model: registration.getModel(), - tools: [calculateTool], + tools: [calculateToolWithUsage], }); const seenToolCalls: Array<{ id: string; name: string; expression: unknown }> = []; + let seenToolUsage: Usage | undefined; harness.on("tool_call", (event) => { seenToolCalls.push({ id: event.toolCallId, name: event.toolName, expression: event.input.expression }); return undefined; @@ -442,9 +475,11 @@ describe("AgentHarness", () => { harness.on("tool_result", (event) => { expect(event.toolCallId).toBe("call-1"); expect(event.toolName).toBe("calculate"); + seenToolUsage = event.usage; return { content: [{ type: "text", text: "patched result" }], details: { patched: true }, + usage: patchedToolUsage, terminate: true, }; }); @@ -455,14 +490,107 @@ describe("AgentHarness", () => { (entry) => entry.type === "message" && entry.message.role === "toolResult", ); expect(seenToolCalls).toEqual([{ id: "call-1", name: "calculate", expression: "2 + 2" }]); + expect(seenToolUsage).toEqual(toolUsage); expect(toolResult).toMatchObject({ type: "message", message: { role: "toolResult", content: [{ type: "text", text: "patched result" }], details: { patched: true }, + usage: patchedToolUsage, + }, + }); + }); + + it("persists generated compaction usage", async () => { + const registration = newFaux(); + registration.setResponses([fauxAssistantMessage("## Goal\nTest summary")]); + const session = new Session(new InMemorySessionStorage()); + await session.appendMessage(createUserMessage("one")); + await session.appendMessage(createAssistantMessage("two")); + const harness = new AgentHarness({ + models, + env: new NodeExecutionEnv({ cwd: process.cwd() }), + session, + model: registration.getModel(), + }); + + const result = await harness.compact(); + const compaction = (await session.getEntries()).find((entry) => entry.type === "compaction"); + + expect(result.usage?.totalTokens).toBeGreaterThan(0); + expect(compaction?.type === "compaction" ? compaction.usage : undefined).toEqual(result.usage); + }); + + it("persists hook-provided compaction usage", async () => { + const registration = newFaux(); + const usage = createUsage(5, 6, 7, 8); + const session = new Session(new InMemorySessionStorage()); + await session.appendMessage(createUserMessage("one")); + await session.appendMessage(createAssistantMessage("two")); + const harness = new AgentHarness({ + models, + env: new NodeExecutionEnv({ cwd: process.cwd() }), + session, + model: registration.getModel(), + }); + harness.on("session_before_compact", (event) => ({ + compaction: { + summary: "hook summary", + firstKeptEntryId: event.preparation.firstKeptEntryId, + tokensBefore: event.preparation.tokensBefore, + usage, }, + })); + + const result = await harness.compact(); + const compaction = (await session.getEntries()).find((entry) => entry.type === "compaction"); + + expect(result.usage).toEqual(usage); + expect(compaction?.type === "compaction" ? compaction.usage : undefined).toEqual(usage); + }); + + it("persists generated branch summary usage", async () => { + const registration = newFaux(); + registration.setResponses([fauxAssistantMessage("## Goal\nBranch summary")]); + const session = new Session(new InMemorySessionStorage()); + const targetId = await session.appendMessage(createUserMessage("first branch")); + await session.appendMessage(createAssistantMessage("first reply")); + await session.appendMessage(createUserMessage("abandoned work")); + await session.appendMessage(createAssistantMessage("abandoned reply")); + const harness = new AgentHarness({ + models, + env: new NodeExecutionEnv({ cwd: process.cwd() }), + session, + model: registration.getModel(), }); + + const result = await harness.navigateTree(targetId, { summarize: true }); + + expect(result.summaryEntry?.usage?.totalTokens).toBeGreaterThan(0); + }); + + it("persists hook-provided branch summary usage", async () => { + const registration = newFaux(); + const usage = createUsage(13, 14, 15, 16); + const session = new Session(new InMemorySessionStorage()); + const targetId = await session.appendMessage(createUserMessage("first branch")); + await session.appendMessage(createAssistantMessage("first reply")); + await session.appendMessage(createUserMessage("abandoned work")); + await session.appendMessage(createAssistantMessage("abandoned reply")); + const harness = new AgentHarness({ + models, + env: new NodeExecutionEnv({ cwd: process.cwd() }), + session, + model: registration.getModel(), + }); + harness.on("session_before_tree", () => ({ + summary: { summary: "hook branch summary", usage }, + })); + + const result = await harness.navigateTree(targetId, { summarize: true }); + + expect(result.summaryEntry?.usage).toEqual(usage); }); it("preserves app tool types for getters and update events", async () => { diff --git a/packages/agent/test/harness/compaction.test.ts b/packages/agent/test/harness/compaction.test.ts index 8a228eeb8..7cdea0c93 100644 --- a/packages/agent/test/harness/compaction.test.ts +++ b/packages/agent/test/harness/compaction.test.ts @@ -6,6 +6,7 @@ import { fauxProvider, type Message, type Model, + type Models, type Usage, } from "@earendil-works/pi-ai"; import { beforeEach, describe, expect, it } from "vitest"; @@ -142,6 +143,17 @@ function createFauxModel(reasoning: boolean, maxTokens = 8192): { faux: FauxProv return { faux, model: faux.getModel() }; } +function createModelsWithSimpleResponses(responses: AssistantMessage[]): Models { + const remaining = [...responses]; + const stub = Object.create(models) as Models; + stub.completeSimple = async () => { + const response = remaining.shift(); + if (!response) throw new Error("No faux completeSimple response queued"); + return response; + }; + return stub; +} + describe("harness compaction", () => { beforeEach(() => { nextId = 0; @@ -501,7 +513,12 @@ describe("harness compaction", () => { await generateSummary(messages, models, model, 2000, undefined, "focus", "old summary"), ); - expect(summary).toContain("Test summary"); + expect(summary.text).toContain("Test summary"); + expect(summary.usage.input).toBeGreaterThan(0); + expect(summary.usage.output).toBeGreaterThan(0); + expect(summary.usage.totalTokens).toBe( + summary.usage.input + summary.usage.output + summary.usage.cacheRead + summary.usage.cacheWrite, + ); expect(promptText).toContain("\nold summary\n"); expect(promptText).toContain("Additional focus: focus"); }); @@ -578,6 +595,30 @@ describe("harness compaction", () => { expect(invalidResult).toMatchObject({ ok: false, error: { code: "invalid_session" } }); }); + it("combines usage for split-turn compaction summaries", async () => { + const messages: AgentMessage[] = [createUserMessage("Summarize this.")]; + const { model } = createFauxModel(false); + const historyUsage = createMockUsage(1, 2, 3, 4); + const turnPrefixUsage = createMockUsage(5, 6, 7, 8); + const usageModels = createModelsWithSimpleResponses([ + { ...fauxAssistantMessage("history summary"), usage: historyUsage }, + { ...fauxAssistantMessage("turn prefix summary"), usage: turnPrefixUsage }, + ]); + const preparation: CompactionPreparation = { + firstKeptEntryId: "entry-keep", + messagesToSummarize: messages, + turnPrefixMessages: messages, + isSplitTurn: true, + tokensBefore: 100, + fileOps: { read: new Set(), written: new Set(), edited: new Set() }, + settings: { enabled: true, reserveTokens: 2000, keepRecentTokens: 20 }, + }; + + const result = getOrThrow(await compact(preparation, usageModels, model)); + + expect(result.usage).toEqual(createMockUsage(6, 8, 10, 12)); + }); + it("passes reasoning through turn-prefix summaries when enabled", async () => { const messages: AgentMessage[] = [createUserMessage("Summarize this.")]; const seenOptions: Array | undefined> = []; @@ -646,6 +687,7 @@ describe("harness compaction", () => { const result = getOrThrow(await compact(preparation!, models, model)); expect(result.summary.length).toBeGreaterThan(0); expect(result.firstKeptEntryId).toBeTruthy(); + expect(result.usage?.totalTokens).toBeGreaterThan(0); expect(result.details).toBeDefined(); }); }); diff --git a/packages/agent/test/harness/session.test.ts b/packages/agent/test/harness/session.test.ts index c009123d6..26af52e2a 100644 --- a/packages/agent/test/harness/session.test.ts +++ b/packages/agent/test/harness/session.test.ts @@ -86,6 +86,49 @@ async function runSessionSuite( expect(context.messages[1]?.role).toBe("branchSummary"); }); + it("persists compaction usage", async () => { + const session = new Session(await createStorage()); + const firstKeptEntryId = await session.appendMessage(createUserMessage("one")); + const usage = { + input: 1, + output: 2, + cacheRead: 3, + cacheWrite: 4, + totalTokens: 10, + cost: { input: 0.1, output: 0.2, cacheRead: 0.3, cacheWrite: 0.4, total: 1 }, + }; + + const compactionId = await session.appendCompaction( + "summary", + firstKeptEntryId, + 1234, + undefined, + false, + usage, + ); + + const compactionEntry = await session.getEntry(compactionId); + expect(compactionEntry?.type === "compaction" ? compactionEntry.usage : undefined).toEqual(usage); + }); + + it("persists branch summary usage", async () => { + const session = new Session(await createStorage()); + const user1 = await session.appendMessage(createUserMessage("one")); + const usage = { + input: 1, + output: 2, + cacheRead: 3, + cacheWrite: 4, + totalTokens: 10, + cost: { input: 0.1, output: 0.2, cacheRead: 0.3, cacheWrite: 0.4, total: 1 }, + }; + + const summaryId = await session.moveTo(user1, { summary: "summary text", usage }); + + const summaryEntry = await session.getEntry(summaryId!); + expect(summaryEntry?.type === "branch_summary" ? summaryEntry.usage : undefined).toEqual(usage); + }); + it("supports custom message entries in context", async () => { const session = new Session(await createStorage()); await session.appendMessage(createUserMessage("one")); diff --git a/packages/agent/test/utils/calculate.ts b/packages/agent/test/utils/calculate.ts index 6158efbde..0a4fefc58 100644 --- a/packages/agent/test/utils/calculate.ts +++ b/packages/agent/test/utils/calculate.ts @@ -1,3 +1,4 @@ +import type { Usage } from "@earendil-works/pi-ai"; import { type Static, Type } from "typebox"; import type { AgentTool, AgentToolResult } from "../../src/types.ts"; @@ -30,3 +31,10 @@ export const calculateTool: AgentTool = { return calculate(args.expression); }, }; + +export function createCalculateToolWithUsage(usage: Usage): AgentTool { + return { + ...calculateTool, + execute: async (_toolCallId: string, args: CalculateParams) => ({ ...calculate(args.expression), usage }), + }; +} diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 9fa1ab20f..aeaf18c67 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -408,6 +408,8 @@ export interface ToolResultMessage { toolName: string; content: (TextContent | ImageContent)[]; // Supports text and images details?: TDetails; + /** Usage from the tool execution itself, if available. Not part of main LLM context accounting. */ + usage?: Usage; /** * Names from `Context.tools` that became available after this result. * Providers with native deferred tool loading use this as the load point; diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index c4afb9b17..eb4f78e90 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -32,6 +32,7 @@ import type { Model, ProviderHeaders, TextContent, + Usage, } from "@earendil-works/pi-ai/compat"; import { clampThinkingLevel, @@ -104,6 +105,7 @@ import { type BuildSystemPromptOptions, buildSystemPrompt } from "./system-promp import { type BashOperations, createLocalBashOperations } from "./tools/bash.ts"; import { createAllToolDefinitions } from "./tools/index.ts"; import { createToolDefinitionFromAgentTool } from "./tools/tool-definition-wrapper.ts"; +import { addUsageToTotals, createUsageTotals } from "./usage-totals.ts"; // ============================================================================ // Skill Block Parsing @@ -482,6 +484,7 @@ export class AgentSession { content: result.content, details: result.details, isError, + usage: result.usage, }); if (!hookResult) { @@ -492,6 +495,7 @@ export class AgentSession { content: hookResult.content, details: hookResult.details, isError: hookResult.isError ?? isError, + usage: hookResult.usage, }; }; } @@ -1812,6 +1816,7 @@ export class AgentSession { let summary: string; let firstKeptEntryId: string; let tokensBefore: number; + let usage: Usage | undefined; let details: unknown; if (extensionCompaction) { @@ -1819,6 +1824,7 @@ export class AgentSession { summary = extensionCompaction.summary; firstKeptEntryId = extensionCompaction.firstKeptEntryId; tokensBefore = extensionCompaction.tokensBefore; + usage = extensionCompaction.usage; details = extensionCompaction.details; } else { // Generate compaction result @@ -1836,6 +1842,7 @@ export class AgentSession { summary = result.summary; firstKeptEntryId = result.firstKeptEntryId; tokensBefore = result.tokensBefore; + usage = result.usage; details = result.details; } @@ -1843,7 +1850,7 @@ export class AgentSession { throw new Error("Compaction cancelled"); } - this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore, details, fromExtension); + this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore, details, fromExtension, usage); const newEntries = this.sessionManager.getEntries(); const sessionContext = this.sessionManager.buildSessionContext(); this.agent.state.messages = sessionContext.messages; @@ -1869,6 +1876,7 @@ export class AgentSession { firstKeptEntryId, tokensBefore, estimatedTokensAfter, + usage, details, }; this._emit({ @@ -2084,6 +2092,7 @@ export class AgentSession { let summary: string; let firstKeptEntryId: string; let tokensBefore: number; + let usage: Usage | undefined; let details: unknown; if (extensionCompaction) { @@ -2091,6 +2100,7 @@ export class AgentSession { summary = extensionCompaction.summary; firstKeptEntryId = extensionCompaction.firstKeptEntryId; tokensBefore = extensionCompaction.tokensBefore; + usage = extensionCompaction.usage; details = extensionCompaction.details; } else { // Generate compaction result @@ -2108,6 +2118,7 @@ export class AgentSession { summary = compactResult.summary; firstKeptEntryId = compactResult.firstKeptEntryId; tokensBefore = compactResult.tokensBefore; + usage = compactResult.usage; details = compactResult.details; } @@ -2122,7 +2133,7 @@ export class AgentSession { return false; } - this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore, details, fromExtension); + this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore, details, fromExtension, usage); const newEntries = this.sessionManager.getEntries(); const sessionContext = this.sessionManager.buildSessionContext(); this.agent.state.messages = sessionContext.messages; @@ -2148,6 +2159,7 @@ export class AgentSession { firstKeptEntryId, tokensBefore, estimatedTokensAfter, + usage, details, }; this._emit({ type: "compaction_end", reason, result, aborted: false, willRetry }); @@ -2872,7 +2884,7 @@ export class AgentSession { this._branchSummaryAbortController = new AbortController(); try { - let extensionSummary: { summary: string; details?: unknown } | undefined; + let extensionSummary: { summary: string; details?: unknown; usage?: Usage } | undefined; let fromExtension = false; // Emit session_before_tree event @@ -2907,6 +2919,7 @@ export class AgentSession { // Run default summarizer if needed let summaryText: string | undefined; let summaryDetails: unknown; + let summaryUsage: Usage | undefined; if (options.summarize && entriesToSummarize.length > 0 && !extensionSummary) { const model = this.model!; const { apiKey, headers, env } = await this._getSummarizationRequestAuth(model); @@ -2929,6 +2942,7 @@ export class AgentSession { throw new Error(result.error); } summaryText = result.summary; + summaryUsage = result.usage; summaryDetails = { readFiles: result.readFiles || [], modifiedFiles: result.modifiedFiles || [], @@ -2936,6 +2950,7 @@ export class AgentSession { } else if (extensionSummary) { summaryText = extensionSummary.summary; summaryDetails = extensionSummary.details; + summaryUsage = extensionSummary.usage; } // Determine the new leaf position based on target type @@ -2965,6 +2980,7 @@ export class AgentSession { summaryText, summaryDetails, fromExtension, + summaryUsage, ); summaryEntry = this.sessionManager.getEntry(summaryId) as BranchSummaryEntry; @@ -3037,13 +3053,12 @@ export class AgentSession { let toolResults = 0; let totalMessages = 0; let toolCalls = 0; - let totalInput = 0; - let totalOutput = 0; - let totalCacheRead = 0; - let totalCacheWrite = 0; - let totalCost = 0; + const usageTotals = createUsageTotals(); for (const entry of this.sessionManager.getEntries()) { + if ((entry.type === "branch_summary" || entry.type === "compaction") && entry.usage) { + addUsageToTotals(usageTotals, entry.usage); + } if (entry.type !== "message") continue; totalMessages++; const message = entry.message; @@ -3051,18 +3066,16 @@ export class AgentSession { userMessages++; } else if (message.role === "toolResult") { toolResults++; + if (message.usage) { + addUsageToTotals(usageTotals, message.usage); + } } else if (message.role === "assistant") { assistantMessages++; const assistantMsg = message as AssistantMessage; if (Array.isArray(assistantMsg.content)) { toolCalls += assistantMsg.content.filter((c) => c.type === "toolCall").length; } - const usage = assistantMsg.usage; - totalInput += usage.input; - totalOutput += usage.output; - totalCacheRead += usage.cacheRead; - totalCacheWrite += usage.cacheWrite; - totalCost += usage.cost.total; + addUsageToTotals(usageTotals, assistantMsg.usage); } } @@ -3075,13 +3088,13 @@ export class AgentSession { toolResults, totalMessages, tokens: { - input: totalInput, - output: totalOutput, - cacheRead: totalCacheRead, - cacheWrite: totalCacheWrite, - total: totalInput + totalOutput + totalCacheRead + totalCacheWrite, + input: usageTotals.input, + output: usageTotals.output, + cacheRead: usageTotals.cacheRead, + cacheWrite: usageTotals.cacheWrite, + total: usageTotals.input + usageTotals.output + usageTotals.cacheRead + usageTotals.cacheWrite, }, - cost: totalCost, + cost: usageTotals.cost, contextUsage: this.getContextUsage(), }; } diff --git a/packages/coding-agent/src/core/compaction/branch-summarization.ts b/packages/coding-agent/src/core/compaction/branch-summarization.ts index f2e06876c..3366f06a7 100644 --- a/packages/coding-agent/src/core/compaction/branch-summarization.ts +++ b/packages/coding-agent/src/core/compaction/branch-summarization.ts @@ -7,7 +7,7 @@ import type { AgentMessage, StreamFn } from "@earendil-works/pi-agent-core"; import { contentText } from "@earendil-works/pi-ai"; -import type { Model, SimpleStreamOptions } from "@earendil-works/pi-ai/compat"; +import type { Model, SimpleStreamOptions, Usage } from "@earendil-works/pi-ai/compat"; import { completeSimple } from "@earendil-works/pi-ai/compat"; import { convertToLlm, @@ -33,6 +33,7 @@ import { export interface BranchSummaryResult { summary?: string; + usage?: Usage; readFiles?: string[]; modifiedFiles?: string[]; aborted?: boolean; @@ -363,6 +364,7 @@ export async function generateBranchSummary( return { summary: summary || "No summary generated", + usage: response.usage, readFiles, modifiedFiles, }; diff --git a/packages/coding-agent/src/core/compaction/compaction.ts b/packages/coding-agent/src/core/compaction/compaction.ts index 5c71c74dc..dafcb09ea 100644 --- a/packages/coding-agent/src/core/compaction/compaction.ts +++ b/packages/coding-agent/src/core/compaction/compaction.ts @@ -90,10 +90,35 @@ export interface CompactionResult { firstKeptEntryId: string; tokensBefore: number; estimatedTokensAfter?: number; + /** Usage from the LLM call(s) that generated this summary, if available */ + usage?: Usage; /** Extension-specific data (e.g., ArtifactIndex, version markers for structured compaction) */ details?: T; } +function combineUsage(first: Usage, second: Usage): Usage { + return { + input: first.input + second.input, + output: first.output + second.output, + cacheRead: first.cacheRead + second.cacheRead, + cacheWrite: first.cacheWrite + second.cacheWrite, + ...(first.cacheWrite1h !== undefined || second.cacheWrite1h !== undefined + ? { cacheWrite1h: (first.cacheWrite1h ?? 0) + (second.cacheWrite1h ?? 0) } + : {}), + ...(first.reasoning !== undefined || second.reasoning !== undefined + ? { reasoning: (first.reasoning ?? 0) + (second.reasoning ?? 0) } + : {}), + totalTokens: first.totalTokens + second.totalTokens, + cost: { + input: first.cost.input + second.cost.input, + output: first.cost.output + second.cost.output, + cacheRead: first.cost.cacheRead + second.cost.cacheRead, + cacheWrite: first.cost.cacheWrite + second.cost.cacheWrite, + total: first.cost.total + second.cost.total, + }, + }; +} + // ============================================================================ // Types // ============================================================================ @@ -556,7 +581,7 @@ export async function generateSummary( thinkingLevel?: ThinkingLevel, streamFn?: StreamFn, env?: Record, -): Promise { +): Promise<{ text: string; usage: Usage }> { const maxTokens = Math.min( Math.floor(0.8 * reserveTokens), model.maxTokens > 0 ? model.maxTokens : Number.POSITIVE_INFINITY, @@ -603,7 +628,7 @@ export async function generateSummary( const textContent = contentText(response.content); - return textContent; + return { text: textContent, usage: response.usage }; } // ============================================================================ @@ -759,24 +784,28 @@ export async function compact( // Generate summaries and merge into one let summary: string; + let summaryUsage: Usage; if (isSplitTurn && turnPrefixMessages.length > 0) { - const historyResult = - messagesToSummarize.length > 0 - ? await generateSummary( - messagesToSummarize, - model, - settings.reserveTokens, - apiKey, - headers, - signal, - customInstructions, - previousSummary, - thinkingLevel, - streamFn, - env, - ) - : "No prior history."; + let historyText = "No prior history."; + let historyUsage: Usage | undefined; + if (messagesToSummarize.length > 0) { + const historyResult = await generateSummary( + messagesToSummarize, + model, + settings.reserveTokens, + apiKey, + headers, + signal, + customInstructions, + previousSummary, + thinkingLevel, + streamFn, + env, + ); + historyText = historyResult.text; + historyUsage = historyResult.usage; + } const turnPrefixResult = await generateTurnPrefixSummary( turnPrefixMessages, model, @@ -789,10 +818,11 @@ export async function compact( streamFn, ); // Merge into single summary - summary = `${historyResult}\n\n---\n\n**Turn Context (split turn):**\n\n${turnPrefixResult}`; + summary = `${historyText}\n\n---\n\n**Turn Context (split turn):**\n\n${turnPrefixResult.text}`; + summaryUsage = historyUsage ? combineUsage(historyUsage, turnPrefixResult.usage) : turnPrefixResult.usage; } else { // Just generate history summary - summary = await generateSummary( + const result = await generateSummary( messagesToSummarize, model, settings.reserveTokens, @@ -805,6 +835,8 @@ export async function compact( streamFn, env, ); + summary = result.text; + summaryUsage = result.usage; } // Compute file lists and append to summary @@ -819,6 +851,7 @@ export async function compact( summary, firstKeptEntryId, tokensBefore, + usage: summaryUsage, details: { readFiles, modifiedFiles } as CompactionDetails, }; } @@ -836,7 +869,7 @@ async function generateTurnPrefixSummary( signal?: AbortSignal, thinkingLevel?: ThinkingLevel, streamFn?: StreamFn, -): Promise { +): Promise<{ text: string; usage: Usage }> { const maxTokens = Math.min( Math.floor(0.5 * reserveTokens), model.maxTokens > 0 ? model.maxTokens : Number.POSITIVE_INFINITY, @@ -863,5 +896,8 @@ async function generateTurnPrefixSummary( throw new Error(`Turn prefix summarization failed: ${response.errorMessage || "Unknown error"}`); } - return contentText(response.content); + return { + text: contentText(response.content), + usage: response.usage, + }; } diff --git a/packages/coding-agent/src/core/extensions/runner.ts b/packages/coding-agent/src/core/extensions/runner.ts index 4bd43dda9..3c320e14f 100644 --- a/packages/coding-agent/src/core/extensions/runner.ts +++ b/packages/coding-agent/src/core/extensions/runner.ts @@ -883,6 +883,10 @@ export class ExtensionRunner { currentEvent.isError = handlerResult.isError; modified = true; } + if (handlerResult.usage !== undefined) { + currentEvent.usage = handlerResult.usage; + modified = true; + } } catch (err) { const message = err instanceof Error ? err.message : String(err); const stack = err instanceof Error ? err.stack : undefined; @@ -904,6 +908,7 @@ export class ExtensionRunner { content: currentEvent.content, details: currentEvent.details, isError: currentEvent.isError, + usage: currentEvent.usage, }; } diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index 514f7af66..e8bf2627c 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/packages/coding-agent/src/core/extensions/types.ts @@ -30,6 +30,7 @@ import type { SimpleStreamOptions, TextContent, ToolResultMessage, + Usage, } from "@earendil-works/pi-ai"; import type { AutocompleteItem, @@ -905,6 +906,8 @@ interface ToolResultEventBase { input: Record; content: (TextContent | ImageContent)[]; isError: boolean; + /** Usage from the tool execution itself, if available. */ + usage?: Usage; } export interface BashToolResultEvent extends ToolResultEventBase { @@ -1072,6 +1075,7 @@ export interface ToolResultEventResult { content?: (TextContent | ImageContent)[]; details?: unknown; isError?: boolean; + usage?: Usage; } export interface MessageEndEventResult { @@ -1104,6 +1108,7 @@ export interface SessionBeforeTreeResult { summary?: { summary: string; details?: unknown; + usage?: Usage; }; /** Override custom instructions for summarization */ customInstructions?: string; diff --git a/packages/coding-agent/src/core/session-manager.ts b/packages/coding-agent/src/core/session-manager.ts index 0ab665957..0dbce18ff 100644 --- a/packages/coding-agent/src/core/session-manager.ts +++ b/packages/coding-agent/src/core/session-manager.ts @@ -1,5 +1,5 @@ import type { AgentMessage } from "@earendil-works/pi-agent-core"; -import { type ImageContent, type Message, type TextContent, uuidv7 } from "@earendil-works/pi-ai"; +import { type ImageContent, type Message, type TextContent, type Usage, uuidv7 } from "@earendil-works/pi-ai"; import { randomUUID } from "crypto"; import { appendFileSync, @@ -73,6 +73,8 @@ export interface CompactionEntry extends SessionEntryBase { tokensBefore: number; /** Extension-specific data (e.g., ArtifactIndex, version markers for structured compaction) */ details?: T; + /** Usage from the LLM call(s) that generated this summary, if available */ + usage?: Usage; /** True if generated by an extension, undefined/false if pi-generated (backward compatible) */ fromHook?: boolean; } @@ -83,6 +85,8 @@ export interface BranchSummaryEntry extends SessionEntryBase { summary: string; /** Extension-specific data (not sent to LLM) */ details?: T; + /** Usage from the LLM call that generated this summary, if available */ + usage?: Usage; /** True if generated by an extension, false if pi-generated */ fromHook?: boolean; } @@ -1096,6 +1100,7 @@ export class SessionManager { tokensBefore: number, details?: T, fromHook?: boolean, + usage?: Usage, ): string { const entry: CompactionEntry = { type: "compaction", @@ -1106,6 +1111,7 @@ export class SessionManager { firstKeptEntryId, tokensBefore, details, + usage, fromHook, }; this._appendEntry(entry); @@ -1372,7 +1378,13 @@ export class SessionManager { * Same as branch(), but also appends a branch_summary entry that captures * context from the abandoned conversation path. */ - branchWithSummary(branchFromId: string | null, summary: string, details?: unknown, fromHook?: boolean): string { + branchWithSummary( + branchFromId: string | null, + summary: string, + details?: unknown, + fromHook?: boolean, + usage?: Usage, + ): string { if (branchFromId !== null && !this.byId.has(branchFromId)) { throw new Error(`Entry ${branchFromId} not found`); } @@ -1385,6 +1397,7 @@ export class SessionManager { fromId: branchFromId ?? "root", summary, details, + usage, fromHook, }; this._appendEntry(entry); diff --git a/packages/coding-agent/src/core/usage-totals.ts b/packages/coding-agent/src/core/usage-totals.ts new file mode 100644 index 000000000..e65e2eb49 --- /dev/null +++ b/packages/coding-agent/src/core/usage-totals.ts @@ -0,0 +1,27 @@ +import type { Usage } from "@earendil-works/pi-ai/compat"; + +export interface UsageTotals { + input: number; + output: number; + cacheRead: number; + cacheWrite: number; + cost: number; +} + +export function createUsageTotals(): UsageTotals { + return { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + cost: 0, + }; +} + +export function addUsageToTotals(totals: UsageTotals, usage: Usage): void { + totals.input += usage.input; + totals.output += usage.output; + totals.cacheRead += usage.cacheRead; + totals.cacheWrite += usage.cacheWrite; + totals.cost += usage.cost.total; +} diff --git a/packages/coding-agent/src/modes/interactive/components/footer.ts b/packages/coding-agent/src/modes/interactive/components/footer.ts index fa9474806..44cf886b3 100644 --- a/packages/coding-agent/src/modes/interactive/components/footer.ts +++ b/packages/coding-agent/src/modes/interactive/components/footer.ts @@ -3,6 +3,7 @@ import { type Component, truncateToWidth, visibleWidth } from "@earendil-works/p import type { AgentSession } from "../../../core/agent-session.ts"; import { areExperimentalFeaturesEnabled } from "../../../core/experimental.ts"; import type { ReadonlyFooterDataProvider } from "../../../core/footer-data-provider.ts"; +import { addUsageToTotals, createUsageTotals } from "../../../core/usage-totals.ts"; import { theme } from "../theme/theme.ts"; /** @@ -84,25 +85,21 @@ export class FooterComponent implements Component { const state = this.session.state; // Calculate cumulative usage from ALL session entries (not just post-compaction messages) - let totalInput = 0; - let totalOutput = 0; - let totalCacheRead = 0; - let totalCacheWrite = 0; - let totalCost = 0; + const usageTotals = createUsageTotals(); let latestCacheHitRate: number | undefined; for (const entry of this.session.sessionManager.getEntries()) { if (entry.type === "message" && entry.message.role === "assistant") { - totalInput += entry.message.usage.input; - totalOutput += entry.message.usage.output; - totalCacheRead += entry.message.usage.cacheRead; - totalCacheWrite += entry.message.usage.cacheWrite; - totalCost += entry.message.usage.cost.total; + addUsageToTotals(usageTotals, entry.message.usage); const latestPromptTokens = entry.message.usage.input + entry.message.usage.cacheRead + entry.message.usage.cacheWrite; latestCacheHitRate = latestPromptTokens > 0 ? (entry.message.usage.cacheRead / latestPromptTokens) * 100 : undefined; + } else if (entry.type === "message" && entry.message.role === "toolResult" && entry.message.usage) { + addUsageToTotals(usageTotals, entry.message.usage); + } else if ((entry.type === "branch_summary" || entry.type === "compaction") && entry.usage) { + addUsageToTotals(usageTotals, entry.usage); } } @@ -130,19 +127,20 @@ export class FooterComponent implements Component { // Build stats line const statsParts = []; - if (totalInput) statsParts.push(`↑${formatTokens(totalInput)}`); - if (totalOutput) statsParts.push(`↓${formatTokens(totalOutput)}`); - if (totalCacheRead) statsParts.push(`R${formatTokens(totalCacheRead)}`); - if (totalCacheWrite) statsParts.push(`W${formatTokens(totalCacheWrite)}`); - if ((totalCacheRead > 0 || totalCacheWrite > 0) && latestCacheHitRate !== undefined) { + if (usageTotals.input) statsParts.push(`↑${formatTokens(usageTotals.input)}`); + if (usageTotals.output) statsParts.push(`↓${formatTokens(usageTotals.output)}`); + if (usageTotals.cacheRead) statsParts.push(`R${formatTokens(usageTotals.cacheRead)}`); + if (usageTotals.cacheWrite) statsParts.push(`W${formatTokens(usageTotals.cacheWrite)}`); + if ((usageTotals.cacheRead > 0 || usageTotals.cacheWrite > 0) && latestCacheHitRate !== undefined) { statsParts.push(`CH${latestCacheHitRate.toFixed(1)}%`); } + // Kimi Coding is subscription-backed despite using API-key authentication. const usingSubscription = state.model ? state.model.provider === "kimi-coding" || this.session.modelRuntime.isUsingOAuth(state.model.provider) : false; - if (totalCost || usingSubscription) { - const costStr = `$${totalCost.toFixed(3)}${usingSubscription ? " (sub)" : ""}`; + if (usageTotals.cost || usingSubscription) { + const costStr = `$${usageTotals.cost.toFixed(3)}${usingSubscription ? " (sub)" : ""}`; statsParts.push(costStr); } diff --git a/packages/coding-agent/test/agent-session-stats.test.ts b/packages/coding-agent/test/agent-session-stats.test.ts index 9a442c73c..684c5e1cc 100644 --- a/packages/coding-agent/test/agent-session-stats.test.ts +++ b/packages/coding-agent/test/agent-session-stats.test.ts @@ -1,5 +1,5 @@ import { Agent } from "@earendil-works/pi-agent-core"; -import { type AssistantMessage, getModel, type Usage } from "@earendil-works/pi-ai/compat"; +import { type AssistantMessage, getModel, type ToolResultMessage, type Usage } from "@earendil-works/pi-ai/compat"; import { describe, expect, it } from "vitest"; import { AgentSession } from "../src/core/agent-session.ts"; import { AuthStorage } from "../src/core/auth-storage.ts"; @@ -48,6 +48,18 @@ function createUserMessage(text: string, timestamp: number) { }; } +function createToolResultMessage(usage: Usage): ToolResultMessage { + return { + role: "toolResult", + toolCallId: "tool-call-1", + toolName: "test_tool", + content: [{ type: "text", text: "tool result" }], + usage, + isError: false, + timestamp: 1, + }; +} + async function createSession() { const settingsManager = SettingsManager.inMemory(); const sessionManager = SessionManager.inMemory(); @@ -143,6 +155,75 @@ describe("AgentSession.getSessionStats", () => { } }); + it("includes branch summary usage in session totals", async () => { + const { session, sessionManager } = await createSession(); + + try { + sessionManager.branchWithSummary(null, "summary", undefined, false, { + input: 10, + output: 20, + cacheRead: 30, + cacheWrite: 40, + totalTokens: 100, + cost: { input: 0.1, output: 0.2, cacheRead: 0.3, cacheWrite: 0.4, total: 1 }, + }); + syncAgentMessages(session, sessionManager); + + const stats = session.getSessionStats(); + expect(stats.tokens).toEqual({ input: 10, output: 20, cacheRead: 30, cacheWrite: 40, total: 100 }); + expect(stats.cost).toBe(1); + } finally { + session.dispose(); + } + }); + + it("includes compaction usage in session totals", async () => { + const { session, sessionManager } = await createSession(); + + try { + const firstKeptEntryId = sessionManager.appendMessage(createUserMessage("hello", 1)); + sessionManager.appendCompaction("summary", firstKeptEntryId, 100, undefined, false, { + input: 10, + output: 20, + cacheRead: 30, + cacheWrite: 40, + totalTokens: 100, + cost: { input: 0.1, output: 0.2, cacheRead: 0.3, cacheWrite: 0.4, total: 1 }, + }); + syncAgentMessages(session, sessionManager); + + const stats = session.getSessionStats(); + expect(stats.tokens).toEqual({ input: 10, output: 20, cacheRead: 30, cacheWrite: 40, total: 100 }); + expect(stats.cost).toBe(1); + } finally { + session.dispose(); + } + }); + + it("includes tool result usage in session totals", async () => { + const { session, sessionManager } = await createSession(); + + try { + sessionManager.appendMessage( + createToolResultMessage({ + input: 10, + output: 20, + cacheRead: 30, + cacheWrite: 40, + totalTokens: 100, + cost: { input: 0.1, output: 0.2, cacheRead: 0.3, cacheWrite: 0.4, total: 1 }, + }), + ); + syncAgentMessages(session, sessionManager); + + const stats = session.getSessionStats(); + expect(stats.tokens).toEqual({ input: 10, output: 20, cacheRead: 30, cacheWrite: 40, total: 100 }); + expect(stats.cost).toBe(1); + } finally { + session.dispose(); + } + }); + it("ignores zero-usage messages when checking for post-compaction context usage", async () => { const { session, sessionManager } = await createSession(); diff --git a/packages/coding-agent/test/branch-summary-extensions.test.ts b/packages/coding-agent/test/branch-summary-extensions.test.ts new file mode 100644 index 000000000..9c47cb81a --- /dev/null +++ b/packages/coding-agent/test/branch-summary-extensions.test.ts @@ -0,0 +1,55 @@ +import type { Usage } from "@earendil-works/pi-ai/compat"; +import { afterEach, describe, expect, it } from "vitest"; +import { createHarness, type Harness } from "./suite/harness.ts"; +import { assistantMsg, userMsg } from "./utilities.ts"; + +describe("Branch summary extensions", () => { + const harnesses: Harness[] = []; + + afterEach(() => { + while (harnesses.length > 0) { + harnesses.pop()?.cleanup(); + } + }); + + it("persists extension-provided summary usage in session totals", async () => { + const usage: Usage = { + input: 10, + output: 20, + cacheRead: 30, + cacheWrite: 40, + totalTokens: 100, + cost: { input: 0.1, output: 0.2, cacheRead: 0.3, cacheWrite: 0.4, total: 1 }, + }; + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.on("session_before_tree", () => ({ + summary: { + summary: "Summary provided by extension", + usage, + }, + })); + }, + ], + }); + harnesses.push(harness); + + const targetId = harness.sessionManager.appendMessage(userMsg("first branch")); + harness.sessionManager.appendMessage(assistantMsg("first reply")); + harness.sessionManager.appendMessage(userMsg("abandoned branch work")); + harness.sessionManager.appendMessage(assistantMsg("abandoned reply")); + + const result = await harness.session.navigateTree(targetId, { summarize: true }); + const summaryEntry = result.summaryEntry; + + expect(summaryEntry?.type).toBe("branch_summary"); + expect(summaryEntry?.fromHook).toBe(true); + expect(summaryEntry?.summary).toBe("Summary provided by extension"); + expect(summaryEntry?.usage).toEqual(usage); + + const stats = harness.session.getSessionStats(); + expect(stats.tokens).toEqual({ input: 12, output: 22, cacheRead: 30, cacheWrite: 40, total: 104 }); + expect(stats.cost).toBe(1); + }); +}); diff --git a/packages/coding-agent/test/compaction-summary-reasoning.test.ts b/packages/coding-agent/test/compaction-summary-reasoning.test.ts index 306a3b1c4..eb1d35794 100644 --- a/packages/coding-agent/test/compaction-summary-reasoning.test.ts +++ b/packages/coding-agent/test/compaction-summary-reasoning.test.ts @@ -57,7 +57,7 @@ describe("generateSummary reasoning options", () => { }); it("uses the provided thinking level for reasoning-capable models", async () => { - await generateSummary( + const result = await generateSummary( messages, createModel(true), 2000, @@ -69,6 +69,9 @@ describe("generateSummary reasoning options", () => { "medium", ); + expect(result.text).toBe("## Goal\nTest summary"); + expect(result.usage).toEqual(mockSummaryResponse.usage); + expect(completeSimpleMock).toHaveBeenCalledTimes(1); expect(completeSimpleMock.mock.calls[0][2]).toMatchObject({ reasoning: "medium", @@ -127,8 +130,15 @@ describe("generateSummary reasoning options", () => { settings: { enabled: true, reserveTokens: 500000, keepRecentTokens: 20000 }, }; - await compact(preparation, createModel(false, 128000), "test-key"); + const result = await compact(preparation, createModel(false, 128000), "test-key"); + expect(result.usage).toEqual({ + ...mockSummaryResponse.usage, + input: 20, + output: 20, + totalTokens: 40, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }); expect(completeSimpleMock.mock.calls.map((call) => call[2]?.maxTokens)).toEqual([128000, 128000]); }); }); diff --git a/packages/coding-agent/test/footer-width.test.ts b/packages/coding-agent/test/footer-width.test.ts index 85956d539..310114a47 100644 --- a/packages/coding-agent/test/footer-width.test.ts +++ b/packages/coding-agent/test/footer-width.test.ts @@ -21,20 +21,46 @@ function createSession(options: { reasoning?: boolean; thinkingLevel?: string; usage?: AssistantUsage; + branchUsage?: AssistantUsage; + compactionUsage?: AssistantUsage; + toolUsage?: AssistantUsage; }): AgentSession { const usage = options.usage; - const entries = - usage === undefined - ? [] - : [ - { - type: "message", - message: { - role: "assistant", - usage, - }, - }, - ]; + const entries: Array> = []; + + if (usage !== undefined) { + entries.push({ + type: "message", + message: { + role: "assistant", + usage, + }, + }); + } + + if (options.branchUsage !== undefined) { + entries.push({ + type: "branch_summary", + usage: options.branchUsage, + }); + } + + if (options.compactionUsage !== undefined) { + entries.push({ + type: "compaction", + usage: options.compactionUsage, + }); + } + + if (options.toolUsage !== undefined) { + entries.push({ + type: "message", + message: { + role: "toolResult", + usage: options.toolUsage, + }, + }); + } const session = { state: { @@ -125,6 +151,44 @@ describe("FooterComponent width handling", () => { } }); + it("includes summary and tool result usage in the total cost", () => { + const session = createSession({ + sessionName: "", + usage: { + input: 100, + output: 10, + cacheRead: 0, + cacheWrite: 0, + cost: { total: 0.5 }, + }, + branchUsage: { + input: 20, + output: 5, + cacheRead: 0, + cacheWrite: 0, + cost: { total: 0.25 }, + }, + compactionUsage: { + input: 5, + output: 2, + cacheRead: 0, + cacheWrite: 0, + cost: { total: 0.125 }, + }, + toolUsage: { + input: 15, + output: 3, + cacheRead: 0, + cacheWrite: 0, + cost: { total: 0.375 }, + }, + }); + const footer = new FooterComponent(session, createFooterData(1)); + + const statsLine = stripAnsi(footer.render(120)[1]); + expect(statsLine).toContain("$1.250"); + }); + it("shows the latest cache hit rate when cache usage is present", () => { const session = createSession({ sessionName: "", diff --git a/packages/coding-agent/test/session-manager/tree-traversal.test.ts b/packages/coding-agent/test/session-manager/tree-traversal.test.ts index e47ec9095..9f447de1b 100644 --- a/packages/coding-agent/test/session-manager/tree-traversal.test.ts +++ b/packages/coding-agent/test/session-manager/tree-traversal.test.ts @@ -71,7 +71,15 @@ describe("SessionManager append and tree traversal", () => { const id1 = session.appendMessage(userMsg("1")); const id2 = session.appendMessage(assistantMsg("2")); - const compactionId = session.appendCompaction("summary", id1, 1000); + const usage = { + input: 10, + output: 20, + cacheRead: 30, + cacheWrite: 40, + totalTokens: 100, + cost: { input: 0.1, output: 0.2, cacheRead: 0.3, cacheWrite: 0.4, total: 1 }, + }; + const compactionId = session.appendCompaction("summary", id1, 1000, undefined, false, usage); const _id3 = session.appendMessage(userMsg("3")); const entries = session.getEntries(); @@ -83,6 +91,7 @@ describe("SessionManager append and tree traversal", () => { expect(compactionEntry.summary).toBe("summary"); expect(compactionEntry.firstKeptEntryId).toBe(id1); expect(compactionEntry.tokensBefore).toBe(1000); + expect(compactionEntry.usage).toEqual(usage); } expect(entries[3].parentId).toBe(compactionId); @@ -319,7 +328,15 @@ describe("SessionManager append and tree traversal", () => { const _id2 = session.appendMessage(assistantMsg("2")); const _id3 = session.appendMessage(userMsg("3")); - const summaryId = session.branchWithSummary(id1, "Summary of abandoned work"); + const usage = { + input: 10, + output: 20, + cacheRead: 30, + cacheWrite: 40, + totalTokens: 100, + cost: { input: 0.1, output: 0.2, cacheRead: 0.3, cacheWrite: 0.4, total: 1 }, + }; + const summaryId = session.branchWithSummary(id1, "Summary of abandoned work", undefined, false, usage); expect(session.getLeafId()).toBe(summaryId); @@ -329,6 +346,7 @@ describe("SessionManager append and tree traversal", () => { expect(summaryEntry?.parentId).toBe(id1); if (summaryEntry?.type === "branch_summary") { expect(summaryEntry.summary).toBe("Summary of abandoned work"); + expect(summaryEntry.usage).toEqual(usage); } }); diff --git a/packages/coding-agent/test/suite/agent-session-compaction.test.ts b/packages/coding-agent/test/suite/agent-session-compaction.test.ts index ba5ae5c72..ee7bea77a 100644 --- a/packages/coding-agent/test/suite/agent-session-compaction.test.ts +++ b/packages/coding-agent/test/suite/agent-session-compaction.test.ts @@ -97,6 +97,14 @@ describe("AgentSession compaction characterization", () => { }); it("manually compacts using an extension-provided summary", async () => { + const summaryUsage = { + input: 10, + output: 20, + cacheRead: 30, + cacheWrite: 40, + totalTokens: 100, + cost: { input: 0.1, output: 0.2, cacheRead: 0.3, cacheWrite: 0.4, total: 1 }, + }; const harness = await createHarness({ settings: { compaction: { keepRecentTokens: 1 } }, extensionFactories: [ @@ -106,6 +114,7 @@ describe("AgentSession compaction characterization", () => { summary: "summary from extension", firstKeptEntryId: event.preparation.firstKeptEntryId, tokensBefore: event.preparation.tokensBefore, + usage: summaryUsage, details: { source: "extension" }, }, })); @@ -116,14 +125,26 @@ describe("AgentSession compaction characterization", () => { await harness.session.prompt("one"); await harness.session.prompt("two"); + const statsBefore = harness.session.getSessionStats(); const result = await harness.session.compact(); const compactionEntries = harness.sessionManager.getEntries().filter((entry) => entry.type === "compaction"); const estimatedTokensAfter = harness.session.messages.reduce((sum, message) => sum + estimateTokens(message), 0); expect(result.summary).toBe("summary from extension"); + expect(result.usage).toEqual(summaryUsage); expect(result.estimatedTokensAfter).toBe(estimatedTokensAfter); expect(compactionEntries).toHaveLength(1); + const compactionEntry = compactionEntries[0]; + if (compactionEntry?.type === "compaction") { + expect(compactionEntry.usage).toEqual(summaryUsage); + } + const statsAfter = harness.session.getSessionStats(); + expect(statsAfter.tokens.input).toBe(statsBefore.tokens.input + summaryUsage.input); + expect(statsAfter.tokens.output).toBe(statsBefore.tokens.output + summaryUsage.output); + expect(statsAfter.tokens.cacheRead).toBe(statsBefore.tokens.cacheRead + summaryUsage.cacheRead); + expect(statsAfter.tokens.cacheWrite).toBe(statsBefore.tokens.cacheWrite + summaryUsage.cacheWrite); + expect(statsAfter.cost).toBe(statsBefore.cost + summaryUsage.cost.total); expect(harness.session.messages[0]?.role).toBe("compactionSummary"); }); @@ -154,6 +175,22 @@ describe("AgentSession compaction characterization", () => { expect(getStreamCallCount()).toBe(1); }); + it("persists usage from pi-generated manual compaction", async () => { + const harness = await createHarness({ withConfiguredAuth: false }); + harnesses.push(harness); + seedCompactableSession(harness); + useSummaryStreamFn(harness, "summary from custom stream"); + + const result = await harness.session.compact(); + + const compactionEntries = harness.sessionManager.getEntries().filter((entry) => entry.type === "compaction"); + expect(result.usage).toEqual(createUsage(10)); + expect(compactionEntries).toHaveLength(1); + expect(compactionEntries[0]?.type === "compaction" ? compactionEntries[0].usage : undefined).toEqual( + createUsage(10), + ); + }); + it("auto-compacts with a custom streamFn when registry auth is absent", async () => { const harness = await createHarness({ withConfiguredAuth: false }); harnesses.push(harness); diff --git a/packages/coding-agent/test/suite/agent-session-model-extension.test.ts b/packages/coding-agent/test/suite/agent-session-model-extension.test.ts index 39a9243c1..b334b1232 100644 --- a/packages/coding-agent/test/suite/agent-session-model-extension.test.ts +++ b/packages/coding-agent/test/suite/agent-session-model-extension.test.ts @@ -1,5 +1,5 @@ import type { AgentTool, ThinkingLevel } from "@earendil-works/pi-agent-core"; -import { fauxAssistantMessage, fauxToolCall, type Model } from "@earendil-works/pi-ai"; +import { fauxAssistantMessage, fauxToolCall, type Model, type Usage } from "@earendil-works/pi-ai"; import { Type } from "typebox"; import { afterEach, describe, expect, it } from "vitest"; import type { BuildSystemPromptOptions, ExtensionAPI } from "../../src/index.ts"; @@ -156,6 +156,23 @@ describe("AgentSession model and extension characterization", () => { }); it("allows extension tool_result handlers to modify tool results", async () => { + const toolUsage: Usage = { + input: 1, + output: 2, + cacheRead: 3, + cacheWrite: 4, + totalTokens: 10, + cost: { input: 0.1, output: 0.2, cacheRead: 0.3, cacheWrite: 0.4, total: 1 }, + }; + const patchedToolUsage: Usage = { + input: 5, + output: 6, + cacheRead: 7, + cacheWrite: 8, + totalTokens: 26, + cost: { input: 0.5, output: 0.6, cacheRead: 0.7, cacheWrite: 0.8, total: 2.6 }, + }; + let observedToolUsage: Usage | undefined; const echoTool: AgentTool = { name: "echo", label: "Echo", @@ -163,17 +180,21 @@ describe("AgentSession model and extension characterization", () => { parameters: Type.Object({ text: Type.String() }), execute: async (_toolCallId, params) => { const text = typeof params === "object" && params !== null && "text" in params ? String(params.text) : ""; - return { content: [{ type: "text", text }], details: { text } }; + return { content: [{ type: "text", text }], details: { text }, usage: toolUsage }; }, }; const harness = await createHarness({ tools: [echoTool], extensionFactories: [ (pi) => { - pi.on("tool_result", async () => ({ - content: [{ type: "text", text: "patched result" }], - details: { patched: true }, - })); + pi.on("tool_result", async (event) => { + observedToolUsage = event.usage; + return { + content: [{ type: "text", text: "patched result" }], + details: { patched: true }, + usage: patchedToolUsage, + }; + }); }, ], }); @@ -196,9 +217,12 @@ describe("AgentSession model and extension characterization", () => { await harness.session.prompt("hi"); expect(getAssistantTexts(harness)).toContain("patched result"); - expect( - harness.session.messages.find((message) => message.role === "toolResult" && message.details?.patched === true), - ).toBeDefined(); + const toolResult = harness.session.messages.find( + (message) => message.role === "toolResult" && message.details?.patched === true, + ); + expect(observedToolUsage).toEqual(toolUsage); + expect(toolResult).toBeDefined(); + expect(toolResult?.role === "toolResult" ? toolResult.usage : undefined).toEqual(patchedToolUsage); }); it("allows extension context handlers to modify messages before the LLM call", async () => { diff --git a/packages/coding-agent/test/suite/regressions/6324-branch-summary-ambient-auth.test.ts b/packages/coding-agent/test/suite/regressions/6324-branch-summary-ambient-auth.test.ts index 121cc006d..bfedbaafd 100644 --- a/packages/coding-agent/test/suite/regressions/6324-branch-summary-ambient-auth.test.ts +++ b/packages/coding-agent/test/suite/regressions/6324-branch-summary-ambient-auth.test.ts @@ -37,7 +37,7 @@ describe("issue #6324 branch summary ambient auth", () => { cacheRead: 0, cacheWrite: 0, totalTokens: 2, - cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0.25 }, }, stopReason: "stop", timestamp: Date.now(), @@ -57,5 +57,6 @@ describe("issue #6324 branch summary ambient auth", () => { expect(streamCallCount).toBe(1); expect(result.summaryEntry?.type).toBe("branch_summary"); expect(result.summaryEntry?.summary).toContain("branch summary text"); + expect(result.summaryEntry?.usage?.cost.total).toBe(0.25); }); }); From f8b74a4507a0505cc8275a8f5afcdd9072cc75f0 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Mon, 20 Jul 2026 17:04:12 +0200 Subject: [PATCH 24/62] fix: complete extension usage accounting closes #6509 --- packages/agent/CHANGELOG.md | 4 ++ .../src/harness/compaction/compaction.ts | 28 ++++++++++- packages/agent/src/index.ts | 1 + .../agent/test/harness/compaction.test.ts | 11 ++++- packages/ai/CHANGELOG.md | 1 + packages/coding-agent/CHANGELOG.md | 1 + packages/coding-agent/README.md | 2 +- packages/coding-agent/docs/compaction.md | 11 +++-- packages/coding-agent/docs/extensions.md | 18 ++++++-- packages/coding-agent/docs/rpc.md | 30 +++++++++++- packages/coding-agent/docs/session-format.md | 3 ++ packages/coding-agent/docs/usage.md | 2 +- .../examples/extensions/custom-compaction.ts | 1 + .../src/core/compaction/compaction.ts | 35 +++++++++++++- .../coding-agent/src/core/usage-totals.ts | 43 +++++++++++++++++ packages/coding-agent/src/index.ts | 1 + .../src/modes/interactive/interactive-mode.ts | 24 +++------- .../test/agent-session-stats.test.ts | 26 +++++++++++ .../test/compaction-summary-reasoning.test.ts | 15 +++++- .../session-manager/tree-traversal.test.ts | 46 +++++++++++++++++++ 20 files changed, 267 insertions(+), 36 deletions(-) diff --git a/packages/agent/CHANGELOG.md b/packages/agent/CHANGELOG.md index eba2a62a8..3e8ec724a 100644 --- a/packages/agent/CHANGELOG.md +++ b/packages/agent/CHANGELOG.md @@ -6,6 +6,10 @@ - Moved the `uuidv7` export to `@earendil-works/pi-ai`. +### Added + +- Added usage metadata to tool results, compaction entries, and branch summaries in the agent harness ([#6671](https://github.com/earendil-works/pi/pull/6671) by [@davidbrai](https://github.com/davidbrai)). + ## [0.80.10] - 2026-07-16 ## [0.80.9] - 2026-07-16 diff --git a/packages/agent/src/harness/compaction/compaction.ts b/packages/agent/src/harness/compaction/compaction.ts index 58938c13c..6c21a492f 100644 --- a/packages/agent/src/harness/compaction/compaction.ts +++ b/packages/agent/src/harness/compaction/compaction.ts @@ -499,6 +499,30 @@ export async function generateSummary( customInstructions?: string, previousSummary?: string, thinkingLevel?: ThinkingLevel, +): Promise> { + const result = await generateSummaryWithUsage( + currentMessages, + models, + model, + reserveTokens, + signal, + customInstructions, + previousSummary, + thinkingLevel, + ); + return result.ok ? ok(result.value.text) : err(result.error); +} + +/** Generate or update a conversation summary and return its provider usage. */ +export async function generateSummaryWithUsage( + currentMessages: AgentMessage[], + models: Models, + model: Model, + reserveTokens: number, + signal?: AbortSignal, + customInstructions?: string, + previousSummary?: string, + thinkingLevel?: ThinkingLevel, ): Promise> { const maxTokens = Math.min( Math.floor(0.8 * reserveTokens), @@ -687,7 +711,7 @@ export async function compact( let historyText = "No prior history."; let historyUsage: Usage | undefined; if (messagesToSummarize.length > 0) { - const historyResult = await generateSummary( + const historyResult = await generateSummaryWithUsage( messagesToSummarize, models, model, @@ -715,7 +739,7 @@ export async function compact( ? combineUsage(historyUsage, turnPrefixResult.value.usage) : turnPrefixResult.value.usage; } else { - const summaryResult = await generateSummary( + const summaryResult = await generateSummaryWithUsage( messagesToSummarize, models, model, diff --git a/packages/agent/src/index.ts b/packages/agent/src/index.ts index 49baf081f..b5ec7807f 100644 --- a/packages/agent/src/index.ts +++ b/packages/agent/src/index.ts @@ -20,6 +20,7 @@ export { findCutPoint, findTurnStartIndex, generateSummary, + generateSummaryWithUsage, getLastAssistantUsage, prepareCompaction, serializeConversation, diff --git a/packages/agent/test/harness/compaction.test.ts b/packages/agent/test/harness/compaction.test.ts index 7cdea0c93..c8160307d 100644 --- a/packages/agent/test/harness/compaction.test.ts +++ b/packages/agent/test/harness/compaction.test.ts @@ -20,6 +20,7 @@ import { findCutPoint, findTurnStartIndex, generateSummary, + generateSummaryWithUsage, getLastAssistantUsage, prepareCompaction, serializeConversation, @@ -510,7 +511,7 @@ describe("harness compaction", () => { ]); const summary = getOrThrow( - await generateSummary(messages, models, model, 2000, undefined, "focus", "old summary"), + await generateSummaryWithUsage(messages, models, model, 2000, undefined, "focus", "old summary"), ); expect(summary.text).toContain("Test summary"); @@ -523,6 +524,14 @@ describe("harness compaction", () => { expect(promptText).toContain("Additional focus: focus"); }); + it("preserves the string result from generateSummary", async () => { + const messages: AgentMessage[] = [createUserMessage("Summarize this.")]; + const { faux, model } = createFauxModel(false); + faux.setResponses([fauxAssistantMessage("## Goal\nTest summary")]); + + expect(getOrThrow(await generateSummary(messages, models, model, 2000))).toBe("## Goal\nTest summary"); + }); + it("returns error results for failed or aborted summary generations", async () => { const messages: AgentMessage[] = [createUserMessage("Summarize this.")]; const { faux: errorFaux, model: errorModel } = createFauxModel(false); diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 7739597b3..56693a287 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -6,6 +6,7 @@ - Added `contentText` for extracting joined text from message content. - Added a shared `uuidv7` utility for time-ordered identifiers. +- Added optional usage metadata to tool result messages ([#6671](https://github.com/earendil-works/pi/pull/6671) by [@davidbrai](https://github.com/davidbrai)). ### Fixed diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 15bdbf80d..cac4d0193 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -6,6 +6,7 @@ - Added built-in llama.cpp router support with `/login` connection setup and `/llama` Hugging Face model search and downloads, explicit loading, unloading, and live progress. See [llama.cpp](docs/llama-cpp.md). - Added extension registration for complete pi-ai providers, including native authentication, model refresh, filtering, and streaming behavior. +- Added usage accounting for tools, compaction, and branch summaries in persisted sessions, footer totals, and session statistics ([#6671](https://github.com/earendil-works/pi/pull/6671) by [@davidbrai](https://github.com/davidbrai)). ### Fixed diff --git a/packages/coding-agent/README.md b/packages/coding-agent/README.md index 3454fac25..2198d7240 100644 --- a/packages/coding-agent/README.md +++ b/packages/coding-agent/README.md @@ -152,7 +152,7 @@ The interface from top to bottom: - **Startup header** - Shows shortcuts (`/hotkeys` for all), loaded AGENTS.md files, prompt templates, skills, and extensions - **Messages** - Your messages, assistant responses, tool calls and results, notifications, errors, and extension UI - **Editor** - Where you type; border color indicates thinking level -- **Footer** - Working directory, session name, total token/cache usage (`↑` input, `↓` output, `R` cache read, `W` cache write, `CH` latest cache hit rate), cost, context usage, current model +- **Footer** - Working directory, session name, total token/cache usage (`↑` input, `↓` output, `R` cache read, `W` cache write, `CH` latest cache hit rate), cost, context usage, current model. Totals include assistant responses, usage reported by tools, and summary generation. The editor can be temporarily replaced by other UI, like built-in `/settings` or custom UI from extensions (e.g., a Q&A tool that lets the user answer model questions in a structured format). [Extensions](#extensions) can also replace the editor, add widgets above/below it, a status line, custom footer, or overlays. diff --git a/packages/coding-agent/docs/compaction.md b/packages/coding-agent/docs/compaction.md index 5e0d4eefe..d55c21364 100644 --- a/packages/coding-agent/docs/compaction.md +++ b/packages/coding-agent/docs/compaction.md @@ -129,6 +129,7 @@ interface CompactionEntry { summary: string; firstKeptEntryId: string; tokensBefore: number; + usage?: Usage; // LLM usage that generated the summary fromHook?: boolean; // true if provided by extension (legacy field name) details?: T; // implementation-specific data } @@ -140,9 +141,9 @@ interface CompactionDetails { } ``` -Extensions can store any JSON-serializable data in `details`. The default compaction tracks file operations, but custom extension implementations can use their own structure. +Extensions can store any JSON-serializable data in `details`. The default compaction tracks file operations, but custom extension implementations can use their own structure. Generated and extension-provided summaries store their LLM `usage` when available so session totals include summarization work. -See [`prepareCompaction()`](https://github.com/earendil-works/pi-mono/blob/main/packages/coding-agent/src/core/compaction/compaction.ts) and [`compact()`](https://github.com/earendil-works/pi-mono/blob/main/packages/coding-agent/src/core/compaction/compaction.ts) for the implementation. +See [`prepareCompaction()`](https://github.com/earendil-works/pi-mono/blob/main/packages/coding-agent/src/core/compaction/compaction.ts) and [`compact()`](https://github.com/earendil-works/pi-mono/blob/main/packages/coding-agent/src/core/compaction/compaction.ts) for the implementation. For direct programmatic summarization, `generateSummary()` returns the summary text and `generateSummaryWithUsage()` returns `{ text, usage }`. ## Branch Summarization @@ -195,6 +196,7 @@ interface BranchSummaryEntry { timestamp: number; summary: string; fromId: string; // Entry we navigated from + usage?: Usage; // LLM usage that generated the summary fromHook?: boolean; // true if provided by extension (legacy field name) details?: T; // implementation-specific data } @@ -300,6 +302,7 @@ pi.on("session_before_compact", async (event, ctx) => { summary: "Your summary...", firstKeptEntryId: preparation.firstKeptEntryId, tokensBefore: preparation.tokensBefore, + // usage: summaryResponse.usage, // Optional; included in session totals details: { /* custom data */ }, } }; @@ -328,13 +331,14 @@ pi.on("session_before_compact", async (event, ctx) => { // [Tool result]: output text // Now send to your model for summarization - const summary = await myModel.summarize(conversationText); + const { summary, usage } = await myModel.summarize(conversationText); return { compaction: { summary, firstKeptEntryId: preparation.firstKeptEntryId, tokensBefore: preparation.tokensBefore, + usage, } }; }); @@ -364,6 +368,7 @@ pi.on("session_before_tree", async (event, ctx) => { return { summary: { summary: "Your summary...", + // usage: summaryResponse.usage, // Optional; included in session totals details: { /* custom data */ }, } }; diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index e51b056e3..6aec11960 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/packages/coding-agent/docs/extensions.md @@ -468,6 +468,7 @@ pi.on("session_before_compact", async (event, ctx) => { summary: "...", firstKeptEntryId: preparation.firstKeptEntryId, tokensBefore: preparation.tokensBefore, + // usage: summaryResponse.usage, // Optional; included in session totals } }; }); @@ -489,7 +490,13 @@ pi.on("session_before_tree", async (event, ctx) => { const { preparation, signal } = event; return { cancel: true }; // OR provide custom summary: - return { summary: { summary: "...", details: {} } }; + return { + summary: { + summary: "...", + // usage: summaryResponse.usage, // Optional; included in session totals + details: {}, + }, + }; }); pi.on("session_tree", async (event, ctx) => { @@ -813,7 +820,7 @@ In parallel tool mode, `tool_result` and `tool_execution_end` may interleave in `tool_result` handlers chain like middleware: - Handlers run in extension load order - Each handler sees the latest result after previous handler changes -- Handlers can return partial patches (`content`, `details`, or `isError`); omitted fields keep their current values +- Handlers can return partial patches (`content`, `details`, `isError`, or `usage`); omitted fields keep their current values Use `ctx.signal` for nested async work inside the handler. This lets Esc cancel model calls, `fetch()`, and other abort-aware operations started by the extension. @@ -822,7 +829,7 @@ import { isBashToolResult } from "@earendil-works/pi-coding-agent"; pi.on("tool_result", async (event, ctx) => { // event.toolName, event.toolCallId, event.input - // event.content, event.details, event.isError + // event.content, event.details, event.isError, event.usage if (isBashToolResult(event)) { // event.details is typed as BashToolDetails @@ -835,7 +842,7 @@ pi.on("tool_result", async (event, ctx) => { }); // Modify result: - return { content: [...], details: {...}, isError: false }; + return { content: [...], details: {...}, isError: false, usage: nestedModelUsage }; }); ``` @@ -1932,6 +1939,7 @@ pi.registerTool({ return { content: [{ type: "text", text: "Done" }], // Sent to LLM details: { data: result }, // For rendering & state + // usage: nestedModelResponse.usage, // Optional nested LLM usage // Optional: stop after this tool batch when every finalized tool result // in the batch also returns terminate: true. terminate: true, @@ -1944,6 +1952,8 @@ pi.registerTool({ }); ``` +**Usage accounting:** If a tool makes nested LLM calls, return their combined `Usage` as `usage`. Pi persists it on the tool result and includes it in footer, `/session`, and RPC session totals. `tool_result` handlers can inspect or replace this value. + **Signaling errors:** To mark a tool execution as failed (sets `isError: true` on the result and reports it to the LLM), throw an error from `execute`. Returning a value never sets the error flag regardless of what properties you include in the return object. **Early termination:** Return `terminate: true` from `execute()` to hint that the automatic follow-up LLM call should be skipped after the current tool batch. This only takes effect when every finalized tool result in that batch is terminating. See [examples/extensions/structured-output.ts](../examples/extensions/structured-output.ts) for a minimal example where the agent ends on a final structured-output tool call. diff --git a/packages/coding-agent/docs/rpc.md b/packages/coding-agent/docs/rpc.md index 8074ceb95..2151b4b3b 100644 --- a/packages/coding-agent/docs/rpc.md +++ b/packages/coding-agent/docs/rpc.md @@ -395,12 +395,20 @@ Response: "firstKeptEntryId": "abc123", "tokensBefore": 150000, "estimatedTokensAfter": 32000, + "usage": { + "input": 32000, + "output": 1200, + "cacheRead": 0, + "cacheWrite": 0, + "totalTokens": 33200, + "cost": {"input": 0.01, "output": 0.02, "cacheRead": 0, "cacheWrite": 0, "total": 0.03} + }, "details": {} } } ``` -`estimatedTokensAfter` is a heuristic estimate over the rebuilt message context immediately after compaction, not a provider-exact token count. +`estimatedTokensAfter` is a heuristic estimate over the rebuilt message context immediately after compaction, not a provider-exact token count. `usage` reports the LLM call or calls that generated the summary and may be omitted by custom compaction handlers. #### set_auto_compaction @@ -557,7 +565,7 @@ Response: } ``` -`tokens` contains assistant usage totals for the current session state. `contextUsage` contains the actual current context-window estimate used for compaction and footer display. +`tokens` and `cost` include assistant messages, usage reported by tools, and compaction/branch-summary generation across the full session. `contextUsage` contains the actual current context-window estimate used for compaction and footer display. `contextUsage` is omitted when no model or context window is available. `contextUsage.tokens` and `contextUsage.percent` are `null` immediately after compaction until a fresh post-compaction assistant response provides valid usage data. @@ -1016,6 +1024,14 @@ The `reason` field is `"manual"`, `"threshold"`, or `"overflow"`. "firstKeptEntryId": "abc123", "tokensBefore": 150000, "estimatedTokensAfter": 32000, + "usage": { + "input": 32000, + "output": 1200, + "cacheRead": 0, + "cacheWrite": 0, + "totalTokens": 33200, + "cost": {"input": 0.01, "output": 0.02, "cacheRead": 0, "cacheWrite": 0, "total": 0.03} + }, "details": {} }, "aborted": false, @@ -1368,11 +1384,21 @@ Stop reasons: `"stop"`, `"length"`, `"toolUse"`, `"error"`, `"aborted"` "toolCallId": "call_123", "toolName": "bash", "content": [{"type": "text", "text": "total 48\ndrwxr-xr-x ..."}], + "usage": { + "input": 100, + "output": 50, + "cacheRead": 0, + "cacheWrite": 0, + "totalTokens": 150, + "cost": {"input": 0.0003, "output": 0.00075, "cacheRead": 0, "cacheWrite": 0, "total": 0.00105} + }, "isError": false, "timestamp": 1733234567890 } ``` +`usage` is optional and reports nested LLM work performed by the tool. When present, it contributes to session token and cost totals. + ### BashExecutionMessage Created by the `bash` RPC command (not by LLM tool calls): diff --git a/packages/coding-agent/docs/session-format.md b/packages/coding-agent/docs/session-format.md index a4afb3180..b27b8b701 100644 --- a/packages/coding-agent/docs/session-format.md +++ b/packages/coding-agent/docs/session-format.md @@ -96,6 +96,7 @@ interface ToolResultMessage { toolName: string; content: (TextContent | ImageContent)[]; details?: any; // Tool-specific metadata + usage?: Usage; // Nested LLM work performed by the tool isError: boolean; timestamp: number; } @@ -232,6 +233,7 @@ Created when context is compacted. Stores a summary of earlier messages. ``` Optional fields: +- `usage`: LLM usage from generating the summary; included in session token and cost totals - `details`: Implementation-specific data (e.g., `{ readFiles: string[], modifiedFiles: string[] }` for default, or custom data for extensions) - `fromHook`: `true` if generated by an extension, `false`/`undefined` if pi-generated (legacy field name) @@ -244,6 +246,7 @@ Created when switching branches via `/tree` with an LLM generated summary of the ``` Optional fields: +- `usage`: LLM usage from generating the summary; included in session token and cost totals - `details`: File tracking data (`{ readFiles: string[], modifiedFiles: string[] }`) for default, or custom data for extensions - `fromHook`: `true` if generated by an extension, `false`/`undefined` if pi-generated (legacy field name) diff --git a/packages/coding-agent/docs/usage.md b/packages/coding-agent/docs/usage.md index 48cf16b17..c9df34efe 100644 --- a/packages/coding-agent/docs/usage.md +++ b/packages/coding-agent/docs/usage.md @@ -11,7 +11,7 @@ The interface has four main areas: - **Startup header** - shortcuts, loaded context files, prompt templates, skills, and extensions - **Messages** - user messages, assistant responses, tool calls, tool results, notifications, errors, and extension UI - **Editor** - where you type; border color indicates the current thinking level -- **Footer** - working directory, session name, token/cache usage, cost, context usage, and current model +- **Footer** - working directory, session name, token/cache usage, cost, context usage, and current model. Totals include assistant responses, usage reported by tools, and summary generation. The editor can be replaced temporarily by built-in UI such as `/settings` or by custom extension UI. diff --git a/packages/coding-agent/examples/extensions/custom-compaction.ts b/packages/coding-agent/examples/extensions/custom-compaction.ts index fdc2128c6..e37ce976a 100644 --- a/packages/coding-agent/examples/extensions/custom-compaction.ts +++ b/packages/coding-agent/examples/extensions/custom-compaction.ts @@ -116,6 +116,7 @@ ${conversationText} summary, firstKeptEntryId, tokensBefore, + usage: response.usage, }, }; } catch (error) { diff --git a/packages/coding-agent/src/core/compaction/compaction.ts b/packages/coding-agent/src/core/compaction/compaction.ts index dafcb09ea..ffad75e85 100644 --- a/packages/coding-agent/src/core/compaction/compaction.ts +++ b/packages/coding-agent/src/core/compaction/compaction.ts @@ -581,6 +581,37 @@ export async function generateSummary( thinkingLevel?: ThinkingLevel, streamFn?: StreamFn, env?: Record, +): Promise { + return ( + await generateSummaryWithUsage( + currentMessages, + model, + reserveTokens, + apiKey, + headers, + signal, + customInstructions, + previousSummary, + thinkingLevel, + streamFn, + env, + ) + ).text; +} + +/** Generate or update a conversation summary and return its provider usage. */ +export async function generateSummaryWithUsage( + currentMessages: AgentMessage[], + model: Model, + reserveTokens: number, + apiKey: string | undefined, + headers?: Record, + signal?: AbortSignal, + customInstructions?: string, + previousSummary?: string, + thinkingLevel?: ThinkingLevel, + streamFn?: StreamFn, + env?: Record, ): Promise<{ text: string; usage: Usage }> { const maxTokens = Math.min( Math.floor(0.8 * reserveTokens), @@ -790,7 +821,7 @@ export async function compact( let historyText = "No prior history."; let historyUsage: Usage | undefined; if (messagesToSummarize.length > 0) { - const historyResult = await generateSummary( + const historyResult = await generateSummaryWithUsage( messagesToSummarize, model, settings.reserveTokens, @@ -822,7 +853,7 @@ export async function compact( summaryUsage = historyUsage ? combineUsage(historyUsage, turnPrefixResult.usage) : turnPrefixResult.usage; } else { // Just generate history summary - const result = await generateSummary( + const result = await generateSummaryWithUsage( messagesToSummarize, model, settings.reserveTokens, diff --git a/packages/coding-agent/src/core/usage-totals.ts b/packages/coding-agent/src/core/usage-totals.ts index e65e2eb49..4ae55d0d2 100644 --- a/packages/coding-agent/src/core/usage-totals.ts +++ b/packages/coding-agent/src/core/usage-totals.ts @@ -1,4 +1,5 @@ import type { Usage } from "@earendil-works/pi-ai/compat"; +import type { SessionEntry } from "./session-manager.ts"; export interface UsageTotals { input: number; @@ -25,3 +26,45 @@ export function addUsageToTotals(totals: UsageTotals, usage: Usage): void { totals.cacheWrite += usage.cacheWrite; totals.cost += usage.cost.total; } + +export interface UsageCostBreakdownEntry { + key: string; + cost: number; + tokens: number; +} + +/** Group attributable assistant usage by model and all other usage into a separate bucket. */ +export function getUsageCostBreakdown(entries: SessionEntry[]): UsageCostBreakdownEntry[] { + const totalsByKey = new Map(); + + for (const entry of entries) { + let key: string | undefined; + let usage: Usage | undefined; + if (entry.type === "message" && entry.message.role === "assistant") { + key = `${entry.message.provider}/${entry.message.responseModel ?? entry.message.model}`; + usage = entry.message.usage; + } else if (entry.type === "message" && entry.message.role === "toolResult" && entry.message.usage) { + key = "Tools/summaries"; + usage = entry.message.usage; + } else if ((entry.type === "branch_summary" || entry.type === "compaction") && entry.usage) { + key = "Tools/summaries"; + usage = entry.usage; + } + if (!key || !usage) continue; + + let totals = totalsByKey.get(key); + if (!totals) { + totals = createUsageTotals(); + totalsByKey.set(key, totals); + } + addUsageToTotals(totals, usage); + } + + return Array.from(totalsByKey, ([key, totals]) => ({ + key, + cost: totals.cost, + tokens: totals.input + totals.output + totals.cacheRead + totals.cacheWrite, + })) + .filter((entry) => entry.cost > 0 || entry.tokens > 0) + .sort((a, b) => b.cost - a.cost); +} diff --git a/packages/coding-agent/src/index.ts b/packages/coding-agent/src/index.ts index 196fc8ca5..e223d3144 100644 --- a/packages/coding-agent/src/index.ts +++ b/packages/coding-agent/src/index.ts @@ -42,6 +42,7 @@ export { type GenerateBranchSummaryOptions, generateBranchSummary, generateSummary, + generateSummaryWithUsage, getLastAssistantUsage, prepareBranchEntries, serializeConversation, diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index b7dc08fc1..54415407c 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -86,6 +86,7 @@ import type { SourceInfo } from "../../core/source-info.ts"; import { isInstallTelemetryEnabled } from "../../core/telemetry.ts"; import type { TruncationResult } from "../../core/tools/truncate.ts"; import { hasTrustRequiringProjectResources, ProjectTrustStore } from "../../core/trust-manager.ts"; +import { getUsageCostBreakdown } from "../../core/usage-totals.ts"; import { getChangelogPath, getNewEntries, normalizeChangelogLinks, parseChangelog } from "../../utils/changelog.ts"; import { copyToClipboard, readClipboardText } from "../../utils/clipboard.ts"; import { extensionForImageMimeType, readClipboardImage } from "../../utils/clipboard-image.ts"; @@ -5597,22 +5598,9 @@ export class InteractiveMode { const cacheWaste = computeCacheWaste(entries, this.session.modelRuntime); // Cost/token totals per provider/model actually used (e.g. OpenRouter `auto` - // resolves to a concrete responseModel), sorted by cost descending. - const perModelMap = new Map(); - for (const entry of entries) { - if (entry.type !== "message" || entry.message.role !== "assistant") continue; - const message = entry.message; - const usage = message.usage; - const key = `${message.provider}/${message.responseModel ?? message.model}`; - let bucket = perModelMap.get(key); - if (!bucket) { - bucket = { key, cost: 0, tokens: 0 }; - perModelMap.set(key, bucket); - } - bucket.cost += usage.cost.total; - bucket.tokens += usage.input + usage.output + usage.cacheRead + usage.cacheWrite; - } - const perModel = Array.from(perModelMap.values()).sort((a, b) => b.cost - a.cost); + // resolves to a concrete responseModel). Usage without model attribution is + // grouped separately so the breakdown reconciles with the session total. + const usageBreakdown = getUsageCostBreakdown(entries); let info = `${theme.bold("Session Info")}\n\n`; if (sessionName) { @@ -5646,8 +5634,8 @@ export class InteractiveMode { if (stats.cost > 0 || cacheWaste.missedTokens > 0) { info += `\n${theme.bold("Cost")}\n`; info += `${theme.fg("dim", "Total:")} $${stats.cost.toFixed(3)}`; - if (perModel.length > 1) { - for (const entry of perModel) { + if (usageBreakdown.length > 1) { + for (const entry of usageBreakdown) { info += `\n ${theme.fg("dim", `${entry.key}:`)} $${entry.cost.toFixed(3)} ${theme.fg("dim", `(${formatTokens(entry.tokens)} tokens)`)}`; } } diff --git a/packages/coding-agent/test/agent-session-stats.test.ts b/packages/coding-agent/test/agent-session-stats.test.ts index 684c5e1cc..2e9d6bf47 100644 --- a/packages/coding-agent/test/agent-session-stats.test.ts +++ b/packages/coding-agent/test/agent-session-stats.test.ts @@ -5,6 +5,7 @@ import { AgentSession } from "../src/core/agent-session.ts"; import { AuthStorage } from "../src/core/auth-storage.ts"; import { SessionManager } from "../src/core/session-manager.ts"; import { SettingsManager } from "../src/core/settings-manager.ts"; +import { getUsageCostBreakdown } from "../src/core/usage-totals.ts"; import { createInMemoryModelRegistry, getModelRuntime } from "./model-runtime-test-utils.ts"; import { createTestResourceLoader } from "./utilities.ts"; @@ -224,6 +225,31 @@ describe("AgentSession.getSessionStats", () => { } }); + it("groups tool and summary usage separately from model-attributed usage", () => { + const sessionManager = SessionManager.inMemory(); + const rootId = sessionManager.appendMessage(createUserMessage("hello", 1)); + sessionManager.appendMessage({ + ...createAssistantMessage("response", 100, 2), + usage: { ...createUsage(100), cost: { ...createUsage(100).cost, total: 0.5 } }, + }); + sessionManager.appendMessage( + createToolResultMessage({ ...createUsage(100), cost: { ...createUsage(100).cost, total: 1 } }), + ); + sessionManager.appendCompaction("summary", rootId, 100, undefined, false, { + ...createUsage(100), + cost: { ...createUsage(100).cost, total: 2 }, + }); + sessionManager.branchWithSummary(null, "branch summary", undefined, false, { + ...createUsage(100), + cost: { ...createUsage(100).cost, total: 3 }, + }); + + expect(getUsageCostBreakdown(sessionManager.getEntries())).toEqual([ + { key: "Tools/summaries", cost: 6, tokens: 300 }, + { key: `${model.provider}/${model.id}`, cost: 0.5, tokens: 100 }, + ]); + }); + it("ignores zero-usage messages when checking for post-compaction context usage", async () => { const { session, sessionManager } = await createSession(); diff --git a/packages/coding-agent/test/compaction-summary-reasoning.test.ts b/packages/coding-agent/test/compaction-summary-reasoning.test.ts index eb1d35794..417662385 100644 --- a/packages/coding-agent/test/compaction-summary-reasoning.test.ts +++ b/packages/coding-agent/test/compaction-summary-reasoning.test.ts @@ -1,7 +1,12 @@ import type { AgentMessage } from "@earendil-works/pi-agent-core"; import type { AssistantMessage, Model } from "@earendil-works/pi-ai"; import { beforeEach, describe, expect, it, vi } from "vitest"; -import { type CompactionPreparation, compact, generateSummary } from "../src/core/compaction/index.ts"; +import { + type CompactionPreparation, + compact, + generateSummary, + generateSummaryWithUsage, +} from "../src/core/compaction/index.ts"; const { completeSimpleMock } = vi.hoisted(() => ({ completeSimpleMock: vi.fn(), @@ -57,7 +62,7 @@ describe("generateSummary reasoning options", () => { }); it("uses the provided thinking level for reasoning-capable models", async () => { - const result = await generateSummary( + const result = await generateSummaryWithUsage( messages, createModel(true), 2000, @@ -79,6 +84,12 @@ describe("generateSummary reasoning options", () => { }); }); + it("preserves the string result from generateSummary", async () => { + await expect(generateSummary(messages, createModel(false), 2000, "test-key")).resolves.toBe( + "## Goal\nTest summary", + ); + }); + it("does not set reasoning when thinking is off", async () => { await generateSummary( messages, diff --git a/packages/coding-agent/test/session-manager/tree-traversal.test.ts b/packages/coding-agent/test/session-manager/tree-traversal.test.ts index 9f447de1b..d6e20f043 100644 --- a/packages/coding-agent/test/session-manager/tree-traversal.test.ts +++ b/packages/coding-agent/test/session-manager/tree-traversal.test.ts @@ -523,6 +523,52 @@ describe("createBranchedSession", () => { } }); + it("preserves tool and summary usage across a file-backed reload", () => { + const tempDir = join(tmpdir(), `session-usage-roundtrip-${Date.now()}`); + mkdirSync(tempDir, { recursive: true }); + + try { + const session = SessionManager.create(tempDir, tempDir); + const rootId = session.appendMessage(userMsg("question")); + session.appendMessage(assistantMsg("answer")); + const usage = { + input: 10, + output: 20, + cacheRead: 30, + cacheWrite: 40, + totalTokens: 100, + cost: { input: 0.1, output: 0.2, cacheRead: 0.3, cacheWrite: 0.4, total: 1 }, + }; + session.appendMessage({ + role: "toolResult", + toolCallId: "call-1", + toolName: "nested-model", + content: [{ type: "text", text: "result" }], + isError: false, + usage, + timestamp: Date.now(), + }); + session.appendCompaction("summary", rootId, 100, undefined, false, usage); + session.branchWithSummary(rootId, "branch summary", undefined, false, usage); + + const file = session.getSessionFile(); + expect(file).toBeDefined(); + const reopened = SessionManager.open(file!, tempDir); + expect(reopened.getEntries()).toEqual( + expect.arrayContaining([ + expect.objectContaining({ type: "compaction", usage }), + expect.objectContaining({ type: "branch_summary", usage }), + expect.objectContaining({ + type: "message", + message: expect.objectContaining({ role: "toolResult", usage }), + }), + ]), + ); + } finally { + rmSync(tempDir, { recursive: true, force: true }); + } + }); + it("writes file immediately when forking from a point with assistant messages", () => { const tempDir = join(tmpdir(), `session-fork-with-assistant-${Date.now()}`); mkdirSync(tempDir, { recursive: true }); From c889eb8809a0f40ccd937dc915b10147bec39115 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Mon, 20 Jul 2026 17:13:06 +0200 Subject: [PATCH 25/62] fix(coding-agent): defer startup model catalog refresh --- packages/coding-agent/CHANGELOG.md | 1 + .../coding-agent/src/core/model-runtime.ts | 35 +++++++++++-------- packages/coding-agent/src/main.ts | 6 +++- packages/coding-agent/test/radius.test.ts | 14 ++++++++ 4 files changed, 41 insertions(+), 15 deletions(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index cac4d0193..f12db68e2 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -10,6 +10,7 @@ ### Fixed +- Moved automatic model catalog network refresh out of startup initialization and into the running interactive and RPC modes. - Fixed persisted sessions being read and parsed twice when opened, reducing startup latency for large sessions ([#6793](https://github.com/earendil-works/pi/issues/6793)). - Fixed prompt-template defaults for all arguments (`${@:-default}` and `${ARGUMENTS:-default}`) ([#6695](https://github.com/earendil-works/pi/issues/6695)). - Fixed obsolete custom UI, custom tool, and custom editor examples in the extension documentation ([#6735](https://github.com/earendil-works/pi/issues/6735)). diff --git a/packages/coding-agent/src/core/model-runtime.ts b/packages/coding-agent/src/core/model-runtime.ts index 2235d4006..64501f731 100644 --- a/packages/coding-agent/src/core/model-runtime.ts +++ b/packages/coding-agent/src/core/model-runtime.ts @@ -62,7 +62,9 @@ export interface CreateModelRuntimeOptions { modelsPath?: string | null; modelsStore?: ModelsStore; modelsStorePath?: string; + /** Allow create() to refresh model catalogs over the network. Defaults to false. */ allowModelNetwork?: boolean; + /** Timeout for the create-time network model refresh. */ modelRefreshTimeoutMs?: number; catalogBaseUrl?: string; } @@ -98,7 +100,7 @@ export class ModelRuntime implements Models { private readonly extensionProviders = new Map(); private readonly compositionErrors = new Map(); private readonly modelsPath: string | undefined; - private readonly allowModelNetwork: boolean; + private readonly modelNetworkEnabled: boolean; private config: ModelConfig; private snapshot: ModelRuntimeSnapshot = { all: [], @@ -116,12 +118,12 @@ export class ModelRuntime implements Models { modelsPath: string | undefined, modelsStore: ModelsStore, providers: readonly Provider[], - allowModelNetwork: boolean, + modelNetworkEnabled: boolean, ) { this.credentials = credentials; this.config = config; this.modelsPath = modelsPath; - this.allowModelNetwork = allowModelNetwork; + this.modelNetworkEnabled = modelNetworkEnabled; this.defaultBuiltins = new Map(providers.map((provider) => [provider.id, provider])); for (const [providerId, provider] of this.defaultBuiltins) this.builtins.set(providerId, provider); this.models = createModels({ credentials, modelsStore }); @@ -149,16 +151,17 @@ export class ModelRuntime implements Models { modelsPath, modelsStore, providers, - options.allowModelNetwork ?? process.env.PI_OFFLINE === undefined, + process.env.PI_OFFLINE === undefined, ); runtime.configureRadiusProviders(); runtime.rebuildProviders(); - const controller = new AbortController(); - const timeout = runtime.allowModelNetwork + const refreshFromNetwork = runtime.modelNetworkEnabled && options.allowModelNetwork === true; + const controller = refreshFromNetwork ? new AbortController() : undefined; + const timeout = controller ? setTimeout(() => controller.abort(), options.modelRefreshTimeoutMs ?? 15_000) : undefined; try { - await runtime.refresh({ allowNetwork: runtime.allowModelNetwork, signal: controller.signal }); + await runtime.refresh({ allowNetwork: refreshFromNetwork, signal: controller?.signal }); } finally { if (timeout) clearTimeout(timeout); } @@ -389,7 +392,11 @@ export class ModelRuntime implements Models { }; } - async setRuntimeApiKey(providerId: string, apiKey: string): Promise { + async setRuntimeApiKey( + providerId: string, + apiKey: string, + refreshOptions: ModelsRefreshOptions = {}, + ): Promise { this.credentials.setRuntimeApiKey(providerId, apiKey); const auth = new Map(this.snapshot.auth).set(providerId, { type: "api_key", source: "runtime API key" }); const configuredProviders = new Set(this.snapshot.configuredProviders).add(providerId); @@ -401,12 +408,12 @@ export class ModelRuntime implements Models { storedProviders, available: this.snapshot.all.filter((model) => configuredProviders.has(model.provider)), }; - await this.refresh({ allowNetwork: this.allowModelNetwork }); + await this.refresh(refreshOptions); } async removeRuntimeApiKey(providerId: string): Promise { this.credentials.removeRuntimeApiKey(providerId); - await this.refresh({ allowNetwork: this.allowModelNetwork }); + await this.refresh({ allowNetwork: this.modelNetworkEnabled }); } listCredentials(): Promise { @@ -492,7 +499,7 @@ export class ModelRuntime implements Models { async login(providerId: string, type: AuthType, interaction: AuthInteraction): Promise { const credential = await this.models.login(providerId, type, interaction); - await this.refresh({ allowNetwork: this.allowModelNetwork }); + await this.refresh({ allowNetwork: this.modelNetworkEnabled }); return credential; } @@ -500,20 +507,20 @@ export class ModelRuntime implements Models { await this.models.logout(providerId); // Reset credential-dependent compatibility projections before the unconfigured provider is skipped by refresh. this.recomposeProvider(providerId); - await this.refresh({ allowNetwork: this.allowModelNetwork }); + await this.refresh({ allowNetwork: this.modelNetworkEnabled }); } async reloadConfig(): Promise { this.config = await ModelConfig.load(this.modelsPath); this.configureRadiusProviders(); this.rebuildProviders(); - await this.refresh({ allowNetwork: this.allowModelNetwork }); + await this.refresh({ allowNetwork: this.modelNetworkEnabled }); } async refresh(options: ModelsRefreshOptions = {}): Promise { const refreshOptions = { ...options, - allowNetwork: options.allowNetwork ?? this.allowModelNetwork, + allowNetwork: options.allowNetwork ?? this.modelNetworkEnabled, }; // Published pi-ai builds before ModelsStore returned void and accepted a provider ID. // The fallback keeps source-mode CLI tests working without rebuilding workspace dependencies. diff --git a/packages/coding-agent/src/main.ts b/packages/coding-agent/src/main.ts index 7d9c55253..e004436f2 100644 --- a/packages/coding-agent/src/main.ts +++ b/packages/coding-agent/src/main.ts @@ -709,7 +709,7 @@ export async function main(args: string[], options?: MainOptions) { message: "--api-key requires a model to be specified via --model, --provider/--model, or --models", }); } else { - await modelRuntime.setRuntimeApiKey(sessionOptions.model.provider, parsed.apiKey); + await modelRuntime.setRuntimeApiKey(sessionOptions.model.provider, parsed.apiKey, { allowNetwork: false }); await services.modelRuntime.getAvailable(); } } @@ -808,6 +808,10 @@ export async function main(args: string[], options?: MainOptions) { process.exit(1); } + if (!offlineMode && (appMode === "interactive" || appMode === "rpc")) { + void modelRuntime.refresh().catch(() => {}); + } + if (appMode === "rpc") { printTimings(); await runRpcMode(runtime); diff --git a/packages/coding-agent/test/radius.test.ts b/packages/coding-agent/test/radius.test.ts index de3836727..8a7f5965a 100644 --- a/packages/coding-agent/test/radius.test.ts +++ b/packages/coding-agent/test/radius.test.ts @@ -91,6 +91,20 @@ describe("Radius provider", () => { expect(vi.mocked(fetch).mock.calls[0]?.[1]?.headers).toMatchObject({ authorization: "Bearer access-token" }); }); + it("does not refresh catalogs over the network by default", async () => { + const fetchSpy = vi.spyOn(globalThis, "fetch"); + const runtime = await ModelRuntime.create({ + credentials: AuthStorage.inMemory({ + [RADIUS_PROVIDER_ID]: radiusOAuthCredential("https://radius.example.com/v1"), + }), + modelsStore: new InMemoryModelsStore(), + modelsPath: null, + }); + + expect(runtime.getModel(RADIUS_PROVIDER_ID, "auto")).toBeDefined(); + expect(fetchSpy).not.toHaveBeenCalled(); + }); + it("does not fetch or expose Radius models without configured auth", async () => { const fetchSpy = vi.spyOn(globalThis, "fetch"); const runtime = await ModelRuntime.create({ From 3a40794ea14c6202586cc203d5b928eca9f6b673 Mon Sep 17 00:00:00 2001 From: Armin Ronacher Date: Mon, 20 Jul 2026 17:19:40 +0200 Subject: [PATCH 26/62] feat(ai): regenerate models --- packages/ai/src/providers/openrouter.models.ts | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/packages/ai/src/providers/openrouter.models.ts b/packages/ai/src/providers/openrouter.models.ts index ad687c266..694e296a0 100644 --- a/packages/ai/src/providers/openrouter.models.ts +++ b/packages/ai/src/providers/openrouter.models.ts @@ -285,6 +285,10 @@ export const OPENROUTER_MODELS = values as { id: "kwaipilot/kat-coder-pro-v2.5"; provider: "openrouter"; }; + "meituan/longcat-2.0": Model<"openai-completions"> & { + id: "meituan/longcat-2.0"; + provider: "openrouter"; + }; "meta-llama/llama-3.1-70b-instruct": Model<"openai-completions"> & { id: "meta-llama/llama-3.1-70b-instruct"; provider: "openrouter"; From 1235c0ec649b5fbd2c7cd2328aa8d7786f565d8c Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Mon, 20 Jul 2026 17:54:17 +0200 Subject: [PATCH 27/62] fix(agent): decouple agent streams from compat closes #6851 --- packages/agent/CHANGELOG.md | 1 + packages/agent/README.md | 41 ++++++++----- packages/agent/src/agent-loop.ts | 35 +++++------ packages/agent/src/agent.ts | 31 +++++----- packages/agent/src/proxy.ts | 4 +- packages/agent/test/agent-loop.test.ts | 6 +- packages/agent/test/agent.test.ts | 45 +++++++------- packages/agent/test/e2e.test.ts | 11 ++++ .../coding-agent/src/core/agent-session.ts | 10 ++-- packages/coding-agent/src/core/sdk.ts | 2 +- ...gent-session-auto-compaction-queue.test.ts | 9 ++- .../test/agent-session-compaction.test.ts | 3 +- .../test/agent-session-concurrent.test.ts | 10 ++-- .../agent-session-dynamic-provider.test.ts | 2 +- .../test/agent-session-retry.test.ts | 6 +- .../test/agent-session-stats.test.ts | 9 ++- .../test/compaction-extensions.test.ts | 3 +- .../rpc-prompt-response-semantics.test.ts | 2 +- .../test/sdk-openrouter-attribution.test.ts | 2 +- .../test/sdk-stream-options.test.ts | 2 +- .../suite/agent-session-compaction.test.ts | 2 +- packages/coding-agent/test/suite/harness.ts | 3 +- .../5596-missing-theme-export.test.ts | 3 +- .../6324-branch-summary-ambient-auth.test.ts | 2 +- packages/coding-agent/test/test-harness.ts | 2 +- packages/coding-agent/test/utilities.ts | 3 +- scripts/agent-treeshake-smoke-entry.ts | 13 +++++ scripts/browser-smoke-entry.ts | 4 +- scripts/check-browser-smoke.mjs | 58 +++++++++++++++++++ 29 files changed, 219 insertions(+), 105 deletions(-) create mode 100644 scripts/agent-treeshake-smoke-entry.ts diff --git a/packages/agent/CHANGELOG.md b/packages/agent/CHANGELOG.md index 3e8ec724a..57bd9ad85 100644 --- a/packages/agent/CHANGELOG.md +++ b/packages/agent/CHANGELOG.md @@ -5,6 +5,7 @@ ### Breaking Changes - Moved the `uuidv7` export to `@earendil-works/pi-ai`. +- Replaced the optional `Agent` `streamFn` fallback with a required `streamFunction` and made low-level loop stream functions required, preventing `@earendil-works/pi-ai/compat` and all built-in providers from entering selective-provider bundles ([#6851](https://github.com/earendil-works/pi/issues/6851)). ### Added diff --git a/packages/agent/README.md b/packages/agent/README.md index 17eecea57..3c935f444 100644 --- a/packages/agent/README.md +++ b/packages/agent/README.md @@ -12,13 +12,20 @@ npm install @earendil-works/pi-agent-core ```typescript import { Agent } from "@earendil-works/pi-agent-core"; -import { getModel } from "@earendil-works/pi-ai"; +import { createModels } from "@earendil-works/pi-ai"; +import { anthropicProvider } from "@earendil-works/pi-ai/providers/anthropic"; + +const models = createModels(); +models.setProvider(anthropicProvider()); +const model = models.getModel("anthropic", "claude-sonnet-4-6"); +if (!model) throw new Error("Model not found"); const agent = new Agent({ initialState: { systemPrompt: "You are a helpful assistant.", - model: getModel("anthropic", "claude-sonnet-4-20250514"), + model, }, + streamFunction: models.streamSimple.bind(models), }); agent.subscribe((event) => { @@ -115,13 +122,19 @@ Tools can also return `terminate: true` to hint that the automatic follow-up LLM Low-level loop callers can set `shouldStopAfterTurn` to stop gracefully after the current turn completes: ```typescript -const stream = agentLoop(prompts, context, { - model, - convertToLlm, - shouldStopAfterTurn: async ({ message, toolResults, context, newMessages }) => { - return shouldCompactBeforeNextTurn(context.messages); +const stream = agentLoop( + prompts, + context, + { + model, + convertToLlm, + shouldStopAfterTurn: async ({ message, toolResults, context, newMessages }) => { + return shouldCompactBeforeNextTurn(context.messages); + }, }, -}); + undefined, + models.streamSimple.bind(models), +); ``` `shouldStopAfterTurn` runs after `turn_end` is emitted and after the assistant response and any tool executions have completed normally. If it returns `true`, the loop emits `agent_end` and exits before polling steering or follow-up queues, and before starting another LLM call. It does not abort the provider stream, does not cancel running tools, and does not alter the assistant message stop reason. @@ -181,8 +194,8 @@ const agent = new Agent({ // Follow-up mode: "one-at-a-time" (default) or "all" followUpMode: "one-at-a-time", - // Custom stream function (for proxy backends) - streamFn: streamProxy, + // Required stream function + streamFunction: models.streamSimple.bind(models), // Session ID for provider caching sessionId: "session-123", @@ -369,6 +382,7 @@ Handle custom types in `convertToLlm`: ```typescript const agent = new Agent({ + streamFunction: models.streamSimple.bind(models), convertToLlm: (messages) => messages.flatMap(m => { if (m.role === "notification") return []; // Filter out return [m]; @@ -439,7 +453,7 @@ For browser apps that proxy through a backend: import { Agent, streamProxy } from "@earendil-works/pi-agent-core"; const agent = new Agent({ - streamFn: (model, context, options) => + streamFunction: (model, context, options) => streamProxy(model, context, { ...options, authToken: "...", @@ -471,12 +485,13 @@ const config: AgentLoopConfig = { const userMessage = { role: "user", content: "Hello", timestamp: Date.now() }; -for await (const event of agentLoop([userMessage], context, config)) { +const streamFunction = models.streamSimple.bind(models); +for await (const event of agentLoop([userMessage], context, config, undefined, streamFunction)) { console.log(event.type); } // Continue from existing context -for await (const event of agentLoopContinue(context, config)) { +for await (const event of agentLoopContinue(context, config, undefined, streamFunction)) { console.log(event.type); } ``` diff --git a/packages/agent/src/agent-loop.ts b/packages/agent/src/agent-loop.ts index c9e731cd1..fff069969 100644 --- a/packages/agent/src/agent-loop.ts +++ b/packages/agent/src/agent-loop.ts @@ -7,10 +7,9 @@ import { type AssistantMessage, type Context, EventStream, - streamSimple, type ToolResultMessage, validateToolArguments, -} from "@earendil-works/pi-ai/compat"; +} from "@earendil-works/pi-ai"; import type { AgentContext, AgentEvent, @@ -32,8 +31,8 @@ export function agentLoop( prompts: AgentMessage[], context: AgentContext, config: AgentLoopConfig, - signal?: AbortSignal, - streamFn?: StreamFn, + signal: AbortSignal | undefined, + streamFunction: StreamFn, ): EventStream { const stream = createAgentStream(); @@ -45,7 +44,7 @@ export function agentLoop( stream.push(event); }, signal, - streamFn, + streamFunction, ).then((messages) => { stream.end(messages); }); @@ -64,8 +63,8 @@ export function agentLoop( export function agentLoopContinue( context: AgentContext, config: AgentLoopConfig, - signal?: AbortSignal, - streamFn?: StreamFn, + signal: AbortSignal | undefined, + streamFunction: StreamFn, ): EventStream { if (context.messages.length === 0) { throw new Error("Cannot continue: no messages in context"); @@ -84,7 +83,7 @@ export function agentLoopContinue( stream.push(event); }, signal, - streamFn, + streamFunction, ).then((messages) => { stream.end(messages); }); @@ -97,8 +96,8 @@ export async function runAgentLoop( context: AgentContext, config: AgentLoopConfig, emit: AgentEventSink, - signal?: AbortSignal, - streamFn?: StreamFn, + signal: AbortSignal | undefined, + streamFunction: StreamFn, ): Promise { const newMessages: AgentMessage[] = [...prompts]; const currentContext: AgentContext = { @@ -113,7 +112,7 @@ export async function runAgentLoop( await emit({ type: "message_end", message: prompt }); } - await runLoop(currentContext, newMessages, config, signal, emit, streamFn); + await runLoop(currentContext, newMessages, config, signal, emit, streamFunction); return newMessages; } @@ -121,8 +120,8 @@ export async function runAgentLoopContinue( context: AgentContext, config: AgentLoopConfig, emit: AgentEventSink, - signal?: AbortSignal, - streamFn?: StreamFn, + signal: AbortSignal | undefined, + streamFunction: StreamFn, ): Promise { if (context.messages.length === 0) { throw new Error("Cannot continue: no messages in context"); @@ -138,7 +137,7 @@ export async function runAgentLoopContinue( await emit({ type: "agent_start" }); await emit({ type: "turn_start" }); - await runLoop(currentContext, newMessages, config, signal, emit, streamFn); + await runLoop(currentContext, newMessages, config, signal, emit, streamFunction); return newMessages; } @@ -158,7 +157,7 @@ async function runLoop( initialConfig: AgentLoopConfig, signal: AbortSignal | undefined, emit: AgentEventSink, - streamFn?: StreamFn, + streamFunction: StreamFn, ): Promise { let currentContext = initialContext; let config = initialConfig; @@ -190,7 +189,7 @@ async function runLoop( } // Stream assistant response - const message = await streamAssistantResponse(currentContext, config, signal, emit, streamFn); + const message = await streamAssistantResponse(currentContext, config, signal, emit, streamFunction); newMessages.push(message); if (message.stopReason === "error" || message.stopReason === "aborted") { @@ -283,7 +282,7 @@ async function streamAssistantResponse( config: AgentLoopConfig, signal: AbortSignal | undefined, emit: AgentEventSink, - streamFn?: StreamFn, + streamFunction: StreamFn, ): Promise { // Apply context transform if configured (AgentMessage[] → AgentMessage[]) let messages = context.messages; @@ -301,8 +300,6 @@ async function streamAssistantResponse( 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; diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index ab10d5ac1..653adf8e8 100644 --- a/packages/agent/src/agent.ts +++ b/packages/agent/src/agent.ts @@ -1,13 +1,12 @@ -import { - type ImageContent, - type Message, - type Model, - type SimpleStreamOptions, - streamSimple, - type TextContent, - type ThinkingBudgets, - type Transport, -} from "@earendil-works/pi-ai/compat"; +import type { + ImageContent, + Message, + Model, + SimpleStreamOptions, + TextContent, + ThinkingBudgets, + Transport, +} from "@earendil-works/pi-ai"; import { runAgentLoop, runAgentLoopContinue } from "./agent-loop.ts"; import type { AfterToolCallContext, @@ -98,7 +97,7 @@ export interface AgentOptions { initialState?: Partial>; convertToLlm?: (messages: AgentMessage[]) => Message[] | Promise; transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise; - streamFn?: StreamFn; + streamFunction: StreamFn; getApiKey?: (provider: string) => Promise | string | undefined; onPayload?: SimpleStreamOptions["onPayload"]; onResponse?: SimpleStreamOptions["onResponse"]; @@ -176,7 +175,7 @@ export class Agent { public convertToLlm: (messages: AgentMessage[]) => Message[] | Promise; public transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise; - public streamFn: StreamFn; + public streamFunction: StreamFn; public getApiKey?: (provider: string) => Promise | string | undefined; public onPayload?: SimpleStreamOptions["onPayload"]; public onResponse?: SimpleStreamOptions["onResponse"]; @@ -207,11 +206,11 @@ export class Agent { /** Tool execution strategy for assistant messages that contain multiple tool calls. */ public toolExecution: ToolExecutionMode; - constructor(options: AgentOptions = {}) { + constructor(options: AgentOptions) { this._state = createMutableAgentState(options.initialState); this.convertToLlm = options.convertToLlm ?? defaultConvertToLlm; this.transformContext = options.transformContext; - this.streamFn = options.streamFn ?? streamSimple; + this.streamFunction = options.streamFunction; this.getApiKey = options.getApiKey; this.onPayload = options.onPayload; this.onResponse = options.onResponse; @@ -404,7 +403,7 @@ export class Agent { this.createLoopConfig(options), (event) => this.processEvents(event), signal, - this.streamFn, + this.streamFunction, ); }); } @@ -416,7 +415,7 @@ export class Agent { this.createLoopConfig(), (event) => this.processEvents(event), signal, - this.streamFn, + this.streamFunction, ); }); } diff --git a/packages/agent/src/proxy.ts b/packages/agent/src/proxy.ts index 5f0925c91..5d4b7a55e 100644 --- a/packages/agent/src/proxy.ts +++ b/packages/agent/src/proxy.ts @@ -84,12 +84,12 @@ export interface ProxyStreamOptions extends ProxySerializableStreamOptions { * 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. + * Use this as the `streamFunction` option when creating an Agent that needs to go through a proxy. * * @example * ```typescript * const agent = new Agent({ - * streamFn: (model, context, options) => + * streamFunction: (model, context, options) => * streamProxy(model, context, { * ...options, * authToken: await getAuthToken(), diff --git a/packages/agent/test/agent-loop.test.ts b/packages/agent/test/agent-loop.test.ts index f9641ce74..47380c42c 100644 --- a/packages/agent/test/agent-loop.test.ts +++ b/packages/agent/test/agent-loop.test.ts @@ -1342,7 +1342,11 @@ describe("agentLoopContinue with AgentMessage", () => { convertToLlm: identityConverter, }; - expect(() => agentLoopContinue(context, config)).toThrow("Cannot continue: no messages in context"); + expect(() => + agentLoopContinue(context, config, undefined, () => { + throw new Error("Unexpected stream call"); + }), + ).toThrow("Cannot continue: no messages in context"); }); it("should continue from existing context without emitting user message events", async () => { diff --git a/packages/agent/test/agent.test.ts b/packages/agent/test/agent.test.ts index f677ee788..6cd77f11e 100644 --- a/packages/agent/test/agent.test.ts +++ b/packages/agent/test/agent.test.ts @@ -1,7 +1,7 @@ import { type AssistantMessage, type AssistantMessageEvent, EventStream, getModel } from "@earendil-works/pi-ai/compat"; import { Type } from "typebox"; import { describe, expect, it } from "vitest"; -import { Agent, type AgentEvent, type AgentTool, type AgentToolUpdateCallback } from "../src/index.ts"; +import { Agent, type AgentEvent, type AgentTool, type AgentToolUpdateCallback, type StreamFn } from "../src/index.ts"; // Mock stream that mimics AssistantMessageEventStream class MockAssistantStream extends EventStream { @@ -59,6 +59,10 @@ function createAssistantToolUseMessage(content: ToolCallContent[]): AssistantMes }; } +const unusedStreamFunction: StreamFn = () => { + throw new Error("Unexpected stream call"); +}; + function createDeferred(): { promise: Promise; resolve: () => void; @@ -72,7 +76,7 @@ function createDeferred(): { describe("Agent", () => { it("should create an agent instance with default state", () => { - const agent = new Agent(); + const agent = new Agent({ streamFunction: unusedStreamFunction }); expect(agent.state).toBeDefined(); expect(agent.state.systemPrompt).toBe(""); @@ -89,6 +93,7 @@ describe("Agent", () => { it("should create an agent instance with custom initial state", () => { const customModel = getModel("openai", "gpt-4o-mini"); const agent = new Agent({ + streamFunction: unusedStreamFunction, initialState: { systemPrompt: "You are a helpful assistant.", model: customModel, @@ -102,7 +107,7 @@ describe("Agent", () => { }); it("should subscribe to events", () => { - const agent = new Agent(); + const agent = new Agent({ streamFunction: unusedStreamFunction }); let eventCount = 0; const unsubscribe = agent.subscribe((_event) => { @@ -125,7 +130,7 @@ describe("Agent", () => { it("emits full lifecycle events for thrown run failures", async () => { const agent = new Agent({ - streamFn: () => { + streamFunction: () => { throw new Error("provider exploded"); }, }); @@ -157,7 +162,7 @@ describe("Agent", () => { it("should await async subscribers before prompt resolves", async () => { const barrier = createDeferred(); const agent = new Agent({ - streamFn: () => { + streamFunction: () => { const stream = new MockAssistantStream(); queueMicrotask(() => { stream.push({ type: "done", reason: "stop", message: createAssistantMessage("ok") }); @@ -195,7 +200,7 @@ describe("Agent", () => { it("waitForIdle should wait for async subscribers", async () => { const barrier = createDeferred(); const agent = new Agent({ - streamFn: () => { + streamFunction: () => { const stream = new MockAssistantStream(); queueMicrotask(() => { stream.push({ type: "done", reason: "stop", message: createAssistantMessage("ok") }); @@ -230,7 +235,7 @@ describe("Agent", () => { it("should pass the active abort signal to subscribers", async () => { let receivedSignal: AbortSignal | undefined; const agent = new Agent({ - streamFn: (_model, _context, options) => { + streamFunction: (_model, _context, options) => { const stream = new MockAssistantStream(); queueMicrotask(() => { stream.push({ type: "start", partial: createAssistantMessage("") }); @@ -293,7 +298,7 @@ describe("Agent", () => { }; const agent = new Agent({ initialState: { tools: [tool] }, - streamFn: () => { + streamFunction: () => { const stream = new MockAssistantStream(); queueMicrotask(() => { stream.push({ @@ -368,7 +373,7 @@ describe("Agent", () => { }; const agent = new Agent({ initialState: { tools: [settledTool, slowTool] }, - streamFn: () => { + streamFunction: () => { const stream = new MockAssistantStream(); queueMicrotask(() => { stream.push({ @@ -407,7 +412,7 @@ describe("Agent", () => { }); it("should update state with mutators", () => { - const agent = new Agent(); + const agent = new Agent({ streamFunction: unusedStreamFunction }); // Test setSystemPrompt agent.state.systemPrompt = "Custom prompt"; @@ -446,7 +451,7 @@ describe("Agent", () => { }); it("should support steering message queue", async () => { - const agent = new Agent(); + const agent = new Agent({ streamFunction: unusedStreamFunction }); const message = { role: "user" as const, content: "Steering message", timestamp: Date.now() }; agent.steer(message); @@ -456,7 +461,7 @@ describe("Agent", () => { }); it("should support follow-up message queue", async () => { - const agent = new Agent(); + const agent = new Agent({ streamFunction: unusedStreamFunction }); const message = { role: "user" as const, content: "Follow-up message", timestamp: Date.now() }; agent.followUp(message); @@ -466,7 +471,7 @@ describe("Agent", () => { }); it("should handle abort controller", () => { - const agent = new Agent(); + const agent = new Agent({ streamFunction: unusedStreamFunction }); // Should not throw even if nothing is running expect(() => agent.abort()).not.toThrow(); @@ -476,7 +481,7 @@ describe("Agent", () => { let abortSignal: AbortSignal | undefined; const agent = new Agent({ // Use a stream function that responds to abort - streamFn: (_model, _context, options) => { + streamFunction: (_model, _context, options) => { abortSignal = options?.signal; const stream = new MockAssistantStream(); queueMicrotask(() => { @@ -515,7 +520,7 @@ describe("Agent", () => { it("should throw when continue() called while streaming", async () => { let abortSignal: AbortSignal | undefined; const agent = new Agent({ - streamFn: (_model, _context, options) => { + streamFunction: (_model, _context, options) => { abortSignal = options?.signal; const stream = new MockAssistantStream(); queueMicrotask(() => { @@ -550,7 +555,7 @@ describe("Agent", () => { it("continue() should process queued follow-up messages after an assistant turn", async () => { const agent = new Agent({ - streamFn: () => { + streamFunction: () => { const stream = new MockAssistantStream(); queueMicrotask(() => { stream.push({ type: "done", reason: "stop", message: createAssistantMessage("Processed") }); @@ -589,7 +594,7 @@ describe("Agent", () => { it("continue() should keep one-at-a-time steering semantics from assistant tail", async () => { let responseCount = 0; const agent = new Agent({ - streamFn: () => { + streamFunction: () => { const stream = new MockAssistantStream(); responseCount++; queueMicrotask(() => { @@ -647,7 +652,7 @@ describe("Agent", () => { sawAbortSignal = signal instanceof AbortSignal; return undefined; }, - streamFn: () => { + streamFunction: () => { requestCount++; const stream = new MockAssistantStream(); queueMicrotask(() => { @@ -671,11 +676,11 @@ describe("Agent", () => { expect(sawAbortSignal).toBe(true); }); - it("forwards sessionId to streamFn options", async () => { + it("forwards sessionId to streamFunction options", async () => { let receivedSessionId: string | undefined; const agent = new Agent({ sessionId: "session-abc", - streamFn: (_model, _context, options) => { + streamFunction: (_model, _context, options) => { receivedSessionId = options?.sessionId; const stream = new MockAssistantStream(); queueMicrotask(() => { diff --git a/packages/agent/test/e2e.test.ts b/packages/agent/test/e2e.test.ts index 57d70585b..484ef0796 100644 --- a/packages/agent/test/e2e.test.ts +++ b/packages/agent/test/e2e.test.ts @@ -7,6 +7,7 @@ import { fauxToolCall, type Model, registerFauxProvider, + streamSimple, type ToolResultMessage, type UserMessage, } from "@earendil-works/pi-ai/compat"; @@ -37,6 +38,7 @@ afterEach(() => { async function basicPrompt(model: Model) { const agent = new Agent({ + streamFunction: streamSimple, initialState: { systemPrompt: "You are a helpful assistant. Keep your responses concise.", model, @@ -59,6 +61,7 @@ async function basicPrompt(model: Model) { async function toolExecution(model: Model) { const agent = new Agent({ + streamFunction: streamSimple, initialState: { systemPrompt: "You are a helpful assistant. Always use the calculator tool for math.", model, @@ -98,6 +101,7 @@ async function toolExecution(model: Model) { async function abortExecution(model: Model) { const agent = new Agent({ + streamFunction: streamSimple, initialState: { systemPrompt: "You are a helpful assistant.", model, @@ -125,6 +129,7 @@ async function abortExecution(model: Model) { async function stateUpdates(model: Model) { const agent = new Agent({ + streamFunction: streamSimple, initialState: { systemPrompt: "You are a helpful assistant.", model, @@ -157,6 +162,7 @@ async function stateUpdates(model: Model) { async function multiTurnConversation(model: Model) { const agent = new Agent({ + streamFunction: streamSimple, initialState: { systemPrompt: "You are a helpful assistant.", model, @@ -238,6 +244,7 @@ describe("Agent integration with faux provider", () => { faux.setResponses([fauxAssistantMessage([fauxThinking("step by step"), fauxText("4")])]); const agent = new Agent({ + streamFunction: streamSimple, initialState: { systemPrompt: "You are a helpful assistant.", model: faux.getModel(), @@ -262,6 +269,7 @@ describe("Agent.continue() with faux provider", () => { it("throws when no messages in context", async () => { const faux = createFauxRegistration(); const agent = new Agent({ + streamFunction: streamSimple, initialState: { systemPrompt: "Test", model: faux.getModel(), @@ -275,6 +283,7 @@ describe("Agent.continue() with faux provider", () => { const faux = createFauxRegistration(); const model = faux.getModel(); const agent = new Agent({ + streamFunction: streamSimple, initialState: { systemPrompt: "Test", model, @@ -309,6 +318,7 @@ describe("Agent.continue() with faux provider", () => { const faux = createFauxRegistration(); faux.setResponses([fauxAssistantMessage("HELLO WORLD")]); const agent = new Agent({ + streamFunction: streamSimple, initialState: { systemPrompt: "You are a helpful assistant. Follow instructions exactly.", model: faux.getModel(), @@ -343,6 +353,7 @@ describe("Agent.continue() with faux provider", () => { const model = faux.getModel(); faux.setResponses([fauxAssistantMessage("The answer is 8.")]); const agent = new Agent({ + streamFunction: streamSimple, initialState: { systemPrompt: "You are a helpful assistant. After getting a calculation result, state the answer clearly.", diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index eb4f78e90..b0fdbb076 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -426,7 +426,7 @@ export class AgentSession { headers?: Record; env?: Record; }> { - if (this.agent.streamFn === streamSimple) { + if (this.agent.streamFunction === streamSimple) { return this._getRequiredRequestAuth(model); } @@ -1836,7 +1836,7 @@ export class AgentSession { customInstructions, this._compactionAbortController.signal, this.thinkingLevel, - this.agent.streamFn, + this.agent.streamFunction, env, ); summary = result.summary; @@ -2037,7 +2037,7 @@ export class AgentSession { let apiKey: string | undefined; let headers: Record | undefined; let env: Record | undefined; - if (this.agent.streamFn === streamSimple) { + if (this.agent.streamFunction === streamSimple) { const authResult = await this._modelRuntime.getAuth(this.model); if (!authResult?.auth.apiKey) return false; apiKey = authResult.auth.apiKey; @@ -2112,7 +2112,7 @@ export class AgentSession { undefined, this._autoCompactionAbortController.signal, this.thinkingLevel, - this.agent.streamFn, + this.agent.streamFunction, env, ); summary = compactResult.summary; @@ -2933,7 +2933,7 @@ export class AgentSession { customInstructions, replaceInstructions, reserveTokens: branchSummarySettings.reserveTokens, - streamFn: this.agent.streamFn, + streamFn: this.agent.streamFunction, }); if (result.aborted) { return { cancelled: true, aborted: true }; diff --git a/packages/coding-agent/src/core/sdk.ts b/packages/coding-agent/src/core/sdk.ts index 9ac4bc765..09138f9fc 100644 --- a/packages/coding-agent/src/core/sdk.ts +++ b/packages/coding-agent/src/core/sdk.ts @@ -294,7 +294,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} tools: [], }, convertToLlm: convertToLlmWithBlockImages, - streamFn: async (model, context, options) => { + streamFunction: async (model, context, options) => { const providerRetrySettings = settingsManager.getProviderRetrySettings(); const httpIdleTimeoutMs = settingsManager.getHttpIdleTimeoutMs(); // SDKs treat timeout=0 as 0ms (immediate timeout), not "no timeout". diff --git a/packages/coding-agent/test/agent-session-auto-compaction-queue.test.ts b/packages/coding-agent/test/agent-session-auto-compaction-queue.test.ts index d34d36f3a..3b12b1efe 100644 --- a/packages/coding-agent/test/agent-session-auto-compaction-queue.test.ts +++ b/packages/coding-agent/test/agent-session-auto-compaction-queue.test.ts @@ -3,7 +3,7 @@ import { tmpdir } from "node:os"; import { join } from "node:path"; import { Agent } from "@earendil-works/pi-agent-core"; import { type AssistantMessage, createAssistantMessageEventStream, fauxAssistantMessage } from "@earendil-works/pi-ai"; -import { getModel } from "@earendil-works/pi-ai/compat"; +import { getModel, streamSimple } from "@earendil-works/pi-ai/compat"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { AgentSession } from "../src/core/agent-session.ts"; import { AuthStorage } from "../src/core/auth-storage.ts"; @@ -21,10 +21,10 @@ describe("AgentSession auto-compaction queue resume", () => { beforeEach(async () => { tempDir = join(tmpdir(), `pi-auto-compaction-queue-${Date.now()}`); mkdirSync(tempDir, { recursive: true }); - vi.useFakeTimers(); const model = getModel("anthropic", "claude-sonnet-4-5")!; const agent = new Agent({ + streamFunction: streamSimple, initialState: { model, systemPrompt: "Test", @@ -50,7 +50,6 @@ describe("AgentSession auto-compaction queue resume", () => { afterEach(() => { session.dispose(); - vi.useRealTimers(); vi.restoreAllMocks(); if (tempDir && existsSync(tempDir)) { rmSync(tempDir, { recursive: true }); @@ -84,9 +83,9 @@ describe("AgentSession auto-compaction queue resume", () => { timestamp: now - 500, }); session.agent.state.messages = sessionManager.buildSessionContext().messages; - session.agent.streamFn = (summaryModel) => { + session.agent.streamFunction = (summaryModel) => { const stream = createAssistantMessageEventStream(); - queueMicrotask(() => { + void Promise.resolve().then(() => { stream.push({ type: "done", reason: "stop", diff --git a/packages/coding-agent/test/agent-session-compaction.test.ts b/packages/coding-agent/test/agent-session-compaction.test.ts index 61ad0a2b0..e35158f80 100644 --- a/packages/coding-agent/test/agent-session-compaction.test.ts +++ b/packages/coding-agent/test/agent-session-compaction.test.ts @@ -12,7 +12,7 @@ import { existsSync, mkdirSync, rmSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { Agent } from "@earendil-works/pi-agent-core"; -import { getModel } from "@earendil-works/pi-ai/compat"; +import { getModel, streamSimple } from "@earendil-works/pi-ai/compat"; import { afterEach, beforeEach, describe, expect, it } from "vitest"; import { AgentSession, type AgentSessionEvent } from "../src/core/agent-session.ts"; import { AuthStorage } from "../src/core/auth-storage.ts"; @@ -49,6 +49,7 @@ describe.skipIf(!API_KEY)("AgentSession compaction e2e", () => { const model = getModel("anthropic", "claude-sonnet-4-5")!; const agent = new Agent({ getApiKey: () => API_KEY, + streamFunction: streamSimple, initialState: { model, systemPrompt: "You are a helpful assistant. Be concise.", diff --git a/packages/coding-agent/test/agent-session-concurrent.test.ts b/packages/coding-agent/test/agent-session-concurrent.test.ts index bd825ec7e..bd3ef463d 100644 --- a/packages/coding-agent/test/agent-session-concurrent.test.ts +++ b/packages/coding-agent/test/agent-session-concurrent.test.ts @@ -90,7 +90,7 @@ describe("AgentSession concurrent prompt guard", () => { systemPrompt: "Test", tools: [], }, - streamFn: (_model, _context, options) => { + streamFunction: (_model, _context, options) => { abortSignal = options?.signal; const stream = new MockAssistantStream(); queueMicrotask(() => { @@ -195,7 +195,7 @@ describe("AgentSession concurrent prompt guard", () => { systemPrompt: "Test", tools: [], }, - streamFn: (_model, context, options) => { + streamFunction: (_model, context, options) => { abortSignal = options?.signal; const stream = new MockAssistantStream(); queueMicrotask(() => { @@ -301,7 +301,7 @@ describe("AgentSession concurrent prompt guard", () => { systemPrompt: "Test", tools: [], }, - streamFn: () => { + streamFunction: () => { const stream = new MockAssistantStream(); queueMicrotask(() => { stream.push({ type: "start", partial: createAssistantMessage("") }); @@ -362,7 +362,7 @@ describe("AgentSession concurrent prompt guard", () => { systemPrompt: "Test", tools: [tool], }, - streamFn: async (_model, context) => { + streamFunction: async (_model, context) => { const stream = new MockAssistantStream(); queueMicrotask(() => { const toolResultCount = context.messages.filter((message) => message.role === "toolResult").length; @@ -508,7 +508,7 @@ describe("AgentSession concurrent prompt guard", () => { systemPrompt: "Test", tools: [tool], }, - streamFn: async (_model, context) => { + streamFunction: async (_model, context) => { const stream = new MockAssistantStream(); queueMicrotask(() => { const hasToolResult = context.messages.some((message) => message.role === "toolResult"); diff --git a/packages/coding-agent/test/agent-session-dynamic-provider.test.ts b/packages/coding-agent/test/agent-session-dynamic-provider.test.ts index ce137d2e8..c806efbc3 100644 --- a/packages/coding-agent/test/agent-session-dynamic-provider.test.ts +++ b/packages/coding-agent/test/agent-session-dynamic-provider.test.ts @@ -84,7 +84,7 @@ describe("AgentSession dynamic provider registration", () => { session: Awaited>, ): Promise { let baseUrl: string | undefined; - session.agent.streamFn = async (model) => { + session.agent.streamFunction = async (model) => { baseUrl = model.baseUrl; throw new Error("stop"); }; diff --git a/packages/coding-agent/test/agent-session-retry.test.ts b/packages/coding-agent/test/agent-session-retry.test.ts index f7d5eced7..f61e71cb4 100644 --- a/packages/coding-agent/test/agent-session-retry.test.ts +++ b/packages/coding-agent/test/agent-session-retry.test.ts @@ -82,7 +82,7 @@ describe("AgentSession retry", () => { const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: "Test", tools: [] }, - streamFn: () => { + streamFunction: () => { callCount++; const stream = new MockAssistantStream(); queueMicrotask(() => { @@ -203,7 +203,7 @@ describe("AgentSession retry", () => { const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: "Test", tools: [] }, - streamFn, + streamFunction: streamFn, }); const sessionManager = SessionManager.inMemory(); const settingsManager = SettingsManager.create(tempDir, tempDir); @@ -255,7 +255,7 @@ describe("AgentSession retry", () => { const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: "Test", tools: [] }, - streamFn: () => { + streamFunction: () => { callCount++; const stream = new MockAssistantStream(); queueMicrotask(() => { diff --git a/packages/coding-agent/test/agent-session-stats.test.ts b/packages/coding-agent/test/agent-session-stats.test.ts index 2e9d6bf47..37706c566 100644 --- a/packages/coding-agent/test/agent-session-stats.test.ts +++ b/packages/coding-agent/test/agent-session-stats.test.ts @@ -1,5 +1,11 @@ import { Agent } from "@earendil-works/pi-agent-core"; -import { type AssistantMessage, getModel, type ToolResultMessage, type Usage } from "@earendil-works/pi-ai/compat"; +import { + type AssistantMessage, + getModel, + streamSimple, + type ToolResultMessage, + type Usage, +} from "@earendil-works/pi-ai/compat"; import { describe, expect, it } from "vitest"; import { AgentSession } from "../src/core/agent-session.ts"; import { AuthStorage } from "../src/core/auth-storage.ts"; @@ -69,6 +75,7 @@ async function createSession() { const session = new AgentSession({ agent: new Agent({ getApiKey: () => "test-key", + streamFunction: streamSimple, initialState: { model, systemPrompt: "You are a helpful assistant.", diff --git a/packages/coding-agent/test/compaction-extensions.test.ts b/packages/coding-agent/test/compaction-extensions.test.ts index fabcb38dc..053be534c 100644 --- a/packages/coding-agent/test/compaction-extensions.test.ts +++ b/packages/coding-agent/test/compaction-extensions.test.ts @@ -7,7 +7,7 @@ import { existsSync, mkdirSync, rmSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { Agent } from "@earendil-works/pi-agent-core"; -import { getModel } from "@earendil-works/pi-ai/compat"; +import { getModel, streamSimple } from "@earendil-works/pi-ai/compat"; import { afterEach, beforeEach, describe, expect, it } from "vitest"; import { AgentSession } from "../src/core/agent-session.ts"; import { AuthStorage } from "../src/core/auth-storage.ts"; @@ -89,6 +89,7 @@ describe.skipIf(!API_KEY)("Compaction extensions", () => { const model = getModel("anthropic", "claude-sonnet-4-5")!; const agent = new Agent({ getApiKey: () => API_KEY, + streamFunction: streamSimple, initialState: { model, systemPrompt: "You are a helpful assistant. Be concise.", diff --git a/packages/coding-agent/test/rpc-prompt-response-semantics.test.ts b/packages/coding-agent/test/rpc-prompt-response-semantics.test.ts index 155d2d384..8279d57c7 100644 --- a/packages/coding-agent/test/rpc-prompt-response-semantics.test.ts +++ b/packages/coding-agent/test/rpc-prompt-response-semantics.test.ts @@ -114,7 +114,7 @@ async function createRuntimeHost(options: { withAuth: boolean; responseDelayMs: systemPrompt: "Test", tools: [], }, - streamFn: (_model, _context, _options) => { + streamFunction: (_model, _context, _options) => { const stream = new MockAssistantStream(); queueMicrotask(() => { stream.push({ type: "start", partial: createAssistantMessage("") }); diff --git a/packages/coding-agent/test/sdk-openrouter-attribution.test.ts b/packages/coding-agent/test/sdk-openrouter-attribution.test.ts index 53f7e9427..8cafaccdc 100644 --- a/packages/coding-agent/test/sdk-openrouter-attribution.test.ts +++ b/packages/coding-agent/test/sdk-openrouter-attribution.test.ts @@ -126,7 +126,7 @@ describe("createAgentSession provider attribution headers", () => { }); try { - const stream = await session.agent.streamFn( + const stream = await session.agent.streamFunction( model, { messages: [] }, { diff --git a/packages/coding-agent/test/sdk-stream-options.test.ts b/packages/coding-agent/test/sdk-stream-options.test.ts index f61cade08..f4c2fd948 100644 --- a/packages/coding-agent/test/sdk-stream-options.test.ts +++ b/packages/coding-agent/test/sdk-stream-options.test.ts @@ -114,7 +114,7 @@ describe("createAgentSession stream options", () => { }); try { - const stream = await session.agent.streamFn(model, { messages: [] }, requestOptions); + const stream = await session.agent.streamFunction(model, { messages: [] }, requestOptions); await stream.result(); return capturedOptions; } finally { diff --git a/packages/coding-agent/test/suite/agent-session-compaction.test.ts b/packages/coding-agent/test/suite/agent-session-compaction.test.ts index ee7bea77a..33649135e 100644 --- a/packages/coding-agent/test/suite/agent-session-compaction.test.ts +++ b/packages/coding-agent/test/suite/agent-session-compaction.test.ts @@ -49,7 +49,7 @@ function createAssistant( function useSummaryStreamFn(harness: Harness, summary: string): () => number { let callCount = 0; - harness.session.agent.streamFn = (model) => { + harness.session.agent.streamFunction = (model) => { callCount++; const stream = createAssistantMessageEventStream(); queueMicrotask(() => { diff --git a/packages/coding-agent/test/suite/harness.ts b/packages/coding-agent/test/suite/harness.ts index b7b279716..95c5dbba6 100644 --- a/packages/coding-agent/test/suite/harness.ts +++ b/packages/coding-agent/test/suite/harness.ts @@ -14,7 +14,7 @@ import type { FauxResponseStep, Model, } from "@earendil-works/pi-ai/compat"; -import { registerFauxProvider } from "@earendil-works/pi-ai/compat"; +import { registerFauxProvider, streamSimple } from "@earendil-works/pi-ai/compat"; import { AgentSession, type AgentSessionEvent } from "../../src/core/agent-session.ts"; import { AuthStorage } from "../../src/core/auth-storage.ts"; import type { ExtensionRunner } from "../../src/core/extensions/index.ts"; @@ -137,6 +137,7 @@ export async function createHarness(options: HarnessOptions = {}): Promise (withConfiguredAuth ? "faux-key" : undefined), + streamFunction: streamSimple, initialState: { model, systemPrompt: options.systemPrompt ?? "You are a test assistant.", diff --git a/packages/coding-agent/test/suite/regressions/5596-missing-theme-export.test.ts b/packages/coding-agent/test/suite/regressions/5596-missing-theme-export.test.ts index e8ef3b354..b6efaa648 100644 --- a/packages/coding-agent/test/suite/regressions/5596-missing-theme-export.test.ts +++ b/packages/coding-agent/test/suite/regressions/5596-missing-theme-export.test.ts @@ -2,7 +2,7 @@ import { existsSync, mkdtempSync, rmSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { Agent } from "@earendil-works/pi-agent-core"; -import { fauxAssistantMessage, registerFauxProvider } from "@earendil-works/pi-ai/compat"; +import { fauxAssistantMessage, registerFauxProvider, streamSimple } from "@earendil-works/pi-ai/compat"; import { afterEach, describe, expect, it } from "vitest"; import { AgentSession } from "../../../src/core/agent-session.ts"; import { AuthStorage } from "../../../src/core/auth-storage.ts"; @@ -61,6 +61,7 @@ describe("regression #5596: missing configured theme export", () => { tools: [], }, convertToLlm, + streamFunction: streamSimple, }); const session = new AgentSession({ agent, diff --git a/packages/coding-agent/test/suite/regressions/6324-branch-summary-ambient-auth.test.ts b/packages/coding-agent/test/suite/regressions/6324-branch-summary-ambient-auth.test.ts index bfedbaafd..d116325d6 100644 --- a/packages/coding-agent/test/suite/regressions/6324-branch-summary-ambient-auth.test.ts +++ b/packages/coding-agent/test/suite/regressions/6324-branch-summary-ambient-auth.test.ts @@ -17,7 +17,7 @@ describe("issue #6324 branch summary ambient auth", () => { harnesses.push(harness); let streamCallCount = 0; - harness.session.agent.streamFn = (model, _context, options) => { + harness.session.agent.streamFunction = (model, _context, options) => { streamCallCount++; expect(options?.apiKey).toBeUndefined(); diff --git a/packages/coding-agent/test/test-harness.ts b/packages/coding-agent/test/test-harness.ts index 3dc9f628e..70fbf30f5 100644 --- a/packages/coding-agent/test/test-harness.ts +++ b/packages/coding-agent/test/test-harness.ts @@ -378,7 +378,7 @@ async function createHarnessWithResourceLoader( systemPrompt: options.systemPrompt ?? "You are a test assistant.", tools: options.tools ?? [], }, - streamFn, + streamFunction: streamFn, }); const sessionManager = SessionManager.inMemory(); diff --git a/packages/coding-agent/test/utilities.ts b/packages/coding-agent/test/utilities.ts index f7568e9f2..e24845059 100644 --- a/packages/coding-agent/test/utilities.ts +++ b/packages/coding-agent/test/utilities.ts @@ -8,7 +8,7 @@ import { homedir, tmpdir } from "node:os"; import { dirname, join } from "node:path"; import { Agent } from "@earendil-works/pi-agent-core"; import type { OAuthCredentials } from "@earendil-works/pi-ai"; -import { getModel } from "@earendil-works/pi-ai/compat"; +import { getModel, streamSimple } from "@earendil-works/pi-ai/compat"; import { builtinProviders } from "@earendil-works/pi-ai/providers/all"; import { AgentSession } from "../src/core/agent-session.ts"; import { AuthStorage } from "../src/core/auth-storage.ts"; @@ -246,6 +246,7 @@ export async function createTestSession(options: TestSessionOptions = {}): Promi systemPrompt: options.systemPrompt ?? "You are a helpful assistant. Be extremely concise.", tools: createCodingTools(process.cwd()), }, + streamFunction: streamSimple, }); const sessionManager = options.inMemory ? SessionManager.inMemory() : SessionManager.create(tempDir); diff --git a/scripts/agent-treeshake-smoke-entry.ts b/scripts/agent-treeshake-smoke-entry.ts new file mode 100644 index 000000000..8a8206ecd --- /dev/null +++ b/scripts/agent-treeshake-smoke-entry.ts @@ -0,0 +1,13 @@ +import { Agent } from "@earendil-works/pi-agent-core"; +import { createModels } from "@earendil-works/pi-ai"; +import { anthropicProvider } from "@earendil-works/pi-ai/providers/anthropic"; + +const models = createModels(); +models.setProvider(anthropicProvider()); +const model = models.getModel("anthropic", "claude-sonnet-4-5"); +if (!model) throw new Error("Anthropic smoke-test model not found"); + +export const agent = new Agent({ + initialState: { model }, + streamFunction: models.streamSimple.bind(models), +}); diff --git a/scripts/browser-smoke-entry.ts b/scripts/browser-smoke-entry.ts index 3ac7bcf02..2a7ccfbe1 100644 --- a/scripts/browser-smoke-entry.ts +++ b/scripts/browser-smoke-entry.ts @@ -1,5 +1,5 @@ import { createAssistantMessageEventStream, Type } from "@earendil-works/pi-ai"; -import { complete, getModel, getProviders } from "@earendil-works/pi-ai/compat"; +import { complete, getModel, getProviders, streamSimple } from "@earendil-works/pi-ai/compat"; import { Agent, bashExecutionToText, @@ -24,7 +24,7 @@ const model = getModel("google", "gemini-2.5-flash"); const schema = Type.Object({ prompt: Type.String() }); const stream = createAssistantMessageEventStream(); -const agent = new Agent({ initialState: { model } }); +const agent = new Agent({ initialState: { model }, streamFunction: streamSimple }); agent.steer({ role: "user", content: [{ type: "text", text: "queued" }], timestamp: 0 }); const repo = new InMemorySessionRepo(); const result = getOrThrow(ok({ value: 1 })); diff --git a/scripts/check-browser-smoke.mjs b/scripts/check-browser-smoke.mjs index a45bd3a9c..eeab27b15 100644 --- a/scripts/check-browser-smoke.mjs +++ b/scripts/check-browser-smoke.mjs @@ -4,6 +4,7 @@ import { dirname, join, resolve } from "node:path"; import { build } from "esbuild"; const outputPath = join(tmpdir(), "pi-browser-smoke.js"); +const agentTreeshakeOutputPath = join(tmpdir(), "pi-agent-treeshake-smoke.js"); const errorLogPath = join(tmpdir(), "pi-browser-smoke-errors.log"); const generatedCatalogDataDir = join(process.cwd(), "packages/ai/src/providers/data"); @@ -23,6 +24,22 @@ const generatedCatalogDataPlugin = { }, }; +function normalizePath(path) { + return path.replaceAll("\\", "/"); +} + +function findInput(inputs, suffix) { + return Object.keys(inputs).find((input) => { + const normalized = normalizePath(input); + return normalized === suffix || normalized.endsWith(`/${suffix}`); + }); +} + +function includesNodePackage(inputs, packageName) { + const marker = `node_modules/${packageName}/`; + return Object.keys(inputs).some((input) => normalizePath(input).includes(marker)); +} + try { await build({ entryPoints: ["scripts/browser-smoke-entry.ts"], @@ -33,6 +50,47 @@ try { outfile: outputPath, plugins: [generatedCatalogDataPlugin], }); + + const agentTreeshakeBuild = await build({ + entryPoints: ["scripts/agent-treeshake-smoke-entry.ts"], + bundle: true, + platform: "browser", + format: "esm", + logLevel: "silent", + metafile: true, + outfile: agentTreeshakeOutputPath, + plugins: [generatedCatalogDataPlugin], + write: false, + }); + const inputs = agentTreeshakeBuild.metafile.inputs; + for (const forbiddenInput of [ + "packages/ai/src/compat.ts", + "packages/ai/src/models.generated.ts", + "packages/ai/src/providers/all.ts", + ]) { + const includedInput = findInput(inputs, forbiddenInput); + if (includedInput) { + throw new Error(`Agent selective-provider bundle unexpectedly includes ${includedInput}`); + } + } + + const aiSdkPackages = [ + "@anthropic-ai/sdk", + "@aws-sdk/client-bedrock-runtime", + "@google/genai", + "@mistralai/mistralai", + "openai", + ]; + const includedAiSdkPackages = aiSdkPackages.filter((packageName) => includesNodePackage(inputs, packageName)); + if ( + includedAiSdkPackages.length !== 1 || + includedAiSdkPackages[0] !== "@anthropic-ai/sdk" + ) { + throw new Error( + `Agent selective-provider bundle SDKs: expected only @anthropic-ai/sdk, found ${includedAiSdkPackages.join(", ") || "none"}`, + ); + } + process.exit(0); } catch (error) { let detailedErrors = ""; From ff992261e2ec349e4257a59d052076355ca0b18b Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Mon, 20 Jul 2026 20:00:14 +0000 Subject: [PATCH 28/62] chore: approve contributor R-Taneja --- .github/APPROVED_CONTRIBUTORS | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/APPROVED_CONTRIBUTORS b/.github/APPROVED_CONTRIBUTORS index 1c70bd0e9..947f19a75 100644 --- a/.github/APPROVED_CONTRIBUTORS +++ b/.github/APPROVED_CONTRIBUTORS @@ -283,3 +283,5 @@ anh-chu pr rsaryev pr QuintinShaw pr + +R-Taneja pr From c8c3cd499f4d35c0f9cfebfec5f4e3822411a49f Mon Sep 17 00:00:00 2001 From: Armin Ronacher Date: Mon, 20 Jul 2026 22:28:48 +0200 Subject: [PATCH 29/62] fix(ai): validate generated model data before builds --- .gitignore | 1 + README.md | 5 +- package.json | 3 + packages/ai/CHANGELOG.md | 4 + packages/ai/package.json | 9 +- packages/ai/scripts/check-model-data.ts | 16 ++ packages/ai/scripts/generate-image-models.ts | 121 +++++--- packages/ai/scripts/generate-models.ts | 202 ++++++++++---- packages/ai/scripts/model-data.ts | 261 ++++++++++++++++++ packages/ai/test/image-model-data.test.ts | 47 ++++ .../ai/test/model-data-validation.test.ts | 123 +++++++++ scripts/check-browser-smoke.mjs | 2 +- scripts/local-release.mjs | 8 +- scripts/release.mjs | 1 + 14 files changed, 700 insertions(+), 103 deletions(-) create mode 100644 packages/ai/scripts/check-model-data.ts create mode 100644 packages/ai/scripts/model-data.ts create mode 100644 packages/ai/test/image-model-data.test.ts create mode 100644 packages/ai/test/model-data-validation.test.ts diff --git a/.gitignore b/.gitignore index 8cfb05c71..0e241a9dc 100644 --- a/.gitignore +++ b/.gitignore @@ -9,6 +9,7 @@ packages/*/dist/ packages/*/dist-chrome/ packages/*/dist-firefox/ packages/ai/src/providers/data/ +packages/ai/src/providers/.model-generation-*/ *.cpuprofile # Environment diff --git a/README.md b/README.md index 130a412f0..900779b0a 100644 --- a/README.md +++ b/README.md @@ -52,8 +52,9 @@ See [CONTRIBUTING.md](CONTRIBUTING.md) for contribution guidelines and [AGENTS.m ```bash npm install --ignore-scripts # Install all dependencies without running lifecycle scripts -npm run build # Build all packages -npm run check # Lint, format, and type check +npm run build # Refresh model data, then build all packages +npm run build:offline # Rebuild using existing model data without network access +npm run check # Lint, format, and type check ./test.sh # Run tests (skips LLM-dependent tests without API keys) ./pi-test.sh # Run pi from sources (can be run from any directory) ``` diff --git a/package.json b/package.json index c919d33e7..fc3dd5dbd 100644 --- a/package.json +++ b/package.json @@ -13,6 +13,7 @@ "scripts": { "clean": "npm run clean --workspaces", "build": "cd packages/tui && npm run build && cd ../ai && npm run build && cd ../agent && npm run build && cd ../coding-agent && npm run build && cd ../orchestrator && npm run build", + "build:offline": "cd packages/tui && npm run build && cd ../ai && npm run build:offline && cd ../agent && npm run build && cd ../coding-agent && npm run build && cd ../orchestrator && npm run build", "check": "biome check --write --error-on-warnings . && npm run check:pinned-deps && npm run check:ts-imports && npm run check:shrinkwrap && npm run check:install-lock:coding-agent && tsgo --noEmit && npm run check:browser-smoke", "check:browser-smoke": "node scripts/check-browser-smoke.mjs", "check:pinned-deps": "node scripts/check-pinned-deps.mjs", @@ -20,6 +21,8 @@ "check:install-lock:coding-agent": "node scripts/generate-coding-agent-install-lock.mjs --check", "check:ts-imports": "node scripts/check-ts-relative-imports.mjs", "generate:models": "npm --prefix packages/ai run generate-models && npm --prefix packages/ai run generate-image-models", + "hydrate:model-data": "npm --prefix packages/ai run hydrate-model-data", + "check:model-data": "npm --prefix packages/ai run check:model-data", "generate:model-catalog": "npm --prefix packages/ai run generate-model-catalog", "diff:model-catalog": "node scripts/diff-model-catalog.mjs", "check:model-catalog": "node scripts/publish-model-catalog.mjs --input .artifacts/model-catalog --dry-run", diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 56693a287..b3a897367 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -8,6 +8,10 @@ - Added a shared `uuidv7` utility for time-ordered identifiers. - Added optional usage metadata to tool result messages ([#6671](https://github.com/earendil-works/pi/pull/6671) by [@davidbrai](https://github.com/davidbrai)). +### Changed + +- Changed model generation to validate ignored provider data before compilation; `npm run build` refreshes model data as before, while `npm run build:offline` reuses existing data without network access. + ### Fixed - Fixed sessionless OpenAI Codex WebSocket requests to use UUIDv7 request IDs, enabling models that reject UUIDv4 IDs. diff --git a/packages/ai/package.json b/packages/ai/package.json index 58b9113f8..ce936c9e1 100644 --- a/packages/ai/package.json +++ b/packages/ai/package.json @@ -49,10 +49,13 @@ ], "scripts": { "clean": "shx rm -rf dist", - "generate-models": "node scripts/generate-models.ts", + "generate-models": "node scripts/generate-models.ts --strict", + "hydrate-model-data": "node scripts/generate-models.ts --strict --data-only", "generate-model-catalog": "node scripts/generate-models.ts --strict --json-only --json-output ../../.artifacts/model-catalog", - "generate-image-models": "node scripts/generate-image-models.ts", - "build": "npm run generate-models && tsgo -p tsconfig.build.json && shx rm -rf dist/providers/data && shx cp -r src/providers/data dist/providers/data", + "generate-image-models": "node scripts/generate-image-models.ts --strict", + "check:model-data": "node scripts/check-model-data.ts", + "build": "npm run generate-models && npm run build:offline", + "build:offline": "npm run check:model-data && tsgo -p tsconfig.build.json && shx rm -rf dist/providers/data && shx cp -r src/providers/data dist/providers/data", "test": "vitest --run", "prepublishOnly": "npm run clean && npm run build" }, diff --git a/packages/ai/scripts/check-model-data.ts b/packages/ai/scripts/check-model-data.ts new file mode 100644 index 000000000..5c61b52c4 --- /dev/null +++ b/packages/ai/scripts/check-model-data.ts @@ -0,0 +1,16 @@ +#!/usr/bin/env node + +import { dirname, join } from "node:path"; +import { fileURLToPath } from "node:url"; +import { validateGeneratedModelData } from "./model-data.ts"; + +const packageRoot = join(dirname(fileURLToPath(import.meta.url)), ".."); + +try { + validateGeneratedModelData(packageRoot); + console.log("Generated model data is valid."); +} catch (error) { + console.error(error instanceof Error ? error.message : String(error)); + console.error("\nModel data is missing or stale. Run `npm run hydrate:model-data` from the repository root."); + process.exitCode = 1; +} diff --git a/packages/ai/scripts/generate-image-models.ts b/packages/ai/scripts/generate-image-models.ts index 2744da487..26e46e450 100644 --- a/packages/ai/scripts/generate-image-models.ts +++ b/packages/ai/scripts/generate-image-models.ts @@ -1,7 +1,7 @@ #!/usr/bin/env node import { writeFileSync } from "fs"; -import { dirname, join } from "path"; +import { dirname, join, resolve } from "path"; import { fileURLToPath } from "url"; import type { ImagesModel } from "../src/types.ts"; @@ -10,6 +10,13 @@ const __dirname = dirname(__filename); const packageRoot = join(__dirname, ".."); const OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1"; +function readStrictOption(args: string[]): boolean { + for (const arg of args) { + if (arg !== "--strict") throw new Error(`Unknown argument: ${arg}`); + } + return args.includes("--strict"); +} + interface OpenRouterModelRecord { id: string; name: string; @@ -26,52 +33,73 @@ interface OpenRouterModelRecord { }; } -async function fetchOpenRouterImageModels(): Promise[]> { - try { - console.log("Fetching image models from OpenRouter API..."); - const response = await fetch(`${OPENROUTER_BASE_URL}/models?output_modalities=image`); - const data = (await response.json()) as { data?: OpenRouterModelRecord[] }; - const models: ImagesModel<"openrouter-images">[] = []; - - for (const model of data.data ?? []) { - const input = Array.from( - new Set( - (model.architecture?.input_modalities ?? []) - .filter((modality): modality is "text" | "image" => modality === "text" || modality === "image"), +export function parseOpenRouterImageModels( + payload: unknown, + strict: boolean, +): ImagesModel<"openrouter-images">[] { + const data = + typeof payload === "object" && payload !== null + ? (payload as { data?: OpenRouterModelRecord[] }).data + : undefined; + if (!Array.isArray(data) || data.length === 0) { + if (strict) throw new Error("OpenRouter API returned a missing or empty image model list"); + return []; + } + + const models: ImagesModel<"openrouter-images">[] = []; + for (const model of data) { + const input = Array.from( + new Set( + (model.architecture?.input_modalities ?? []).filter( + (modality): modality is "text" | "image" => modality === "text" || modality === "image", ), - ); - const output = Array.from( - new Set( - (model.architecture?.output_modalities ?? []).filter( - (modality): modality is "text" | "image" => modality === "text" || modality === "image", - ), + ), + ); + const output = Array.from( + new Set( + (model.architecture?.output_modalities ?? []).filter( + (modality): modality is "text" | "image" => modality === "text" || modality === "image", ), - ); - - if (!output.includes("image")) continue; - if (input.length === 0) input.push("text"); - - models.push({ - id: model.id, - name: model.name, - api: "openrouter-images", - provider: "openrouter", - baseUrl: OPENROUTER_BASE_URL, - input, - output, - cost: { - input: parseFloat(model.pricing?.prompt || "0") * 1_000_000, - output: parseFloat(model.pricing?.completion || "0") * 1_000_000, - cacheRead: parseFloat(model.pricing?.input_cache_read || "0") * 1_000_000, - cacheWrite: parseFloat(model.pricing?.input_cache_write || "0") * 1_000_000, - }, - }); - } + ), + ); + + if (!output.includes("image")) continue; + if (input.length === 0) input.push("text"); + models.push({ + id: model.id, + name: model.name, + api: "openrouter-images", + provider: "openrouter", + baseUrl: OPENROUTER_BASE_URL, + input, + output, + cost: { + input: parseFloat(model.pricing?.prompt || "0") * 1_000_000, + output: parseFloat(model.pricing?.completion || "0") * 1_000_000, + cacheRead: parseFloat(model.pricing?.input_cache_read || "0") * 1_000_000, + cacheWrite: parseFloat(model.pricing?.input_cache_write || "0") * 1_000_000, + }, + }); + } + + if (strict && models.length === 0) { + throw new Error("OpenRouter API returned no usable image models"); + } + return models; +} + +async function fetchOpenRouterImageModels(strict: boolean): Promise[]> { + try { + console.log("Fetching image models from OpenRouter API..."); + const response = await fetch(`${OPENROUTER_BASE_URL}/models?output_modalities=image`); + if (!response.ok) throw new Error(`OpenRouter API returned ${response.status}`); + const models = parseOpenRouterImageModels(await response.json(), strict); console.log(`Fetched ${models.length} image models from OpenRouter`); return models; } catch (error) { console.error("Failed to fetch OpenRouter image models:", error); + if (strict) throw error; return []; } } @@ -118,14 +146,17 @@ ${providerEntries} } async function main(): Promise { - const models = await fetchOpenRouterImageModels(); + const strict = readStrictOption(process.argv.slice(2)); + const models = await fetchOpenRouterImageModels(strict); const output = generateImageModelsFile(models); const outputPath = join(packageRoot, "src", "image-models.generated.ts"); writeFileSync(outputPath, output, "utf-8"); console.log(`Generated ${outputPath}`); } -main().catch((error) => { - console.error(error); - process.exit(1); -}); +if (process.argv[1] && resolve(process.argv[1]) === __filename) { + main().catch((error) => { + console.error(error); + process.exit(1); + }); +} diff --git a/packages/ai/scripts/generate-models.ts b/packages/ai/scripts/generate-models.ts index 82700ac99..66e64e5a1 100644 --- a/packages/ai/scripts/generate-models.ts +++ b/packages/ai/scripts/generate-models.ts @@ -1,6 +1,6 @@ #!/usr/bin/env node -import { mkdirSync, readdirSync, rmSync, writeFileSync } from "fs"; +import { existsSync, mkdirSync, mkdtempSync, readFileSync, readdirSync, renameSync, rmSync, writeFileSync } from "fs"; import { dirname, join, resolve } from "path"; import { fileURLToPath } from "url"; import { @@ -18,6 +18,14 @@ import type { OpenAICompletionsCompat, OpenAIResponsesCompat, } from "../src/types.ts"; +import { + createModelDataManifest, + type ModelDataStructure, + MODEL_DATA_MANIFEST_FILE, + readModelDataStructure, + validateGeneratedModelData, + validateModelDataDirectory, +} from "./model-data.ts"; const __filename = fileURLToPath(import.meta.url); const __dirname = dirname(__filename); @@ -25,11 +33,13 @@ const packageRoot = join(__dirname, ".."); function readGeneratorOptions(args: string[]): { strict: boolean; + dataOnly: boolean; jsonOnly: boolean; jsonOutputDir: string | undefined; pretty: boolean; } { let strict = false; + let dataOnly = false; let jsonOnly = false; let jsonOutputDir: string | undefined; let pretty = false; @@ -40,6 +50,10 @@ function readGeneratorOptions(args: string[]): { strict = true; continue; } + if (arg === "--data-only") { + dataOnly = true; + continue; + } if (arg === "--json-only") { jsonOnly = true; continue; @@ -58,7 +72,8 @@ function readGeneratorOptions(args: string[]): { } if (jsonOnly && !jsonOutputDir) throw new Error("--json-only requires --json-output"); - return { strict, jsonOnly, jsonOutputDir, pretty }; + if (dataOnly && (jsonOnly || jsonOutputDir)) throw new Error("--data-only cannot be combined with JSON catalog output"); + return { strict, dataOnly, jsonOnly, jsonOutputDir, pretty }; } const generatorOptions = readGeneratorOptions(process.argv.slice(2)); @@ -2403,63 +2418,154 @@ async function generateModels() { jsonProviders[providerId][modelId] = providers[providerId][modelId]; } } - const writeJson = (path: string, value: unknown) => - writeFileSync(path, `${JSON.stringify(value, null, generatorOptions.pretty ? 2 : undefined)}\n`); + + const serializeJson = (value: unknown) => `${JSON.stringify(value, null, generatorOptions.pretty ? 2 : undefined)}\n`; + const writeJson = (path: string, value: unknown) => writeFileSync(path, serializeJson(value)); + let generatedDataProviderIds = sortedProviderIds; + let generatedDataProviders = jsonProviders; + let modelDataStructure: ModelDataStructure = Object.fromEntries( + sortedProviderIds.map((providerId) => [ + providerId, + Object.fromEntries( + Object.entries(jsonProviders[providerId]).map(([modelId, model]) => [modelId, model.api]), + ), + ]), + ); + + if (generatorOptions.dataOnly) { + modelDataStructure = readModelDataStructure(packageRoot); + generatedDataProviderIds = Object.keys(modelDataStructure); + const hydratedProviders: typeof jsonProviders = {}; + const hydrationErrors: string[] = []; + for (const [providerId, expectedModels] of Object.entries(modelDataStructure)) { + hydratedProviders[providerId] = {}; + for (const [modelId, expectedApi] of Object.entries(expectedModels)) { + const model = jsonProviders[providerId]?.[modelId]; + if (!model) { + hydrationErrors.push(`missing ${providerId}/${modelId}`); + continue; + } + if (model.api !== expectedApi) { + hydrationErrors.push(`${providerId}/${modelId} uses ${model.api}, expected ${expectedApi}`); + continue; + } + hydratedProviders[providerId][modelId] = model; + } + } + if (hydrationErrors.length > 0) { + throw new Error(`Cannot hydrate the committed model catalog:\n${hydrationErrors.map((error) => ` - ${error}`).join("\n")}`); + } + generatedDataProviders = hydratedProviders; + } if (!generatorOptions.jsonOnly) { - // Generate TypeScript structural catalogs and adjacent JSON values. - const generatedHeader = `// This file is auto-generated by scripts/generate-models.ts + // Stage and validate all provider values before replacing the current generated data. + const providersDir = join(packageRoot, "src/providers"); + const dataDir = join(providersDir, "data"); + const stagingRoot = mkdtempSync(join(providersDir, ".model-generation-")); + const stagedDataDir = join(stagingRoot, "data"); + const previousDataDir = join(stagingRoot, "previous-data"); + let restoreStructuralCatalog: (() => void) | undefined; + try { + mkdirSync(stagedDataDir, { recursive: true }); + const fileContents: Record = {}; + for (const providerId of generatedDataProviderIds) { + const filename = `${providerId}.json`; + const content = serializeJson(generatedDataProviders[providerId]); + fileContents[filename] = content; + writeFileSync(join(stagedDataDir, filename), content); + } + writeJson( + join(stagedDataDir, MODEL_DATA_MANIFEST_FILE), + createModelDataManifest(modelDataStructure, fileContents), + ); + validateModelDataDirectory(modelDataStructure, stagedDataDir); + + if (!generatorOptions.dataOnly) { + // Generate TypeScript structural catalogs only after the model data is complete and valid. + const previousShardContents = new Map( + readdirSync(providersDir) + .filter((entry) => entry.endsWith(".models.ts")) + .map((entry) => [entry, readFileSync(join(providersDir, entry), "utf8")] as const), + ); + const aggregatorPath = join(packageRoot, "src/models.generated.ts"); + const previousAggregator = readFileSync(aggregatorPath, "utf8"); + restoreStructuralCatalog = () => { + for (const entry of readdirSync(providersDir)) { + if (entry.endsWith(".models.ts")) rmSync(join(providersDir, entry)); + } + for (const [entry, content] of previousShardContents) { + writeFileSync(join(providersDir, entry), content); + } + writeFileSync(aggregatorPath, previousAggregator); + }; + + const generatedHeader = `// This file is auto-generated by scripts/generate-models.ts // Do not edit manually - run 'npm run generate-models' to update `; - const catalogConstName = (providerId: string) => - `${providerId.toUpperCase().replace(/[^A-Z0-9]+/g, "_")}_MODELS`; - const providersDir = join(packageRoot, "src/providers"); - const dataDir = join(providersDir, "data"); + const catalogConstName = (providerId: string) => + `${providerId.toUpperCase().replace(/[^A-Z0-9]+/g, "_")}_MODELS`; + const generatedShardFiles = new Set(); - function emitModelShape(model: Model, indent: string): string { - return `${indent}${JSON.stringify(model.id)}: Model<${JSON.stringify(model.api)}> & {\n${indent}\tid: ${JSON.stringify(model.id)};\n${indent}\tprovider: ${JSON.stringify(model.provider)};\n${indent}};\n`; - } + function emitModelShape(model: Model, indent: string): string { + return `${indent}${JSON.stringify(model.id)}: Model<${JSON.stringify(model.api)}> & {\n${indent}\tid: ${JSON.stringify(model.id)};\n${indent}\tprovider: ${JSON.stringify(model.provider)};\n${indent}};\n`; + } - // Remove stale per-provider catalogs and their generated values. - for (const entry of readdirSync(providersDir)) { - if (entry.endsWith(".models.ts")) { - rmSync(join(providersDir, entry)); - } - } - rmSync(dataDir, { recursive: true, force: true }); - mkdirSync(dataDir, { recursive: true }); + for (const providerId of sortedProviderIds) { + const models = providers[providerId]; + let output = generatedHeader; + output += `import values from "./data/${providerId}.json" with { type: "json" };\n`; + output += `import type { Model } from "../types.ts";\n\n`; + output += `export const ${catalogConstName(providerId)} = values as {\n`; + for (const modelId of Object.keys(models).sort()) { + output += emitModelShape(models[modelId], "\t"); + } + output += `};\n`; + const filename = `${providerId}.models.ts`; + generatedShardFiles.add(filename); + writeFileSync(join(providersDir, filename), output); + } + for (const entry of readdirSync(providersDir)) { + if (entry.endsWith(".models.ts") && !generatedShardFiles.has(entry)) rmSync(join(providersDir, entry)); + } + console.log(`Generated ${sortedProviderIds.length} catalog structures under src/providers/`); - // Per-provider catalog structure and values (sorted for deterministic output). - for (const providerId of sortedProviderIds) { - const models = providers[providerId]; - const sortedModelIds = Object.keys(models).sort(); - let output = generatedHeader; - output += `import values from "./data/${providerId}.json" with { type: "json" };\n`; - output += `import type { Model } from "../types.ts";\n\n`; - output += `export const ${catalogConstName(providerId)} = values as {\n`; - for (const modelId of sortedModelIds) { - output += emitModelShape(models[modelId], "\t"); + let output = generatedHeader; + for (const providerId of sortedProviderIds) { + output += `import { ${catalogConstName(providerId)} } from "./providers/${providerId}.models.ts";\n`; + } + output += `\nexport const MODELS = {\n`; + for (const providerId of sortedProviderIds) { + output += `\t${JSON.stringify(providerId)}: ${catalogConstName(providerId)},\n`; + } + output += `} as const;\n`; + writeFileSync(aggregatorPath, output); + console.log("Generated src/models.generated.ts"); } - output += `};\n`; - writeFileSync(join(providersDir, `${providerId}.models.ts`), output); - writeJson(join(dataDir, `${providerId}.json`), jsonProviders[providerId]); - } - console.log(`Generated ${sortedProviderIds.length} catalog structures under src/providers/`); - console.log("Generated JSON model values under src/providers/data/"); - // Aggregator - let output = generatedHeader; - for (const providerId of sortedProviderIds) { - output += `import { ${catalogConstName(providerId)} } from "./providers/${providerId}.models.ts";\n`; - } - output += `\nexport const MODELS = {\n`; - for (const providerId of sortedProviderIds) { - output += `\t${JSON.stringify(providerId)}: ${catalogConstName(providerId)},\n`; + const hadPreviousData = existsSync(dataDir); + if (hadPreviousData) renameSync(dataDir, previousDataDir); + try { + renameSync(stagedDataDir, dataDir); + validateGeneratedModelData(packageRoot); + } catch (error) { + rmSync(dataDir, { recursive: true, force: true }); + if (hadPreviousData && existsSync(previousDataDir)) renameSync(previousDataDir, dataDir); + throw error; + } + restoreStructuralCatalog = undefined; + console.log( + generatorOptions.dataOnly + ? "Hydrated JSON model values under src/providers/data/" + : "Generated JSON model values under src/providers/data/", + ); + } catch (error) { + restoreStructuralCatalog?.(); + throw error; + } finally { + rmSync(stagingRoot, { recursive: true, force: true }); } - output += `} as const;\n`; - writeFileSync(join(packageRoot, "src/models.generated.ts"), output); - console.log("Generated src/models.generated.ts"); } if (generatorOptions.jsonOutputDir) { diff --git a/packages/ai/scripts/model-data.ts b/packages/ai/scripts/model-data.ts new file mode 100644 index 000000000..588bc6956 --- /dev/null +++ b/packages/ai/scripts/model-data.ts @@ -0,0 +1,261 @@ +import { createHash } from "node:crypto"; +import { existsSync, readFileSync, readdirSync, statSync } from "node:fs"; +import { join } from "node:path"; + +export const MODEL_DATA_SCHEMA_VERSION = 1; +export const MODEL_DATA_MANIFEST_FILE = ".manifest.json"; + +export type ModelDataStructure = Record>; + +export interface ModelDataManifest { + schemaVersion: number; + structureHash: string; + files: Record; +} + +const JSON_STRING_PATTERN = '"(?:\\\\.|[^"\\\\])*"'; +const MODEL_SHAPE_PATTERN = new RegExp(`^\\t(${JSON_STRING_PATTERN}): Model<(${JSON_STRING_PATTERN})> & \\{$`); +const MODEL_ID_PATTERN = new RegExp(`^\\t\\tid: (${JSON_STRING_PATTERN});$`); +const MODEL_PROVIDER_PATTERN = new RegExp(`^\\t\\tprovider: (${JSON_STRING_PATTERN});$`); + +function sha256(value: string): string { + return createHash("sha256").update(value).digest("hex"); +} + +function parseJsonString(value: string, description: string): string { + const parsed: unknown = JSON.parse(value); + if (typeof parsed !== "string") throw new Error(`${description} is not a string`); + return parsed; +} + +function sortedRecord(entries: Iterable): Record { + return Object.fromEntries(Array.from(entries).sort(([a], [b]) => (a < b ? -1 : a > b ? 1 : 0))); +} + +function sameStrings(a: readonly string[], b: readonly string[]): boolean { + return a.length === b.length && a.every((value, index) => value === b[index]); +} + +function describeSetDifference(expected: readonly string[], actual: readonly string[]): string { + const expectedSet = new Set(expected); + const actualSet = new Set(actual); + const missing = expected.filter((value) => !actualSet.has(value)); + const extra = actual.filter((value) => !expectedSet.has(value)); + return [missing.length > 0 ? `missing: ${missing.join(", ")}` : "", extra.length > 0 ? `extra: ${extra.join(", ")}` : ""] + .filter(Boolean) + .join("; "); +} + +function parseProviderStructure(path: string, providerId: string): Record { + const source = readFileSync(path, "utf8"); + const expectedImport = `import values from "./data/${providerId}.json" with { type: "json" };`; + if (!source.includes(expectedImport)) { + throw new Error(`${path} does not import ${providerId}.json`); + } + + const models = new Map(); + const lines = source.split("\n"); + for (let index = 0; index < lines.length; index++) { + const shapeMatch = MODEL_SHAPE_PATTERN.exec(lines[index]); + if (!shapeMatch) continue; + + const idMatch = MODEL_ID_PATTERN.exec(lines[index + 1] ?? ""); + const providerMatch = MODEL_PROVIDER_PATTERN.exec(lines[index + 2] ?? ""); + if (!idMatch || !providerMatch || lines[index + 3] !== "\t};") { + throw new Error(`${path}:${index + 1} has a malformed generated model declaration`); + } + + const key = parseJsonString(shapeMatch[1], `${path}:${index + 1} model key`); + const api = parseJsonString(shapeMatch[2], `${path}:${index + 1} model API`); + const id = parseJsonString(idMatch[1], `${path}:${index + 2} model ID`); + const provider = parseJsonString(providerMatch[1], `${path}:${index + 3} provider ID`); + if (id !== key) throw new Error(`${path}:${index + 1} declares key ${key} with ID ${id}`); + if (provider !== providerId) { + throw new Error(`${path}:${index + 1} declares provider ${provider} instead of ${providerId}`); + } + if (models.has(key)) throw new Error(`${path} declares model ${key} more than once`); + models.set(key, api); + index += 3; + } + + if (models.size === 0) throw new Error(`${path} contains no generated model declarations`); + return sortedRecord(models); +} + +export function readModelDataStructure(packageRoot: string): ModelDataStructure { + const providersDir = join(packageRoot, "src", "providers"); + const shardProviderIds = readdirSync(providersDir) + .filter((entry) => entry.endsWith(".models.ts")) + .map((entry) => entry.slice(0, -".models.ts".length)) + .sort(); + if (shardProviderIds.length === 0) throw new Error(`No generated provider shards found under ${providersDir}`); + + const aggregator = readFileSync(join(packageRoot, "src", "models.generated.ts"), "utf8"); + const importedProviderIds = Array.from( + aggregator.matchAll(/^import \{ [A-Z0-9_]+_MODELS \} from "\.\/providers\/([^"/]+)\.models\.ts";$/gm), + (match) => match[1], + ).sort(); + if (!sameStrings(shardProviderIds, importedProviderIds)) { + throw new Error( + `Generated model aggregator and provider shards do not match (${describeSetDifference(shardProviderIds, importedProviderIds)})`, + ); + } + + return sortedRecord( + shardProviderIds.map((providerId) => [ + providerId, + parseProviderStructure(join(providersDir, `${providerId}.models.ts`), providerId), + ] as const), + ); +} + +export function modelDataStructureHash(structure: ModelDataStructure): string { + return sha256(JSON.stringify(structure)); +} + +export function createModelDataManifest( + structure: ModelDataStructure, + fileContents: Readonly>, +): ModelDataManifest { + return { + schemaVersion: MODEL_DATA_SCHEMA_VERSION, + structureHash: modelDataStructureHash(structure), + files: sortedRecord(Object.entries(fileContents).map(([file, content]) => [file, sha256(content)] as const)), + }; +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +function readJsonObject(path: string, description: string, errors: string[]): Record | undefined { + let parsed: unknown; + try { + parsed = JSON.parse(readFileSync(path, "utf8")); + } catch (error) { + errors.push(`${description} is not valid JSON: ${error instanceof Error ? error.message : String(error)}`); + return undefined; + } + if (!isRecord(parsed)) { + errors.push(`${description} must contain a JSON object`); + return undefined; + } + return parsed; +} + +function validateModelValue( + value: unknown, + providerId: string, + modelId: string, + expectedApi: string, + errors: string[], +): void { + const label = `${providerId}/${modelId}`; + if (!isRecord(value)) { + errors.push(`${label} must be an object`); + return; + } + if (value.id !== modelId) errors.push(`${label} has id ${JSON.stringify(value.id)}, expected ${JSON.stringify(modelId)}`); + if (value.provider !== providerId) { + errors.push(`${label} has provider ${JSON.stringify(value.provider)}, expected ${JSON.stringify(providerId)}`); + } + if (value.api !== expectedApi) { + errors.push(`${label} has api ${JSON.stringify(value.api)}, expected ${JSON.stringify(expectedApi)}`); + } + if (typeof value.name !== "string" || value.name.length === 0) errors.push(`${label} has no model name`); + if (typeof value.baseUrl !== "string") errors.push(`${label} has no baseUrl string`); + if (typeof value.reasoning !== "boolean") errors.push(`${label} has no reasoning boolean`); + if ( + !Array.isArray(value.input) || + value.input.length === 0 || + value.input.some((entry) => entry !== "text" && entry !== "image") + ) { + errors.push(`${label} has invalid input modalities`); + } + if (typeof value.contextWindow !== "number" || !Number.isFinite(value.contextWindow) || value.contextWindow <= 0) { + errors.push(`${label} has invalid contextWindow`); + } + if (typeof value.maxTokens !== "number" || !Number.isFinite(value.maxTokens) || value.maxTokens <= 0) { + errors.push(`${label} has invalid maxTokens`); + } + if (!isRecord(value.cost)) { + errors.push(`${label} has invalid cost metadata`); + } else { + for (const field of ["input", "output", "cacheRead", "cacheWrite"] as const) { + const cost = value.cost[field]; + if (typeof cost !== "number" || !Number.isFinite(cost)) { + errors.push(`${label} has invalid cost.${field}`); + } + } + } +} + +function throwValidationErrors(errors: string[]): never { + const visible = errors.slice(0, 30); + const suffix = errors.length > visible.length ? `\n ... and ${errors.length - visible.length} more` : ""; + throw new Error(`Invalid generated model data:\n${visible.map((error) => ` - ${error}`).join("\n")}${suffix}`); +} + +export function validateModelDataDirectory(structure: ModelDataStructure, dataDir: string): void { + if (!existsSync(dataDir) || !statSync(dataDir).isDirectory()) { + throw new Error(`Generated model data directory does not exist: ${dataDir}`); + } + + const errors: string[] = []; + const expectedFiles = Object.keys(structure) + .map((providerId) => `${providerId}.json`) + .sort(); + const actualFiles = readdirSync(dataDir) + .filter((entry) => entry.endsWith(".json") && entry !== MODEL_DATA_MANIFEST_FILE) + .sort(); + if (!sameStrings(expectedFiles, actualFiles)) { + errors.push(`provider data files do not match the structural catalog (${describeSetDifference(expectedFiles, actualFiles)})`); + } + + const manifestPath = join(dataDir, MODEL_DATA_MANIFEST_FILE); + const manifest = readJsonObject(manifestPath, "model data manifest", errors); + if (manifest?.schemaVersion !== MODEL_DATA_SCHEMA_VERSION) { + errors.push( + `model data schema is ${JSON.stringify(manifest?.schemaVersion)}, expected ${MODEL_DATA_SCHEMA_VERSION}`, + ); + } + const expectedStructureHash = modelDataStructureHash(structure); + if (manifest?.structureHash !== expectedStructureHash) { + errors.push("model data generation stamp does not match the structural catalog"); + } + const manifestFiles = isRecord(manifest?.files) ? manifest.files : undefined; + if (!manifestFiles) errors.push("model data manifest has no file hashes"); + else { + const manifestFileNames = Object.keys(manifestFiles).sort(); + if (!sameStrings(expectedFiles, manifestFileNames)) { + errors.push(`manifest file hashes do not match provider data files (${describeSetDifference(expectedFiles, manifestFileNames)})`); + } + } + + for (const [providerId, expectedModels] of Object.entries(structure)) { + const filename = `${providerId}.json`; + const path = join(dataDir, filename); + if (!existsSync(path)) continue; + const content = readFileSync(path, "utf8"); + if (manifestFiles && manifestFiles[filename] !== sha256(content)) { + errors.push(`${filename} does not match its manifest hash`); + } + const values = readJsonObject(path, filename, errors); + if (!values) continue; + const expectedModelIds = Object.keys(expectedModels).sort(); + const actualModelIds = Object.keys(values).sort(); + if (!sameStrings(expectedModelIds, actualModelIds)) { + errors.push(`${filename} model IDs do not match the structural catalog (${describeSetDifference(expectedModelIds, actualModelIds)})`); + } + for (const [modelId, api] of Object.entries(expectedModels)) { + if (modelId in values) validateModelValue(values[modelId], providerId, modelId, api, errors); + } + } + + if (errors.length > 0) throwValidationErrors(errors); +} + +export function validateGeneratedModelData(packageRoot: string): void { + const structure = readModelDataStructure(packageRoot); + validateModelDataDirectory(structure, join(packageRoot, "src", "providers", "data")); +} diff --git a/packages/ai/test/image-model-data.test.ts b/packages/ai/test/image-model-data.test.ts new file mode 100644 index 000000000..21d468050 --- /dev/null +++ b/packages/ai/test/image-model-data.test.ts @@ -0,0 +1,47 @@ +import { describe, expect, it } from "vitest"; +import { parseOpenRouterImageModels } from "../scripts/generate-image-models.ts"; + +const validImageModel = { + id: "example/image-model", + name: "Example Image Model", + architecture: { + input_modalities: ["text", "image"], + output_modalities: ["image"], + }, + pricing: { + prompt: "0.000001", + completion: "0.000002", + }, +}; + +describe("OpenRouter image model parsing", () => { + it.each([{}, { data: [] }, { data: "invalid" }])("rejects a missing or empty strict catalog", (payload) => { + expect(() => parseOpenRouterImageModels(payload, true)).toThrow("missing or empty image model list"); + }); + + it("rejects a strict catalog with no usable image models", () => { + expect(() => + parseOpenRouterImageModels( + { + data: [ + { + ...validImageModel, + architecture: { input_modalities: ["text"], output_modalities: ["text"] }, + }, + ], + }, + true, + ), + ).toThrow("no usable image models"); + }); + + it("parses a non-empty image model catalog", () => { + expect(parseOpenRouterImageModels({ data: [validImageModel] }, true)).toEqual([ + expect.objectContaining({ + id: "example/image-model", + input: ["text", "image"], + output: ["image"], + }), + ]); + }); +}); diff --git a/packages/ai/test/model-data-validation.test.ts b/packages/ai/test/model-data-validation.test.ts new file mode 100644 index 000000000..1e7798d3f --- /dev/null +++ b/packages/ai/test/model-data-validation.test.ts @@ -0,0 +1,123 @@ +import { mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { afterEach, describe, expect, it } from "vitest"; +import { + createModelDataManifest, + MODEL_DATA_MANIFEST_FILE, + MODEL_DATA_SCHEMA_VERSION, + type ModelDataStructure, + readModelDataStructure, + validateModelDataDirectory, +} from "../scripts/model-data.ts"; + +const temporaryRoots: string[] = []; + +afterEach(() => { + for (const root of temporaryRoots.splice(0)) rmSync(root, { force: true, recursive: true }); +}); + +function createFixture(): { + dataDir: string; + packageRoot: string; + structure: ModelDataStructure; + values: Record; +} { + const packageRoot = mkdtempSync(join(tmpdir(), "pi-model-data-")); + temporaryRoots.push(packageRoot); + const providersDir = join(packageRoot, "src", "providers"); + const dataDir = join(providersDir, "data"); + mkdirSync(dataDir, { recursive: true }); + writeFileSync( + join(packageRoot, "src", "models.generated.ts"), + 'import { TEST_PROVIDER_MODELS } from "./providers/test-provider.models.ts";\n\nexport const MODELS = {\n\t"test-provider": TEST_PROVIDER_MODELS,\n} as const;\n', + ); + writeFileSync( + join(providersDir, "test-provider.models.ts"), + '// generated\n\nimport values from "./data/test-provider.json" with { type: "json" };\nimport type { Model } from "../types.ts";\n\nexport const TEST_PROVIDER_MODELS = values as {\n\t"model-a": Model<"openai-completions"> & {\n\t\tid: "model-a";\n\t\tprovider: "test-provider";\n\t};\n};\n', + ); + + const structure = readModelDataStructure(packageRoot); + const values: Record = { + "model-a": { + id: "model-a", + name: "Model A", + api: "openai-completions", + provider: "test-provider", + baseUrl: "https://example.test/v1", + reasoning: false, + input: ["text"], + cost: { input: 1, output: 2, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 1000, + maxTokens: 100, + }, + }; + writeFixtureData(dataDir, structure, values); + return { dataDir, packageRoot, structure, values }; +} + +function writeFixtureData( + dataDir: string, + structure: ModelDataStructure, + values: Record, + manifestSchemaVersion = MODEL_DATA_SCHEMA_VERSION, +): void { + const filename = "test-provider.json"; + const content = `${JSON.stringify(values)}\n`; + writeFileSync(join(dataDir, filename), content); + const manifest = createModelDataManifest(structure, { [filename]: content }); + manifest.schemaVersion = manifestSchemaVersion; + writeFileSync(join(dataDir, MODEL_DATA_MANIFEST_FILE), `${JSON.stringify(manifest)}\n`); +} + +describe("generated model data validation", () => { + it("validates complete data against generated structural catalogs", () => { + const { dataDir, structure } = createFixture(); + expect(() => validateModelDataDirectory(structure, dataDir)).not.toThrow(); + }); + + it("rejects a missing model data directory", () => { + const { dataDir, structure } = createFixture(); + rmSync(dataDir, { recursive: true }); + expect(() => validateModelDataDirectory(structure, dataDir)).toThrow("does not exist"); + }); + + it.each([ + ["id", "wrong-id", "has id"], + ["provider", "wrong-provider", "has provider"], + ["api", "anthropic-messages", "has api"], + ] as const)("rejects a wrong model %s", (field, value, expectedMessage) => { + const fixture = createFixture(); + const model = fixture.values["model-a"] as Record; + model[field] = value; + writeFixtureData(fixture.dataDir, fixture.structure, fixture.values); + expect(() => validateModelDataDirectory(fixture.structure, fixture.dataDir)).toThrow(expectedMessage); + }); + + it("rejects missing model IDs and stale file hashes", () => { + const fixture = createFixture(); + writeFileSync(join(fixture.dataDir, "test-provider.json"), "{}\n"); + expect(() => validateModelDataDirectory(fixture.structure, fixture.dataDir)).toThrow(/manifest hash|model IDs/); + }); + + it("rejects incompatible schema and generation stamps", () => { + const fixture = createFixture(); + writeFixtureData(fixture.dataDir, fixture.structure, fixture.values, MODEL_DATA_SCHEMA_VERSION + 1); + expect(() => validateModelDataDirectory(fixture.structure, fixture.dataDir)).toThrow("model data schema"); + + const manifestPath = join(fixture.dataDir, MODEL_DATA_MANIFEST_FILE); + const manifest = JSON.parse(readFileSync(manifestPath, "utf8")) as Record; + manifest.structureHash = "stale"; + writeFileSync(manifestPath, `${JSON.stringify(manifest)}\n`); + expect(() => validateModelDataDirectory(fixture.structure, fixture.dataDir)).toThrow("generation stamp"); + }); + + it("rejects missing provider shards referenced by the aggregator", () => { + const { packageRoot } = createFixture(); + writeFileSync( + join(packageRoot, "src", "models.generated.ts"), + 'import { TEST_PROVIDER_MODELS } from "./providers/test-provider.models.ts";\nimport { MISSING_MODELS } from "./providers/missing.models.ts";\n', + ); + expect(() => readModelDataStructure(packageRoot)).toThrow("aggregator and provider shards do not match"); + }); +}); diff --git a/scripts/check-browser-smoke.mjs b/scripts/check-browser-smoke.mjs index eeab27b15..d11b22d22 100644 --- a/scripts/check-browser-smoke.mjs +++ b/scripts/check-browser-smoke.mjs @@ -8,7 +8,7 @@ const agentTreeshakeOutputPath = join(tmpdir(), "pi-agent-treeshake-smoke.js"); const errorLogPath = join(tmpdir(), "pi-browser-smoke-errors.log"); const generatedCatalogDataDir = join(process.cwd(), "packages/ai/src/providers/data"); -// Fresh checkouts do not materialize provider JSON until npm run build. +// Fresh checkouts do not materialize provider JSON until model data is hydrated. const generatedCatalogDataPlugin = { name: "generated-model-catalog", setup(build) { diff --git a/scripts/local-release.mjs b/scripts/local-release.mjs index e06996d03..181379d17 100644 --- a/scripts/local-release.mjs +++ b/scripts/local-release.mjs @@ -209,9 +209,9 @@ const bunInstallDirectory = join(outDir, "bun-install"); const binaryDirectory = join(outDir, "bun"); mkdirSync(tarballDirectory, { recursive: true }); -if (!options.skipCheck || !options.skipTest) { - run("npm", ["--prefix", "packages/ai", "run", "generate-models"], { cwd: repoRoot }); -} +// Release artifacts always use a freshly generated, strictly validated catalog, +// including when checks or tests are explicitly skipped. +run("npm", ["run", "generate:models"], { cwd: repoRoot }); if (!options.skipCheck) { run("npm", ["run", "check"], { cwd: repoRoot }); @@ -223,7 +223,7 @@ if (!options.skipTest) { for (const pkg of packages) { run("npm", ["run", "clean"], { cwd: pkg.directory }); - run("npm", ["run", "build"], { cwd: pkg.directory }); + run("npm", ["run", pkg.directory === "packages/ai" ? "build:offline" : "build"], { cwd: pkg.directory }); } const tarballs = new Map(); diff --git a/scripts/release.mjs b/scripts/release.mjs index e8a3d7258..95cac6e69 100755 --- a/scripts/release.mjs +++ b/scripts/release.mjs @@ -167,6 +167,7 @@ console.log(); // 4. Regenerate release artifacts console.log("Regenerating release artifacts..."); run("npm run generate:models"); +run("npm run check:model-data"); run("npm run shrinkwrap:coding-agent"); run("npm run install-lock:coding-agent"); console.log(); From 31dc078bf8921ac6e6fbdd6f52cbdfad37f5080a Mon Sep 17 00:00:00 2001 From: David Brailovsky Date: Tue, 21 Jul 2026 06:07:44 +0000 Subject: [PATCH 30/62] update brace-expansion version --- package-lock.json | 6 +++--- packages/coding-agent/install-lock/package-lock.json | 6 +++--- packages/coding-agent/npm-shrinkwrap.json | 6 +++--- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/package-lock.json b/package-lock.json index d6a1480ec..1190eb579 100644 --- a/package-lock.json +++ b/package-lock.json @@ -2451,9 +2451,9 @@ "license": "MIT" }, "node_modules/brace-expansion": { - "version": "5.0.6", - "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.6.tgz", - "integrity": "sha512-kLpxurY4Z4r9sgMsyG0Z9uzsBlgiU/EFKhj/h91/8yHu0edo7XuixOIH3VcJ8kkxs6/jPzoI6U9Vj3WqbMQ94g==", + "version": "5.0.7", + "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.7.tgz", + "integrity": "sha512-7oFy703dxfY3/NLxC1fh2SUCQ0H9rmAY+5EpDVfXjUTTs+HEwR2nYaqLv+GWcTsumwxPfiz6CzCNkwXwBUwqCA==", "license": "MIT", "dependencies": { "balanced-match": "^4.0.2" diff --git a/packages/coding-agent/install-lock/package-lock.json b/packages/coding-agent/install-lock/package-lock.json index 31c29e590..f66fd9b37 100644 --- a/packages/coding-agent/install-lock/package-lock.json +++ b/packages/coding-agent/install-lock/package-lock.json @@ -1047,9 +1047,9 @@ "license": "MIT" }, "node_modules/brace-expansion": { - "version": "5.0.6", - "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.6.tgz", - "integrity": "sha512-kLpxurY4Z4r9sgMsyG0Z9uzsBlgiU/EFKhj/h91/8yHu0edo7XuixOIH3VcJ8kkxs6/jPzoI6U9Vj3WqbMQ94g==", + "version": "5.0.7", + "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.7.tgz", + "integrity": "sha512-7oFy703dxfY3/NLxC1fh2SUCQ0H9rmAY+5EpDVfXjUTTs+HEwR2nYaqLv+GWcTsumwxPfiz6CzCNkwXwBUwqCA==", "license": "MIT", "dependencies": { "balanced-match": "^4.0.2" diff --git a/packages/coding-agent/npm-shrinkwrap.json b/packages/coding-agent/npm-shrinkwrap.json index fa901aeb5..96c93de83 100644 --- a/packages/coding-agent/npm-shrinkwrap.json +++ b/packages/coding-agent/npm-shrinkwrap.json @@ -1037,9 +1037,9 @@ "license": "MIT" }, "node_modules/brace-expansion": { - "version": "5.0.6", - "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.6.tgz", - "integrity": "sha512-kLpxurY4Z4r9sgMsyG0Z9uzsBlgiU/EFKhj/h91/8yHu0edo7XuixOIH3VcJ8kkxs6/jPzoI6U9Vj3WqbMQ94g==", + "version": "5.0.7", + "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.7.tgz", + "integrity": "sha512-7oFy703dxfY3/NLxC1fh2SUCQ0H9rmAY+5EpDVfXjUTTs+HEwR2nYaqLv+GWcTsumwxPfiz6CzCNkwXwBUwqCA==", "license": "MIT", "dependencies": { "balanced-match": "^4.0.2" From bb019731fc23db8e8889fbbbddbafc015a1f3fc2 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Tue, 21 Jul 2026 07:33:18 +0000 Subject: [PATCH 31/62] chore: approve contributor zaycruz --- .github/APPROVED_CONTRIBUTORS | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/APPROVED_CONTRIBUTORS b/.github/APPROVED_CONTRIBUTORS index 947f19a75..c69fbdbaf 100644 --- a/.github/APPROVED_CONTRIBUTORS +++ b/.github/APPROVED_CONTRIBUTORS @@ -285,3 +285,5 @@ rsaryev pr QuintinShaw pr R-Taneja pr + +zaycruz pr From 890b3547af230d1ad75a4ec2a8989387fae2bd2b Mon Sep 17 00:00:00 2001 From: David Brailovsky Date: Tue, 21 Jul 2026 07:48:16 +0000 Subject: [PATCH 32/62] run generate-models fixes #6891 --- packages/ai/src/providers/opencode.models.ts | 4 ---- packages/ai/src/providers/openrouter.models.ts | 4 ---- 2 files changed, 8 deletions(-) diff --git a/packages/ai/src/providers/opencode.models.ts b/packages/ai/src/providers/opencode.models.ts index 4611ee2cb..6907cb861 100644 --- a/packages/ai/src/providers/opencode.models.ts +++ b/packages/ai/src/providers/opencode.models.ts @@ -173,10 +173,6 @@ export const OPENCODE_MODELS = values as { id: "grok-build-0.1"; provider: "opencode"; }; - "hy3-free": Model<"openai-completions"> & { - id: "hy3-free"; - provider: "opencode"; - }; "kimi-k2.5": Model<"openai-completions"> & { id: "kimi-k2.5"; provider: "opencode"; diff --git a/packages/ai/src/providers/openrouter.models.ts b/packages/ai/src/providers/openrouter.models.ts index 694e296a0..c2358136c 100644 --- a/packages/ai/src/providers/openrouter.models.ts +++ b/packages/ai/src/providers/openrouter.models.ts @@ -957,10 +957,6 @@ export const OPENROUTER_MODELS = values as { id: "tencent/hy3-preview"; provider: "openrouter"; }; - "tencent/hy3:free": Model<"openai-completions"> & { - id: "tencent/hy3:free"; - provider: "openrouter"; - }; "thedrummer/unslopnemo-12b": Model<"openai-completions"> & { id: "thedrummer/unslopnemo-12b"; provider: "openrouter"; From 54fad505b9d8cbc8922ff55d7e2938f70cbf6a3d Mon Sep 17 00:00:00 2001 From: Armin Ronacher Date: Tue, 21 Jul 2026 11:11:08 +0200 Subject: [PATCH 33/62] fix(coding-agent): prefer newer generated model catalogs --- packages/ai/src/models-store.ts | 2 + packages/ai/src/providers/all.ts | 5 + .../coding-agent/src/core/model-runtime.ts | 10 +- .../src/core/remote-catalog-provider.ts | 48 ++++++- .../test/remote-catalog-provider.test.ts | 132 +++++++++--------- 5 files changed, 125 insertions(+), 72 deletions(-) diff --git a/packages/ai/src/models-store.ts b/packages/ai/src/models-store.ts index 020edb1f2..3bb2b57e8 100644 --- a/packages/ai/src/models-store.ts +++ b/packages/ai/src/models-store.ts @@ -2,6 +2,8 @@ import type { Api, Model } from "./types.ts"; export interface ModelsStoreEntry { models: readonly Model[]; + /** Unix timestamp from the remote catalog's Last-Modified header. */ + lastModified?: number; /** Unix timestamp of the last completed remote check. */ checkedAt?: number; } diff --git a/packages/ai/src/providers/all.ts b/packages/ai/src/providers/all.ts index fd7bea58d..549880e33 100644 --- a/packages/ai/src/providers/all.ts +++ b/packages/ai/src/providers/all.ts @@ -67,6 +67,11 @@ export function getBuiltinProviders(): BuiltinProvider[] { return Object.keys(MODELS) as BuiltinProvider[]; } +/** URL of a generated provider catalog, used to compare its mtime with remote catalogs during development. */ +export function getBuiltinModelDataUrl(provider: BuiltinProvider): URL { + return new URL(`./data/${provider}.json`, import.meta.url); +} + export function getBuiltinModels( provider: TProvider, ): Model>[] { diff --git a/packages/coding-agent/src/core/model-runtime.ts b/packages/coding-agent/src/core/model-runtime.ts index 64501f731..2cd85b8d6 100644 --- a/packages/coding-agent/src/core/model-runtime.ts +++ b/packages/coding-agent/src/core/model-runtime.ts @@ -143,7 +143,15 @@ export class ModelRuntime implements Models { const providers = builtinProviderCatalog .builtinProviders() .map((provider) => - provider.id === "radius" ? provider : withRemoteCatalog(provider, options.catalogBaseUrl), + provider.id === "radius" + ? provider + : withRemoteCatalog( + provider, + options.catalogBaseUrl, + builtinProviderCatalog.getBuiltinModelDataUrl( + provider.id as builtinProviderCatalog.BuiltinProvider, + ), + ), ); const runtime = new ModelRuntime( credentials, diff --git a/packages/coding-agent/src/core/remote-catalog-provider.ts b/packages/coding-agent/src/core/remote-catalog-provider.ts index 6667916d8..b823cf7da 100644 --- a/packages/coding-agent/src/core/remote-catalog-provider.ts +++ b/packages/coding-agent/src/core/remote-catalog-provider.ts @@ -1,4 +1,5 @@ -import type { Api, Model, Provider } from "@earendil-works/pi-ai"; +import { stat } from "node:fs/promises"; +import type { Api, Model, ModelsStoreEntry, Provider } from "@earendil-works/pi-ai"; import { VERSION } from "../config.ts"; import { getPiUserAgent } from "../utils/pi-user-agent.ts"; @@ -29,8 +30,26 @@ function parseCatalog(providerId: string, value: unknown): Model[] { .map((model) => ({ ...model, provider: providerId })); } +function remoteModels( + entry: ModelsStoreEntry | undefined, + localLastModified: number | undefined, +): readonly Model[] { + if (!entry) return []; + if ( + localLastModified !== undefined && + (entry.lastModified === undefined || entry.lastModified <= localLastModified) + ) { + return []; + } + return entry.models; +} + /** Add a persisted pi.dev catalog overlay to a static built-in provider. */ -export function withRemoteCatalog(provider: Provider, catalogBaseUrl: string = DEFAULT_CATALOG_BASE_URL): Provider { +export function withRemoteCatalog( + provider: Provider, + catalogBaseUrl: string = DEFAULT_CATALOG_BASE_URL, + localCatalogUrl?: URL, +): Provider { let dynamicModels: readonly Model[] = []; let inflightRefresh: Promise | undefined; @@ -40,12 +59,21 @@ export function withRemoteCatalog(provider: Provider, catalogBaseUrl: string = D refreshModels: (context) => { inflightRefresh ??= (async () => { try { + const localLastModified = localCatalogUrl + ? await stat(localCatalogUrl).then( + (value) => value.mtimeMs, + () => undefined, + ) + : undefined; const stored = await context.store.read(); - if (stored) dynamicModels = stored.models.filter((model) => model.provider === provider.id); + dynamicModels = remoteModels(stored, localLastModified).filter( + (model) => model.provider === provider.id, + ); if (!context.allowNetwork || context.signal?.aborted) return; if ( !context.force && stored?.checkedAt !== undefined && + stored.lastModified !== undefined && Date.now() - stored.checkedAt < REMOTE_CATALOG_REFRESH_INTERVAL_MS ) { return; @@ -62,17 +90,23 @@ export function withRemoteCatalog(provider: Provider, catalogBaseUrl: string = D if (context.signal?.aborted) return; const checkedAt = Date.now(); if (response.status === 404 || response.status === 501) { - await context.store.write({ models: dynamicModels, checkedAt }); + await context.store.write({ ...(stored ?? { models: [] }), checkedAt, lastModified: 0 }); return; } if (!response.ok) { - await context.store.write({ models: dynamicModels, checkedAt }); + await context.store.write({ ...(stored ?? { models: [] }), checkedAt }); throw new Error(`Model catalog request failed for ${provider.id}: ${response.status}`); } const refreshed = parseCatalog(provider.id, await response.json()); + const lastModified = Date.parse(response.headers.get("last-modified") ?? ""); if (context.signal?.aborted) return; - dynamicModels = refreshed; - await context.store.write({ models: refreshed, checkedAt }); + const entry = { + models: refreshed, + checkedAt, + lastModified: Number.isNaN(lastModified) ? 0 : lastModified, + }; + dynamicModels = remoteModels(entry, localLastModified); + await context.store.write(entry); } finally { inflightRefresh = undefined; } diff --git a/packages/coding-agent/test/remote-catalog-provider.test.ts b/packages/coding-agent/test/remote-catalog-provider.test.ts index c8c5f8698..8b43a4f32 100644 --- a/packages/coding-agent/test/remote-catalog-provider.test.ts +++ b/packages/coding-agent/test/remote-catalog-provider.test.ts @@ -1,4 +1,11 @@ -import { createProvider, InMemoryModelsStore, type Model } from "@earendil-works/pi-ai"; +import { statSync } from "node:fs"; +import { + createProvider, + InMemoryModelsStore, + type Model, + type ModelsStoreEntry, + type ProviderModelsStore, +} from "@earendil-works/pi-ai"; import { afterEach, describe, expect, it, vi } from "vitest"; import { VERSION } from "../src/config.ts"; import { withRemoteCatalog } from "../src/core/remote-catalog-provider.ts"; @@ -18,6 +25,34 @@ function model(id: string): Model<"openai-completions"> { }; } +function testProvider(localCatalogUrl?: URL) { + return withRemoteCatalog( + createProvider({ + id: "test-provider", + auth: { apiKey: { name: "Test", resolve: async () => ({ auth: {} }) } }, + models: [model("static")], + api: { + stream: () => { + throw new Error("not used"); + }, + streamSimple: () => { + throw new Error("not used"); + }, + }, + }), + "https://pi.dev", + localCatalogUrl, + ); +} + +function scopedStore(store: InMemoryModelsStore): ProviderModelsStore { + return { + read: () => store.read("test-provider"), + write: (entry: ModelsStoreEntry) => store.write("test-provider", entry), + delete: () => store.delete("test-provider"), + }; +} + afterEach(() => vi.restoreAllMocks()); describe("remote catalog provider", () => { @@ -29,50 +64,12 @@ describe("remote catalog provider", () => { headers: { "content-type": "application/json" }, }), ); - const provider = withRemoteCatalog( - createProvider({ - id: "test-provider", - auth: { apiKey: { name: "Test", resolve: async () => ({ auth: {} }) } }, - models: [model("static")], - api: { - stream: () => { - throw new Error("not used"); - }, - streamSimple: () => { - throw new Error("not used"); - }, - }, - }), - ); + const provider = testProvider(); const store = new InMemoryModelsStore(); - await provider.refreshModels?.({ - credential: { type: "api_key" }, - store: { - read: () => store.read(provider.id), - write: (entry) => store.write(provider.id, entry), - delete: () => store.delete(provider.id), - }, - allowNetwork: true, - }); - await provider.refreshModels?.({ - credential: { type: "api_key" }, - store: { - read: () => store.read(provider.id), - write: (entry) => store.write(provider.id, entry), - delete: () => store.delete(provider.id), - }, - allowNetwork: true, - }); - await provider.refreshModels?.({ - credential: { type: "api_key" }, - store: { - read: () => store.read(provider.id), - write: (entry) => store.write(provider.id, entry), - delete: () => store.delete(provider.id), - }, - allowNetwork: true, - force: true, - }); + const refresh = { credential: { type: "api_key" } as const, store: scopedStore(store), allowNetwork: true }; + await provider.refreshModels?.(refresh); + await provider.refreshModels?.(refresh); + await provider.refreshModels?.({ ...refresh, force: true }); expect(provider.getModels().map((entry) => entry.id)).toEqual(["static", "dynamic"]); expect((await store.read(provider.id))?.models.map((entry) => entry.id)).toEqual(["dynamic"]); @@ -82,33 +79,40 @@ describe("remote catalog provider", () => { }); }); + it("prefers the newer of the generated and remote catalogs", async () => { + const localCatalogUrl = new URL(import.meta.url); + const localMtime = statSync(localCatalogUrl).mtimeMs; + const newerHeader = new Date(localMtime + 60_000).toUTCString(); + const responses = [ + new Response(JSON.stringify({ old: model("old") }), { + headers: { "last-modified": new Date(localMtime - 60_000).toUTCString() }, + }), + new Response(JSON.stringify({ newer: model("newer") }), { + headers: { "last-modified": newerHeader }, + }), + ]; + vi.spyOn(globalThis, "fetch").mockImplementation(async () => responses.shift() as Response); + const provider = testProvider(localCatalogUrl); + const store = new InMemoryModelsStore(); + const refresh = { credential: { type: "api_key" } as const, store: scopedStore(store), allowNetwork: true }; + + await provider.refreshModels?.(refresh); + expect(provider.getModels().map((entry) => entry.id)).toEqual(["static"]); + + await provider.refreshModels?.({ ...refresh, force: true }); + expect(provider.getModels().map((entry) => entry.id)).toEqual(["static", "newer"]); + expect(await store.read(provider.id)).toMatchObject({ lastModified: Date.parse(newerHeader) }); + }); + it("treats unimplemented pi.dev catalog routes as an unavailable overlay", async () => { vi.spyOn(globalThis, "fetch").mockResolvedValue(new Response("not implemented", { status: 501 })); - const provider = withRemoteCatalog( - createProvider({ - id: "test-provider", - auth: { apiKey: { name: "Test", resolve: async () => ({ auth: {} }) } }, - models: [model("static")], - api: { - stream: () => { - throw new Error("not used"); - }, - streamSimple: () => { - throw new Error("not used"); - }, - }, - }), - ); + const provider = testProvider(); const store = new InMemoryModelsStore(); await expect( provider.refreshModels?.({ credential: { type: "api_key" }, - store: { - read: () => store.read(provider.id), - write: (entry) => store.write(provider.id, entry), - delete: () => store.delete(provider.id), - }, + store: scopedStore(store), allowNetwork: true, }), ).resolves.toBeUndefined(); From 9e7582aa03e54f410fa9688197a3b64514e93400 Mon Sep 17 00:00:00 2001 From: Cristina Poncela Cubeiro <140309543+cristinaponcela@users.noreply.github.com> Date: Tue, 21 Jul 2026 11:36:31 +0200 Subject: [PATCH 34/62] feat: sqlite session storage (#6594) This PR: - Adds retainedTail to compaction entries in the new agent harness so we don't have to walk up the tree for the 2000 tokens before compaction, - Changes getPathToRoot to getPathToRootOrCompaction to only load until last compaction, as unnecessary to access all nodes where it is called, - Adds a SQLite storage backend, in a separate packages/session-backend-sqlite, with a migration system and schemas as per on-site discussions: sessions to match session header messages (except for metadata, which I couldn't understand what it's used for or where it gets written, so I omitted it), session_entries for shared entry types as columns plus payload as a json for what remains, session_sequences to represent the append-only, serialized nature of the jsonl files, branch_entries to attribute nodes to branches (relationship one-to-many), and session_materialized with the session info (see /session in TUI) to act as a "cache" or quick-access for costs, message count, token info, labels, session name, and model-thinking-level config (e.g. for fast resume). - This is compatible with the new agent harness Session abstraction. --- biome.json | 2 + package-lock.json | 96 ++-- package.json | 5 +- packages/agent/README.md | 4 + packages/agent/package.json | 3 +- packages/agent/src/harness/agent-harness.ts | 1 + .../src/harness/compaction/compaction.ts | 20 +- .../src/harness/session/jsonl-storage.ts | 70 ++- .../src/harness/session/memory-storage.ts | 62 ++- .../agent/src/harness/session/repo-utils.ts | 2 +- packages/agent/src/harness/session/session.ts | 42 +- packages/agent/src/harness/types.ts | 26 +- packages/agent/src/index.ts | 1 + .../agent/test/harness/compaction.test.ts | 37 +- packages/agent/test/harness/session.test.ts | 11 +- .../test/harness/sqlite-migrations.test.ts | 441 +++++++++++++++++ .../agent/test/harness/sqlite-node.test.ts | 21 + packages/agent/test/harness/storage.test.ts | 175 ++++++- packages/agent/vitest.harness.config.ts | 2 + packages/coding-agent/docs/session-format.md | 17 +- .../install-lock/package-lock.json | 21 +- packages/coding-agent/npm-shrinkwrap.json | 21 +- packages/storage/sqlite-node/README.md | 5 + packages/storage/sqlite-node/package.json | 37 ++ .../sqlite-node/scripts/prepare-dist.mjs | 35 ++ packages/storage/sqlite-node/src/index.ts | 97 ++++ .../storage/sqlite-node/src/sqlite/index.ts | 4 + .../sqlite-node/src/sqlite/migrations.ts | 50 ++ .../src/sqlite/migrations/001_initial.sql | 59 +++ .../storage/sqlite-node/src/sqlite/repo.ts | 192 ++++++++ .../src/sqlite/storage/branch-entries.ts | 57 +++ .../sqlite-node/src/sqlite/storage/index.ts | 449 ++++++++++++++++++ .../src/sqlite/storage/session-entries.ts | 217 +++++++++ .../sqlite/storage/session-materialized.ts | 368 ++++++++++++++ .../src/sqlite/storage/session-sequences.ts | 16 + .../src/sqlite/storage/sessions.ts | 40 ++ .../sqlite-node/src/sqlite/storage/shared.ts | 29 ++ .../storage/sqlite-node/src/sqlite/types.ts | 55 +++ .../storage/sqlite-node/tsconfig.build.json | 11 + tsconfig.json | 3 +- 40 files changed, 2659 insertions(+), 145 deletions(-) create mode 100644 packages/agent/test/harness/sqlite-migrations.test.ts create mode 100644 packages/agent/test/harness/sqlite-node.test.ts create mode 100644 packages/storage/sqlite-node/README.md create mode 100644 packages/storage/sqlite-node/package.json create mode 100644 packages/storage/sqlite-node/scripts/prepare-dist.mjs create mode 100644 packages/storage/sqlite-node/src/index.ts create mode 100644 packages/storage/sqlite-node/src/sqlite/index.ts create mode 100644 packages/storage/sqlite-node/src/sqlite/migrations.ts create mode 100644 packages/storage/sqlite-node/src/sqlite/migrations/001_initial.sql create mode 100644 packages/storage/sqlite-node/src/sqlite/repo.ts create mode 100644 packages/storage/sqlite-node/src/sqlite/storage/branch-entries.ts create mode 100644 packages/storage/sqlite-node/src/sqlite/storage/index.ts create mode 100644 packages/storage/sqlite-node/src/sqlite/storage/session-entries.ts create mode 100644 packages/storage/sqlite-node/src/sqlite/storage/session-materialized.ts create mode 100644 packages/storage/sqlite-node/src/sqlite/storage/session-sequences.ts create mode 100644 packages/storage/sqlite-node/src/sqlite/storage/sessions.ts create mode 100644 packages/storage/sqlite-node/src/sqlite/storage/shared.ts create mode 100644 packages/storage/sqlite-node/src/sqlite/types.ts create mode 100644 packages/storage/sqlite-node/tsconfig.build.json diff --git a/biome.json b/biome.json index b54514029..ef3cec421 100644 --- a/biome.json +++ b/biome.json @@ -27,6 +27,8 @@ "includes": [ "packages/*/src/**/*.ts", "packages/*/test/**/*.ts", + "packages/storage/*/src/**/*.ts", + "packages/storage/*/test/**/*.ts", "packages/coding-agent/examples/**/*.ts", "!**/node_modules/**/*", "!**/test-sessions.ts", diff --git a/package-lock.json b/package-lock.json index 1190eb579..2799f6d7c 100644 --- a/package-lock.json +++ b/package-lock.json @@ -9,6 +9,7 @@ "version": "0.0.3", "workspaces": [ "packages/*", + "packages/storage/*", "packages/coding-agent/examples/extensions/with-deps", "packages/coding-agent/examples/extensions/custom-provider-anthropic", "packages/coding-agent/examples/extensions/custom-provider-gitlab-duo", @@ -785,6 +786,10 @@ "resolved": "packages/agent", "link": true }, + "node_modules/@earendil-works/pi-agent-sqlite-node": { + "resolved": "packages/storage/sqlite-node", + "link": true + }, "node_modules/@earendil-works/pi-ai": { "resolved": "packages/ai", "link": true @@ -801,29 +806,6 @@ "resolved": "packages/tui", "link": true }, - "node_modules/@emnapi/core": { - "version": "1.10.0", - "resolved": "https://registry.npmjs.org/@emnapi/core/-/core-1.10.0.tgz", - "integrity": "sha512-yq6OkJ4p82CAfPl0u9mQebQHKPJkY7WrIuk205cTYnYe+k2Z8YBh11FrbRG/H6ihirqcacOgl2BIO8oyMQLeXw==", - "dev": true, - "license": "MIT", - "optional": true, - "dependencies": { - "@emnapi/wasi-threads": "1.2.1", - "tslib": "^2.4.0" - } - }, - "node_modules/@emnapi/runtime": { - "version": "1.10.0", - "resolved": "https://registry.npmjs.org/@emnapi/runtime/-/runtime-1.10.0.tgz", - "integrity": "sha512-ewvYlk86xUoGI0zQRNq/mC+16R1QeDlKQy21Ki3oSYXNgLb45GV1P6A0M+/s6nyCuNDqe5VpaY84BzXGwVbwFA==", - "dev": true, - "license": "MIT", - "optional": true, - "dependencies": { - "tslib": "^2.4.0" - } - }, "node_modules/@emnapi/wasi-threads": { "version": "1.2.1", "resolved": "https://registry.npmjs.org/@emnapi/wasi-threads/-/wasi-threads-1.2.1.tgz", @@ -1403,9 +1385,6 @@ "cpu": [ "arm64" ], - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -1422,9 +1401,6 @@ "cpu": [ "arm64" ], - "libc": [ - "musl" - ], "license": "MIT", "optional": true, "os": [ @@ -1441,9 +1417,6 @@ "cpu": [ "riscv64" ], - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -1460,9 +1433,6 @@ "cpu": [ "x64" ], - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -1479,9 +1449,6 @@ "cpu": [ "x64" ], - "libc": [ - "musl" - ], "license": "MIT", "optional": true, "os": [ @@ -1617,6 +1584,7 @@ "resolved": "https://registry.npmjs.org/@opentelemetry/api/-/api-1.9.0.tgz", "integrity": "sha512-3giAOQvZiH5F9bMlMiv8+GSPMeqg0dbaeo58/0SlA9sxSqZhnUtxzX9/2FzyhS9sWQf5S0GJE0AKBrFqjpeYcg==", "license": "Apache-2.0", + "peer": true, "engines": { "node": ">=8.0.0" } @@ -1796,9 +1764,6 @@ "arm64" ], "dev": true, - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -1816,9 +1781,6 @@ "arm64" ], "dev": true, - "libc": [ - "musl" - ], "license": "MIT", "optional": true, "os": [ @@ -1836,9 +1798,6 @@ "ppc64" ], "dev": true, - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -1856,9 +1815,6 @@ "s390x" ], "dev": true, - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -1876,9 +1832,6 @@ "x64" ], "dev": true, - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -1896,9 +1849,6 @@ "x64" ], "dev": true, - "libc": [ - "musl" - ], "license": "MIT", "optional": true, "os": [ @@ -2712,6 +2662,7 @@ "dev": true, "hasInstallScript": true, "license": "MIT", + "peer": true, "bin": { "esbuild": "bin/esbuild" }, @@ -3607,9 +3558,6 @@ "arm64" ], "dev": true, - "libc": [ - "glibc" - ], "license": "MPL-2.0", "optional": true, "os": [ @@ -3631,9 +3579,6 @@ "arm64" ], "dev": true, - "libc": [ - "musl" - ], "license": "MPL-2.0", "optional": true, "os": [ @@ -3655,9 +3600,6 @@ "x64" ], "dev": true, - "libc": [ - "glibc" - ], "license": "MPL-2.0", "optional": true, "os": [ @@ -3679,9 +3621,6 @@ "x64" ], "dev": true, - "libc": [ - "musl" - ], "license": "MPL-2.0", "optional": true, "os": [ @@ -4793,6 +4732,7 @@ "integrity": "sha512-QP88BAKvMam/3NxH6vj2o21R6MjxZUAd6nlwAS/pnGvN9IVLocLHxGYIzFhg6fUQ+5th6P4dv4eW9jX3DSIj7A==", "dev": true, "license": "MIT", + "peer": true, "engines": { "node": ">=12" }, @@ -4831,6 +4771,7 @@ "integrity": "sha512-TvncJykhxAzFCk0VQZKBTClall4Pm7qXDSodb6uxi8QFa8X8mT6ABjxxsQ2opDRYxG7AzcRWXaFtruz5HJKuWg==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "esbuild": "~0.28.0" }, @@ -4911,6 +4852,7 @@ "integrity": "sha512-h9bXPmJichP5fLmVQo3PyaGSDE2n3aPuomeAlVRm0JLmt4rY6zmPKd59HYI4LNW8oTK7tlTsuC7l/m7awx9Jcw==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "lightningcss": "^1.32.0", "picomatch": "^4.0.4", @@ -5100,6 +5042,7 @@ "resolved": "https://registry.npmjs.org/zod/-/zod-3.25.76.tgz", "integrity": "sha512-gzUt/qt81nXsFGKIFcC3YnfEAx5NkunCfnDlvuBSSFS02bcXu4Lmea0AFIUwbLWxWPx3d9p8S5QoaujKcNQxcQ==", "license": "MIT", + "peer": true, "funding": { "url": "https://github.com/sponsors/colinhacks" } @@ -5139,6 +5082,7 @@ "integrity": "sha512-GUUEShf+PBCGW2KaXwcIt3Yk+e3pkKwWKb9GSyM9WQVE+ep2jzmHdGsHzu4wgcZy5fN9FBdVzjpBQsYlpfpgLA==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "undici-types": "~7.16.0" } @@ -5149,6 +5093,7 @@ "integrity": "sha512-G9/lgqibheLVBDRuya45EbsEXTYcWoSG+TLg7i2axuzx0Eq62eXn+aWXyaVdV5vKvFSWd6ywcX8hA7la9Pvu8g==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "@bcoe/v8-coverage": "^1.0.2", "@vitest/utils": "4.1.9", @@ -5354,6 +5299,7 @@ "integrity": "sha512-nE3/LEyc0z87uHYLZebqCUOaJr2hdtuPp7BQ4BosVFnfltxgAvMG08NyrSGlPpOUWvR27c5flSmYFTNr78L9GQ==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "@vitest/expect": "4.1.9", "@vitest/mocker": "4.1.9", @@ -5500,6 +5446,7 @@ "integrity": "sha512-GUUEShf+PBCGW2KaXwcIt3Yk+e3pkKwWKb9GSyM9WQVE+ep2jzmHdGsHzu4wgcZy5fN9FBdVzjpBQsYlpfpgLA==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "undici-types": "~7.16.0" } @@ -5867,6 +5814,7 @@ "integrity": "sha512-GUUEShf+PBCGW2KaXwcIt3Yk+e3pkKwWKb9GSyM9WQVE+ep2jzmHdGsHzu4wgcZy5fN9FBdVzjpBQsYlpfpgLA==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "undici-types": "~7.16.0" } @@ -6152,6 +6100,18 @@ "node": ">=22.19.0" } }, + "packages/storage/sqlite-node": { + "name": "@earendil-works/pi-agent-sqlite-node", + "version": "0.80.10", + "license": "MIT", + "dependencies": { + "@earendil-works/pi-agent-core": "^0.80.10", + "@earendil-works/pi-ai": "^0.80.10" + }, + "engines": { + "node": ">=22.19.0" + } + }, "packages/tui": { "name": "@earendil-works/pi-tui", "version": "0.80.10", diff --git a/package.json b/package.json index fc3dd5dbd..8cdca8557 100644 --- a/package.json +++ b/package.json @@ -4,6 +4,7 @@ "type": "module", "workspaces": [ "packages/*", + "packages/storage/*", "packages/coding-agent/examples/extensions/with-deps", "packages/coding-agent/examples/extensions/custom-provider-anthropic", "packages/coding-agent/examples/extensions/custom-provider-gitlab-duo", @@ -12,8 +13,8 @@ ], "scripts": { "clean": "npm run clean --workspaces", - "build": "cd packages/tui && npm run build && cd ../ai && npm run build && cd ../agent && npm run build && cd ../coding-agent && npm run build && cd ../orchestrator && npm run build", - "build:offline": "cd packages/tui && npm run build && cd ../ai && npm run build:offline && cd ../agent && npm run build && cd ../coding-agent && npm run build && cd ../orchestrator && npm run build", + "build": "cd packages/tui && npm run build && cd ../ai && npm run build && cd ../agent && npm run build && cd ../storage/sqlite-node && npm run build && cd ../../coding-agent && npm run build && cd ../orchestrator && npm run build", + "build:offline": "cd packages/tui && npm run build && cd ../ai && npm run build:offline && cd ../agent && npm run build && cd ../storage/sqlite-node && npm run build && cd ../../coding-agent && npm run build && cd ../orchestrator && npm run build", "check": "biome check --write --error-on-warnings . && npm run check:pinned-deps && npm run check:ts-imports && npm run check:shrinkwrap && npm run check:install-lock:coding-agent && tsgo --noEmit && npm run check:browser-smoke", "check:browser-smoke": "node scripts/check-browser-smoke.mjs", "check:pinned-deps": "node scripts/check-pinned-deps.mjs", diff --git a/packages/agent/README.md b/packages/agent/README.md index 3c935f444..00dac621a 100644 --- a/packages/agent/README.md +++ b/packages/agent/README.md @@ -8,6 +8,10 @@ Stateful agent with tool execution and event streaming. Built on `@earendil-work npm install @earendil-works/pi-agent-core ``` +### SQLite session backends + +The SQLite session backend and the `node:sqlite` adapter live in a separate package, `@earendil-works/pi-agent-sqlite-node`, so the core package does not pull in runtime builtins or native SQLite dependencies by default. The backend accepts a runtime-specific SQLite factory, allowing other storage backends to ship as their own packages in the future. + ## Quick Start ```typescript diff --git a/packages/agent/package.json b/packages/agent/package.json index 7334e3a64..d130df50f 100644 --- a/packages/agent/package.json +++ b/packages/agent/package.json @@ -21,12 +21,11 @@ "README.md" ], "scripts": { - "clean": "shx rm -rf dist", "build": "tsgo -p tsconfig.build.json", "test": "vitest --run", "test:harness": "vitest --run --config vitest.harness.config.ts", "coverage:harness": "vitest --run --config vitest.harness.config.ts --coverage", - "prepublishOnly": "npm run clean && npm run build" + "prepublishOnly": "npm run build" }, "dependencies": { "@earendil-works/pi-ai": "^0.80.10", diff --git a/packages/agent/src/harness/agent-harness.ts b/packages/agent/src/harness/agent-harness.ts index 85af6b28c..5ec7d5742 100644 --- a/packages/agent/src/harness/agent-harness.ts +++ b/packages/agent/src/harness/agent-harness.ts @@ -730,6 +730,7 @@ export class AgentHarness< result.details, provided !== undefined, result.usage, + result.retainedTail, ); const entry = await this.session.getEntry(entryId); if (entry?.type === "compaction") { diff --git a/packages/agent/src/harness/compaction/compaction.ts b/packages/agent/src/harness/compaction/compaction.ts index 6c21a492f..0dc82c7cd 100644 --- a/packages/agent/src/harness/compaction/compaction.ts +++ b/packages/agent/src/harness/compaction/compaction.ts @@ -97,12 +97,14 @@ function getMessageFromEntryForCompaction(entry: SessionTreeEntry): AgentMessage export interface CompactionResult { /** Summary text that replaces compacted history in future context. */ summary: string; - /** Entry id where retained history starts. */ - firstKeptEntryId: string; + /** Entry id where retained history starts. Optional during Pi 2.0 transition. */ + firstKeptEntryId?: string; /** Estimated context tokens before compaction. */ tokensBefore: number; /** Usage from the LLM call(s) that generated this summary, if available. */ usage?: Usage; + /** Retained recent messages stored directly on the compaction entry. Optional during Pi 2.0 transition. */ + retainedTail?: AgentMessage[]; /** Optional implementation-specific details stored with the compaction entry. */ details?: T; } @@ -583,6 +585,8 @@ export interface CompactionPreparation { messagesToSummarize: AgentMessage[]; /** Prefix messages summarized separately when compaction splits a turn. */ turnPrefixMessages: AgentMessage[]; + /** Recent messages retained after compaction and stored on the compaction entry. */ + retainedTail: AgentMessage[]; /** Whether compaction splits a turn. */ isSplitTurn: boolean; /** Estimated context tokens before compaction. */ @@ -617,7 +621,9 @@ export function prepareCompaction( if (prevCompactionIndex >= 0) { const prevCompaction = pathEntries[prevCompactionIndex] as CompactionEntry; previousSummary = prevCompaction.summary; - const firstKeptEntryIndex = pathEntries.findIndex((entry) => entry.id === prevCompaction.firstKeptEntryId); + const firstKeptEntryIndex = prevCompaction.firstKeptEntryId + ? pathEntries.findIndex((entry) => entry.id === prevCompaction.firstKeptEntryId) + : -1; boundaryStart = firstKeptEntryIndex >= 0 ? firstKeptEntryIndex : prevCompactionIndex + 1; } const boundaryEnd = pathEntries.length; @@ -644,6 +650,11 @@ export function prepareCompaction( if (msg) turnPrefixMessages.push(msg); } } + const retainedTail: AgentMessage[] = []; + for (let i = cutPoint.firstKeptEntryIndex; i < boundaryEnd; i++) { + const msg = getMessageFromEntryForCompaction(pathEntries[i]); + if (msg) retainedTail.push(msg); + } const fileOps = extractFileOperations(messagesToSummarize, pathEntries, prevCompactionIndex); if (cutPoint.isSplitTurn) { for (const msg of turnPrefixMessages) { @@ -655,6 +666,7 @@ export function prepareCompaction( firstKeptEntryId, messagesToSummarize, turnPrefixMessages, + retainedTail, isSplitTurn: cutPoint.isSplitTurn, tokensBefore, previousSummary, @@ -693,6 +705,7 @@ export async function compact( firstKeptEntryId, messagesToSummarize, turnPrefixMessages, + retainedTail, isSplitTurn, tokensBefore, previousSummary, @@ -762,6 +775,7 @@ export async function compact( firstKeptEntryId, tokensBefore, usage: summaryUsage, + retainedTail, details: { readFiles, modifiedFiles } as CompactionDetails, }); } diff --git a/packages/agent/src/harness/session/jsonl-storage.ts b/packages/agent/src/harness/session/jsonl-storage.ts index d57398d53..fea97fbdc 100644 --- a/packages/agent/src/harness/session/jsonl-storage.ts +++ b/packages/agent/src/harness/session/jsonl-storage.ts @@ -1,5 +1,12 @@ import { uuidv7 } from "@earendil-works/pi-ai"; -import type { FileSystem, JsonlSessionMetadata, LeafEntry, SessionStorage, SessionTreeEntry } from "../types.ts"; +import type { + FileSystem, + JsonlSessionMetadata, + LeafEntry, + SessionEntryCursorOptions, + SessionStorage, + SessionTreeEntry, +} from "../types.ts"; import { SessionError, toError } from "../types.ts"; import { getFileSystemResultOrThrow } from "./repo-utils.ts"; @@ -293,13 +300,66 @@ export class JsonlSessionStorage implements SessionStorage return this.labelsById.get(id); } - async getPathToRoot(leafId: string | null): Promise { + async getSessionName(): Promise { + const entries = await this.findEntries("session_info"); + return entries[entries.length - 1]?.name?.trim() || undefined; + } + + async getSessionStats() { + let messageCount = 0; + let cachedTokens = 0; + let uncachedTokens = 0; + let totalTokens = 0; + let costTotal = 0; + for (const entry of this.entries) { + if (entry.type === "message") { + messageCount += 1; + } + const usage = + entry.type === "message" + ? entry.message.role === "assistant" + ? entry.message.usage + : undefined + : entry.type === "compaction" || entry.type === "branch_summary" + ? entry.usage + : undefined; + if ( + !usage || + typeof usage.input !== "number" || + typeof usage.output !== "number" || + typeof usage.cacheRead !== "number" || + typeof usage.cacheWrite !== "number" || + typeof usage.cost?.total !== "number" + ) { + continue; + } + cachedTokens += usage.cacheRead; + uncachedTokens += usage.input + usage.cacheWrite; + totalTokens += usage.input + usage.output + usage.cacheRead + usage.cacheWrite; + costTotal += usage.cost.total; + } + return { + messageCount, + cachedTokens, + uncachedTokens, + totalTokens, + costTotal, + }; + } + + async getPathToRootOrCompaction(leafId: string | null): Promise { if (leafId === null) return []; const path: SessionTreeEntry[] = []; + let stopAtEntryId: string | null = null; let current = this.byId.get(leafId); if (!current) throw new SessionError("not_found", `Entry ${leafId} not found`); while (current) { path.unshift(current); + if (stopAtEntryId !== null && current.id === stopAtEntryId) break; + if (current.type === "compaction") { + if (current.retainedTail) break; + stopAtEntryId = current.firstKeptEntryId ?? null; + } if (!current.parentId) break; const parent = this.byId.get(current.parentId); if (!parent) throw new SessionError("invalid_session", `Entry ${current.parentId} not found`); @@ -308,7 +368,9 @@ export class JsonlSessionStorage implements SessionStorage return path; } - async getEntries(): Promise { - return [...this.entries]; + async getEntries(options?: SessionEntryCursorOptions): Promise { + const start = options?.afterEntrySeq ?? 0; + const end = options?.limit === undefined ? undefined : start + options.limit; + return this.entries.slice(start, end); } } diff --git a/packages/agent/src/harness/session/memory-storage.ts b/packages/agent/src/harness/session/memory-storage.ts index 68ec76231..1bc24f04f 100644 --- a/packages/agent/src/harness/session/memory-storage.ts +++ b/packages/agent/src/harness/session/memory-storage.ts @@ -1,6 +1,7 @@ import { uuidv7 } from "@earendil-works/pi-ai"; import { type LeafEntry, + type SessionEntryCursorOptions, SessionError, type SessionMetadata, type SessionStorage, @@ -112,13 +113,66 @@ export class InMemorySessionStorage { + async getSessionName(): Promise { + const entries = await this.findEntries("session_info"); + return entries[entries.length - 1]?.name?.trim() || undefined; + } + + async getSessionStats() { + let messageCount = 0; + let cachedTokens = 0; + let uncachedTokens = 0; + let totalTokens = 0; + let costTotal = 0; + for (const entry of this.entries) { + if (entry.type === "message") { + messageCount += 1; + } + const usage = + entry.type === "message" + ? entry.message.role === "assistant" + ? entry.message.usage + : undefined + : entry.type === "compaction" || entry.type === "branch_summary" + ? entry.usage + : undefined; + if ( + !usage || + typeof usage.input !== "number" || + typeof usage.output !== "number" || + typeof usage.cacheRead !== "number" || + typeof usage.cacheWrite !== "number" || + typeof usage.cost?.total !== "number" + ) { + continue; + } + cachedTokens += usage.cacheRead; + uncachedTokens += usage.input + usage.cacheWrite; + totalTokens += usage.input + usage.output + usage.cacheRead + usage.cacheWrite; + costTotal += usage.cost.total; + } + return { + messageCount, + cachedTokens, + uncachedTokens, + totalTokens, + costTotal, + }; + } + + async getPathToRootOrCompaction(leafId: string | null): Promise { if (leafId === null) return []; const path: SessionTreeEntry[] = []; + let stopAtEntryId: string | null = null; let current = this.byId.get(leafId); if (!current) throw new SessionError("not_found", `Entry ${leafId} not found`); while (current) { path.unshift(current); + if (stopAtEntryId !== null && current.id === stopAtEntryId) break; + if (current.type === "compaction") { + if (current.retainedTail) break; + stopAtEntryId = current.firstKeptEntryId ?? null; + } if (!current.parentId) break; const parent = this.byId.get(current.parentId); if (!parent) throw new SessionError("invalid_session", `Entry ${current.parentId} not found`); @@ -127,7 +181,9 @@ export class InMemorySessionStorage { - return [...this.entries]; + async getEntries(options?: SessionEntryCursorOptions): Promise { + const start = options?.afterEntrySeq ?? 0; + const end = options?.limit === undefined ? undefined : start + options.limit; + return this.entries.slice(start, end); } } diff --git a/packages/agent/src/harness/session/repo-utils.ts b/packages/agent/src/harness/session/repo-utils.ts index dd2e0f9aa..fd6ee4f2a 100644 --- a/packages/agent/src/harness/session/repo-utils.ts +++ b/packages/agent/src/harness/session/repo-utils.ts @@ -47,5 +47,5 @@ export async function getEntriesToFork( } effectiveLeafId = target.parentId; } - return storage.getPathToRoot(effectiveLeafId); + return storage.getPathToRootOrCompaction(effectiveLeafId); } diff --git a/packages/agent/src/harness/session/session.ts b/packages/agent/src/harness/session/session.ts index bff41fdd7..13007bc65 100644 --- a/packages/agent/src/harness/session/session.ts +++ b/packages/agent/src/harness/session/session.ts @@ -11,8 +11,10 @@ import type { MessageEntry, ModelChangeEntry, SessionContext, + SessionEntryCursorOptions, SessionInfoEntry, SessionMetadata, + SessionStats, SessionStorage, SessionTreeEntry, ThinkingLevelChangeEntry, @@ -67,11 +69,19 @@ export function defaultContextEntryTransform(pathEntries: readonly SessionTreeEn const entries: SessionTreeEntry[] = [compaction]; const compactionIdx = pathEntries.findIndex((entry) => entry.type === "compaction" && entry.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) entries.push(entry); + if (compaction.retainedTail) { + for (let i = compactionIdx + 1; i < pathEntries.length; i++) { + entries.push(pathEntries[i]!); + } + return entries; + } + if (compaction.firstKeptEntryId) { + let foundFirstKept = false; + for (let i = 0; i < compactionIdx; i++) { + const entry = pathEntries[i]!; + if (entry.id === compaction.firstKeptEntryId) foundFirstKept = true; + if (foundFirstKept) entries.push(entry); + } } for (let i = compactionIdx + 1; i < pathEntries.length; i++) { entries.push(pathEntries[i]!); @@ -111,7 +121,10 @@ export function sessionEntryToContextMessages( ]; } if (entry.type === "compaction") { - return [createCompactionSummaryMessage(entry.summary, entry.tokensBefore, entry.timestamp)]; + return [ + createCompactionSummaryMessage(entry.summary, entry.tokensBefore, entry.timestamp), + ...(entry.retainedTail ?? []), + ]; } if (entry.type === "branch_summary" && entry.summary) { return [createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp)]; @@ -159,13 +172,13 @@ export class Session { return this.storage.getEntry(id); } - getEntries(): Promise { - return this.storage.getEntries(); + getEntries(options?: SessionEntryCursorOptions): Promise { + return this.storage.getEntries(options); } async getBranch(fromId?: string): Promise { const leafId = fromId ?? (await this.storage.getLeafId()); - return this.storage.getPathToRoot(leafId); + return this.storage.getPathToRootOrCompaction(leafId); } async buildContextEntries(options: SessionContextBuildOptions = {}): Promise { @@ -190,9 +203,12 @@ export class Session { return this.storage.getLabel(id); } + getSessionStats(): Promise { + return this.storage.getSessionStats(); + } + async getSessionName(): Promise { - const entries = await this.storage.findEntries("session_info"); - return entries[entries.length - 1]?.name?.trim() || undefined; + return this.storage.getSessionName(); } private async appendTypedEntry(entry: TEntry): Promise { @@ -243,11 +259,12 @@ export class Session { async appendCompaction( summary: string, - firstKeptEntryId: string, + firstKeptEntryId: string | undefined, tokensBefore: number, details?: T, fromHook?: boolean, usage?: Usage, + retainedTail?: AgentMessage[], ): Promise { return this.appendTypedEntry({ type: "compaction", @@ -257,6 +274,7 @@ export class Session { summary, firstKeptEntryId, tokensBefore, + retainedTail, details, usage, fromHook, diff --git a/packages/agent/src/harness/types.ts b/packages/agent/src/harness/types.ts index cdc7a5426..4a2f31d9b 100644 --- a/packages/agent/src/harness/types.ts +++ b/packages/agent/src/harness/types.ts @@ -370,8 +370,9 @@ export interface ActiveToolsChangeEntry extends SessionTreeEntryBase { export interface CompactionEntry extends SessionTreeEntryBase { type: "compaction"; summary: string; - firstKeptEntryId: string; + firstKeptEntryId?: string; tokensBefore: number; + retainedTail?: AgentMessage[]; details?: T; usage?: Usage; fromHook?: boolean; @@ -436,6 +437,14 @@ export interface SessionContext { activeToolNames: string[] | null; } +export interface SessionStats { + messageCount: number; + cachedTokens: number; + uncachedTokens: number; + totalTokens: number; + costTotal: number; +} + export interface SessionMetadata { id: string; createdAt: string; @@ -448,6 +457,11 @@ export interface JsonlSessionMetadata extends SessionMetadata { metadata?: Record; } +export interface SessionEntryCursorOptions { + afterEntrySeq?: number; + limit?: number; +} + export interface SessionStorage { getMetadata(): Promise; getLeafId(): Promise; @@ -460,8 +474,10 @@ export interface SessionStorage>>; getLabel(id: string): Promise; - getPathToRoot(leafId: string | null): Promise; - getEntries(): Promise; + getSessionName(): Promise; + getSessionStats(): Promise; + getPathToRootOrCompaction(leafId: string | null): Promise; + getEntries(options?: SessionEntryCursorOptions): Promise; } export type { Session } from "./session/session.ts"; @@ -753,10 +769,11 @@ export interface AbortResult { export interface CompactResult { summary: string; - firstKeptEntryId: string; + firstKeptEntryId?: string; tokensBefore: number; /** Usage from the LLM call(s) that generated this summary, if available. */ usage?: Usage; + retainedTail?: AgentMessage[]; details?: unknown; } @@ -776,6 +793,7 @@ export interface CompactionPreparation { firstKeptEntryId: string; messagesToSummarize: AgentMessage[]; turnPrefixMessages: AgentMessage[]; + retainedTail: AgentMessage[]; isSplitTurn: boolean; tokensBefore: number; previousSummary?: string; diff --git a/packages/agent/src/index.ts b/packages/agent/src/index.ts index b5ec7807f..b98023b04 100644 --- a/packages/agent/src/index.ts +++ b/packages/agent/src/index.ts @@ -1,4 +1,5 @@ // Core Agent +export { uuidv7 } from "@earendil-works/pi-ai"; export * from "./agent.ts"; // Loop functions export * from "./agent-loop.ts"; diff --git a/packages/agent/test/harness/compaction.test.ts b/packages/agent/test/harness/compaction.test.ts index c8160307d..e55be45a5 100644 --- a/packages/agent/test/harness/compaction.test.ts +++ b/packages/agent/test/harness/compaction.test.ts @@ -91,6 +91,7 @@ function createCompactionEntry( summary: string, firstKeptEntryId: string, parentId: string | null = null, + retainedTail?: AgentMessage[], ): CompactionEntry { return { type: "compaction", @@ -100,6 +101,7 @@ function createCompactionEntry( summary, firstKeptEntryId, tokensBefore: 1234, + retainedTail, }; } @@ -343,12 +345,38 @@ describe("harness compaction", () => { const a1 = createMessageEntry(createAssistantMessage("a"), u1.id); const u2 = createMessageEntry(createUserMessage("2"), a1.id); const a2 = createMessageEntry(createAssistantMessage("b"), u2.id); - const compaction = createCompactionEntry("Summary of 1,a,2,b", u2.id, a2.id); + const compaction = createCompactionEntry("Summary of 1,a,2,b", u2.id, a2.id, [ + createUserMessage("2"), + createAssistantMessage("b"), + ]); const u3 = createMessageEntry(createUserMessage("3"), compaction.id); const a3 = createMessageEntry(createAssistantMessage("c"), u3.id); const loaded = buildSessionContext([u1, a1, u2, a2, compaction, u3, a3]); expect(loaded.messages).toHaveLength(5); expect(loaded.messages[0]?.role).toBe("compactionSummary"); + expect(loaded.messages.map((message) => message.role)).toEqual([ + "compactionSummary", + "user", + "assistant", + "user", + "assistant", + ]); + }); + + it("falls back to firstKeptEntryId when a compaction has no retained tail", () => { + const u1 = createMessageEntry(createUserMessage("1")); + const a1 = createMessageEntry(createAssistantMessage("a"), u1.id); + const u2 = createMessageEntry(createUserMessage("2"), a1.id); + const a2 = createMessageEntry(createAssistantMessage("b"), u2.id); + const compaction = createCompactionEntry("Summary of 1,a,2,b", u2.id, a2.id); + const u3 = createMessageEntry(createUserMessage("3"), compaction.id); + const loaded = buildSessionContext([u1, a1, u2, a2, compaction, u3]); + expect(loaded.messages.map((message) => message.role)).toEqual([ + "compactionSummary", + "user", + "assistant", + "user", + ]); }); it("tracks model and thinking level changes in built context", () => { @@ -374,6 +402,7 @@ describe("harness compaction", () => { expect(preparation).toBeDefined(); expect(preparation?.previousSummary).toBe("First summary"); expect(preparation?.firstKeptEntryId).toBeTruthy(); + expect(preparation?.retainedTail.length).toBeGreaterThan(0); expect(preparation?.tokensBefore).toBe(estimateContextTokens(buildSessionContext(pathEntries).messages).tokens); }); @@ -566,6 +595,7 @@ describe("harness compaction", () => { firstKeptEntryId: "entry-keep", messagesToSummarize: messages, turnPrefixMessages: messages, + retainedTail: messages, isSplitTurn: true, tokensBefore: 600000, fileOps: { read: new Set(), written: new Set(), edited: new Set() }, @@ -583,6 +613,7 @@ describe("harness compaction", () => { firstKeptEntryId: "entry-keep", messagesToSummarize: messages, turnPrefixMessages: [], + retainedTail: messages, isSplitTurn: false, tokensBefore: 100, fileOps: { read: new Set(), written: new Set(), edited: new Set() }, @@ -619,6 +650,7 @@ describe("harness compaction", () => { turnPrefixMessages: messages, isSplitTurn: true, tokensBefore: 100, + retainedTail: messages, fileOps: { read: new Set(), written: new Set(), edited: new Set() }, settings: { enabled: true, reserveTokens: 2000, keepRecentTokens: 20 }, }; @@ -642,6 +674,7 @@ describe("harness compaction", () => { firstKeptEntryId: "entry-keep", messagesToSummarize: [], turnPrefixMessages: messages, + retainedTail: messages, isSplitTurn: true, tokensBefore: 100, fileOps: { read: new Set(), written: new Set(), edited: new Set() }, @@ -659,6 +692,7 @@ describe("harness compaction", () => { firstKeptEntryId: "entry-keep", messagesToSummarize: [], turnPrefixMessages: messages, + retainedTail: messages, isSplitTurn: true, tokensBefore: 100, fileOps: { read: new Set(), written: new Set(), edited: new Set() }, @@ -697,6 +731,7 @@ describe("harness compaction", () => { expect(result.summary.length).toBeGreaterThan(0); expect(result.firstKeptEntryId).toBeTruthy(); expect(result.usage?.totalTokens).toBeGreaterThan(0); + expect(result.retainedTail?.length).toBeGreaterThan(0); expect(result.details).toBeDefined(); }); }); diff --git a/packages/agent/test/harness/session.test.ts b/packages/agent/test/harness/session.test.ts index 26af52e2a..70adbb386 100644 --- a/packages/agent/test/harness/session.test.ts +++ b/packages/agent/test/harness/session.test.ts @@ -68,11 +68,20 @@ async function runSessionSuite( await session.appendMessage(createAssistantMessage("two")); const user2 = await session.appendMessage(createUserMessage("three")); await session.appendMessage(createAssistantMessage("four")); - await session.appendCompaction("summary", user2, 1234); + await session.appendCompaction("summary", user2, 1234, undefined, undefined, undefined, [ + createUserMessage("three"), + createAssistantMessage("four"), + ]); await session.appendMessage(createUserMessage("five")); const context = await session.buildContext(); expect(context.messages[0]?.role).toBe("compactionSummary"); expect(context.messages).toHaveLength(4); + expect(context.messages.map((message) => message.role)).toEqual([ + "compactionSummary", + "user", + "assistant", + "user", + ]); }); it("supports moving with branch summary entries in context", async () => { diff --git a/packages/agent/test/harness/sqlite-migrations.test.ts b/packages/agent/test/harness/sqlite-migrations.test.ts new file mode 100644 index 000000000..0cf74938d --- /dev/null +++ b/packages/agent/test/harness/sqlite-migrations.test.ts @@ -0,0 +1,441 @@ +import { mkdtempSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { describe, expect, it } from "vitest"; +import { + applyMigrations, + createNodeSqliteFactory, + type SqliteDatabase, + type SqliteDatabaseFactory, + type SqliteRunResult, + type SqliteSessionMetadata, + SqliteSessionRepo, + SqliteSessionStorage, + type SqliteStatement, +} from "../../../storage/sqlite-node/src/index.ts"; +import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts"; +import { createAssistantMessage, createUserMessage } from "./session-test-utils.ts"; + +function createTempDir(): string { + return mkdtempSync(join(tmpdir(), "pi-agent-sqlite-")); +} + +class ThrowingStatement implements SqliteStatement { + private readonly onRun: () => Promise; + + constructor(onRun: () => Promise) { + this.onRun = onRun; + } + + async run(..._params: unknown[]): Promise { + return this.onRun(); + } + + async get(..._params: unknown[]): Promise { + return undefined; + } + + async all(..._params: unknown[]): Promise { + return []; + } +} + +class CountingDatabase implements SqliteDatabase { + closeCount = 0; + private readonly statementFactory: (sql: string) => SqliteStatement; + + constructor(statementFactory: (sql: string) => SqliteStatement) { + this.statementFactory = statementFactory; + } + + async exec(_sql: string): Promise {} + + prepare(sql: string): SqliteStatement { + return this.statementFactory(sql); + } + + async transaction(fn: () => Promise): Promise { + return fn(); + } + + async close(): Promise { + this.closeCount += 1; + } +} + +describe("SQLite migrations", () => { + it("applies file-based migrations and records them", async () => { + const root = createTempDir(); + const databasePath = join(root, "sessions.sqlite"); + const env = new NodeExecutionEnv({ cwd: root }); + const sqlite = createNodeSqliteFactory(); + const repo = new SqliteSessionRepo({ env, sqlite, databasePath }); + await repo.create({ cwd: root, id: "session-1" }); + + const db = await sqlite.open(databasePath); + try { + const rows = await db.prepare("SELECT id FROM migrations ORDER BY id").all<{ id: string }>(); + expect(rows.map((row) => row.id)).toEqual(["001_initial.sql"]); + const tables = await db + .prepare("SELECT name, sql FROM sqlite_master WHERE type = 'table' ORDER BY name") + .all<{ name: string; sql: string | null }>(); + expect(tables.map((row) => row.name)).toEqual( + expect.arrayContaining([ + "migrations", + "sessions", + "session_entries", + "session_sequences", + "branch_entries", + "session_materialized", + "entry_materialized", + ]), + ); + const sessionColumns = await db.prepare("PRAGMA table_info(sessions)").all<{ name: string }>(); + expect(sessionColumns.map((column) => column.name)).toContain("active_leaf_id"); + for (const tableName of [ + "sessions", + "session_sequences", + "branch_entries", + "session_materialized", + "entry_materialized", + ]) { + const table = tables.find((row) => row.name === tableName); + expect(table?.sql).toContain("WITHOUT ROWID"); + } + } finally { + await db.close(); + } + }); + + it("persists session metadata through create, list, open, and fork", async () => { + const root = createTempDir(); + const databasePath = join(root, "sessions.sqlite"); + const env = new NodeExecutionEnv({ cwd: root }); + const repo = new SqliteSessionRepo({ env, sqlite: createNodeSqliteFactory(), databasePath }); + const source = await repo.create({ + cwd: root, + id: "session-1", + metadata: { profile: "reviewer" }, + }); + const sourceMetadata = await source.getMetadata(); + expect(sourceMetadata.metadata).toEqual({ profile: "reviewer" }); + expect((await repo.list({ cwd: root })).map((listed) => listed.metadata)).toEqual([{ profile: "reviewer" }]); + expect((await (await repo.open(sourceMetadata)).getMetadata()).metadata).toEqual({ profile: "reviewer" }); + const fork = await repo.fork(sourceMetadata, { cwd: root, id: "session-2" }); + expect((await fork.getMetadata()).metadata).toEqual({ profile: "reviewer" }); + const overridden = await repo.fork(sourceMetadata, { + cwd: root, + id: "session-3", + metadata: { profile: "writer" }, + }); + expect((await overridden.getMetadata()).metadata).toEqual({ profile: "writer" }); + }); + + it("materializes active leaf id in sessions transactionally", async () => { + const root = createTempDir(); + const databasePath = join(root, "sessions.sqlite"); + const env = new NodeExecutionEnv({ cwd: root }); + const sqlite = createNodeSqliteFactory(); + const repo = new SqliteSessionRepo({ env, sqlite, databasePath }); + const session = await repo.create({ cwd: root, id: "session-1" }); + const rootId = await session.appendMessage(createUserMessage("root")); + const childId = await session.appendMessage(createAssistantMessage("child")); + await session.getStorage().setLeafId(rootId); + + const db = await sqlite.open(databasePath); + try { + const row = await db + .prepare("SELECT active_leaf_id FROM sessions WHERE id = ?") + .get<{ active_leaf_id: string | null }>("session-1"); + expect(row?.active_leaf_id).toBe(rootId); + const latestBranchRow = await db + .prepare( + "SELECT branch_id, entry_id, entry_seq FROM branch_entries WHERE session_id = ? ORDER BY entry_seq DESC LIMIT 1", + ) + .get<{ branch_id: string; entry_id: string; entry_seq: number }>("session-1"); + const latestSessionEntry = await db + .prepare("SELECT id, type FROM session_entries WHERE session_id = ? ORDER BY entry_seq DESC LIMIT 1") + .get<{ id: string; type: string }>("session-1"); + expect(latestSessionEntry?.type).toBe("leaf"); + expect(latestBranchRow?.entry_id).toBe(latestSessionEntry?.id); + } finally { + await db.close(); + } + + const reopened = await repo.open(await session.getMetadata()); + expect(await reopened.getLeafId()).toBe(rootId); + expect(childId).not.toBe(rootId); + }); + + it("materializes a new branch when appending from a parent with an existing child", async () => { + const root = createTempDir(); + const databasePath = join(root, "sessions.sqlite"); + const env = new NodeExecutionEnv({ cwd: root }); + const sqlite = createNodeSqliteFactory(); + const repo = new SqliteSessionRepo({ env, sqlite, databasePath }); + const session = await repo.create({ cwd: root, id: "session-1" }); + const rootId = await session.appendMessage(createUserMessage("root")); + const firstChildId = await session.appendMessage(createAssistantMessage("first child")); + await session.getStorage().setLeafId(rootId); + const secondChildId = await session.appendMessage(createAssistantMessage("second child")); + + const db = await sqlite.open(databasePath); + try { + const branchRows = await db + .prepare( + "SELECT branch_id, entry_id, entry_seq FROM branch_entries WHERE session_id = ? ORDER BY branch_id, entry_seq", + ) + .all<{ branch_id: string; entry_id: string; entry_seq: number }>("session-1"); + const branchIds = [...new Set(branchRows.map((row) => row.branch_id))]; + expect(branchIds).toHaveLength(3); + expect(branchRows.filter((row) => row.entry_id === rootId)).toHaveLength(3); + expect(branchRows.filter((row) => row.entry_id === firstChildId)).toHaveLength(1); + expect(branchRows.filter((row) => row.entry_id === secondChildId)).toHaveLength(1); + } finally { + await db.close(); + } + }); + + it("reopens using branch materialization and session summary state", async () => { + const root = createTempDir(); + const databasePath = join(root, "sessions.sqlite"); + const env = new NodeExecutionEnv({ cwd: root }); + const repo = new SqliteSessionRepo({ env, sqlite: createNodeSqliteFactory(), databasePath }); + const session = await repo.create({ cwd: root, id: "session-1" }); + const rootId = await session.appendMessage(createUserMessage("root")); + await session.appendMessage(createAssistantMessage("first child")); + await session.appendSessionName(" Reopened Session "); + await session.getStorage().setLeafId(rootId); + await session.appendMessage(createAssistantMessage("branched child")); + + const reopened = await repo.open(await session.getMetadata()); + expect(await reopened.getSessionName()).toBe("Reopened Session"); + expect((await reopened.buildContext()).messages.map((message) => message.role)).toEqual(["user", "assistant"]); + expect((await reopened.buildContext()).messages.at(-1)).toMatchObject({ + content: [{ type: "text", text: "branched child" }], + }); + }); + + it("pages entries by entry_seq cursor", async () => { + const root = createTempDir(); + const databasePath = join(root, "sessions.sqlite"); + const env = new NodeExecutionEnv({ cwd: root }); + const repo = new SqliteSessionRepo({ env, sqlite: createNodeSqliteFactory(), databasePath }); + const session = await repo.create({ cwd: root, id: "session-1" }); + await session.appendMessage(createUserMessage("one")); + await session.appendMessage(createAssistantMessage("two")); + await session.appendMessage(createUserMessage("three")); + + expect((await session.getEntries({ limit: 2 })).map((entry) => entry.type)).toEqual(["message", "message"]); + expect((await session.getEntries({ afterEntrySeq: 2, limit: 2 })).map((entry) => entry.type)).toEqual([ + "message", + "message", + ]); + }); + + it("closes the database when create fails after openDatabase succeeds", async () => { + const root = createTempDir(); + const db = new CountingDatabase((sql) => { + if (sql.startsWith("INSERT INTO sessions")) { + return new ThrowingStatement(async () => { + throw new Error("insert failed"); + }); + } + return new ThrowingStatement(async () => ({ changes: 1 })); + }); + const sqlite: SqliteDatabaseFactory = { + open: async () => db, + }; + const env = new NodeExecutionEnv({ cwd: root }); + const repo = new SqliteSessionRepo({ env, sqlite, databasePath: join(root, "sessions.sqlite") }); + + await expect(repo.create({ cwd: root, id: "session-1" })).rejects.toThrow("insert failed"); + expect(db.closeCount).toBe(1); + }); + + it("closes the database when open fails after openDatabase succeeds", async () => { + const root = createTempDir(); + const db = new CountingDatabase((sql) => { + if (sql.includes("FROM sessions WHERE id = ?")) { + return new ThrowingStatement(async () => ({ changes: 0 })); + } + return new ThrowingStatement(async () => ({ changes: 1 })); + }); + const sqlite: SqliteDatabaseFactory = { + open: async () => db, + }; + const env = new NodeExecutionEnv({ cwd: root }); + const repo = new SqliteSessionRepo({ env, sqlite, databasePath: join(root, "sessions.sqlite") }); + const metadata: SqliteSessionMetadata = { + id: "missing", + createdAt: new Date().toISOString(), + cwd: root, + path: join(root, "sessions.sqlite"), + }; + writeFileSync(metadata.path, ""); + + await expect(repo.open(metadata)).rejects.toThrow("Session not found: missing"); + expect(db.closeCount).toBe(1); + }); + + it("closes the source storage after fork reads its entries", async () => { + const root = createTempDir(); + const databasePath = join(root, "sessions.sqlite"); + const env = new NodeExecutionEnv({ cwd: root }); + const repo = new SqliteSessionRepo({ env, sqlite: createNodeSqliteFactory(), databasePath }); + let cleanupCount = 0; + const sourceStorage = { + async getEntries() { + return []; + }, + async getPathToRootOrCompaction() { + return []; + }, + async cleanup() { + cleanupCount += 1; + }, + } as const; + const originalOpen = repo.open.bind(repo); + repo.open = async () => + ({ + getStorage() { + return sourceStorage; + }, + }) as never; + + try { + await repo.fork( + { + id: "session-1", + createdAt: new Date().toISOString(), + cwd: root, + path: databasePath, + }, + { cwd: root, id: "session-2" }, + ); + } finally { + repo.open = originalOpen; + } + + expect(cleanupCount).toBe(1); + }); + + it("restores in-memory state when appendEntry fails after mutating caches", async () => { + const root = createTempDir(); + const databasePath = join(root, "sessions.sqlite"); + const sqlite = createNodeSqliteFactory(); + const db = await sqlite.open(databasePath); + await applyMigrations(db); + const storage = await SqliteSessionStorage.create(db, databasePath, { + cwd: root, + sessionId: "session-1", + }); + const originalPrepare = db.prepare.bind(db); + db.prepare = (sql: string) => { + if (sql.startsWith("UPDATE sessions SET active_leaf_id = ?")) { + return new ThrowingStatement(async () => { + throw new Error("active leaf update failed"); + }); + } + return originalPrepare(sql); + }; + + await expect( + storage.appendEntry({ + type: "message", + id: "root", + parentId: null, + timestamp: new Date().toISOString(), + message: createUserMessage("root"), + }), + ).rejects.toMatchObject({ code: "storage" }); + expect(await storage.getLeafId()).toBeNull(); + expect(await storage.getEntry("root")).toBeUndefined(); + expect(await storage.getEntries()).toEqual([]); + await db.close(); + }); + + it("materializes session summary fields transactionally", async () => { + const root = createTempDir(); + const databasePath = join(root, "sessions.sqlite"); + const env = new NodeExecutionEnv({ cwd: root }); + const sqlite = createNodeSqliteFactory(); + const repo = new SqliteSessionRepo({ env, sqlite, databasePath }); + const session = await repo.create({ cwd: root, id: "session-1" }); + const userId = await session.appendMessage(createUserMessage("one")); + await session.appendThinkingLevelChange("high"); + await session.appendModelChange("anthropic", "claude-sonnet-4-5"); + const assistant = { + ...createAssistantMessage("two"), + provider: "anthropic", + model: "claude-sonnet-4-5", + usage: { + input: 100, + output: 25, + cacheRead: 40, + cacheWrite: 10, + totalTokens: 175, + cost: { input: 0.1, output: 0.2, cacheRead: 0.03, cacheWrite: 0.04, total: 0.37 }, + }, + }; + await session.appendMessage(assistant); + await session.appendCompaction("summary", userId, 200, undefined, false, { + input: 1, + output: 2, + cacheRead: 3, + cacheWrite: 4, + totalTokens: 10, + cost: { input: 0.01, output: 0.02, cacheRead: 0.03, cacheWrite: 0.04, total: 0.1 }, + }); + await session.moveTo(userId, { + summary: "branch summary", + usage: { + input: 5, + output: 6, + cacheRead: 7, + cacheWrite: 8, + totalTokens: 26, + cost: { input: 0.05, output: 0.06, cacheRead: 0.07, cacheWrite: 0.08, total: 0.26 }, + }, + }); + await session.appendSessionName(" My Session "); + await session.appendLabel(userId, "checkpoint"); + + const db = await sqlite.open(databasePath); + try { + const row = await db.prepare("SELECT session_id, payload FROM session_materialized WHERE session_id = ?").get<{ + session_id: string; + payload: string; + }>("session-1"); + expect(row).toBeDefined(); + expect(row?.session_id).toBe("session-1"); + expect(JSON.parse(row?.payload ?? "null")).toMatchObject({ + name: "My Session", + messageCount: 2, + cachedTokens: 50, + uncachedTokens: 128, + totalTokens: 211, + costTotal: 0.73, + currentModel: { provider: "anthropic", modelId: "claude-sonnet-4-5" }, + currentThinkingLevel: "high", + }); + const entryRows = await db + .prepare( + "SELECT session_id, entry_seq, type, payload FROM entry_materialized WHERE session_id = ? ORDER BY entry_seq, type", + ) + .all<{ + session_id: string; + entry_seq: number; + type: string; + payload: string; + }>("session-1"); + expect( + entryRows.some((entryRow) => entryRow.type === "label" && JSON.parse(entryRow.payload).targetId === userId), + ).toBe(true); + expect(entryRows.some((entryRow) => entryRow.type === "thinking")).toBe(false); + expect(entryRows.some((entryRow) => entryRow.type === "model")).toBe(false); + } finally { + await db.close(); + } + }); +}); diff --git a/packages/agent/test/harness/sqlite-node.test.ts b/packages/agent/test/harness/sqlite-node.test.ts new file mode 100644 index 000000000..a94d37858 --- /dev/null +++ b/packages/agent/test/harness/sqlite-node.test.ts @@ -0,0 +1,21 @@ +import { join } from "node:path"; +import { describe, expect, it } from "vitest"; +import { createNodeSqliteFactory } from "../../../storage/sqlite-node/src/index.ts"; +import { createTempDir } from "./session-test-utils.ts"; + +describe("sqlite-node adapter", () => { + it("supports node:sqlite-style named parameters", async () => { + const root = createTempDir(); + const databasePath = join(root, "adapter.sqlite"); + const sqlite = createNodeSqliteFactory(); + const db = await sqlite.open(databasePath); + try { + await db.exec("CREATE TABLE items (id INTEGER PRIMARY KEY, text TEXT NOT NULL)"); + await db.prepare("INSERT INTO items (id, text) VALUES ($id, $text)").run({ $id: 1, $text: "hello" }); + const row = await db.prepare("SELECT text FROM items WHERE id = $id").get<{ text: string }>({ $id: 1 }); + expect(row).toEqual({ text: "hello" }); + } finally { + await db.close(); + } + }); +}); diff --git a/packages/agent/test/harness/storage.test.ts b/packages/agent/test/harness/storage.test.ts index 07f07df26..50f4fc59a 100644 --- a/packages/agent/test/harness/storage.test.ts +++ b/packages/agent/test/harness/storage.test.ts @@ -4,7 +4,13 @@ import { describe, expect, it } from "vitest"; import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts"; import { JsonlSessionStorage, loadJsonlSessionMetadata } from "../../src/harness/session/jsonl-storage.ts"; import { InMemorySessionStorage } from "../../src/harness/session/memory-storage.ts"; -import { type MessageEntry, ok, type SessionMetadata } from "../../src/harness/types.ts"; +import { + type BranchSummaryEntry, + type CompactionEntry, + type MessageEntry, + ok, + type SessionMetadata, +} from "../../src/harness/types.ts"; import { createAssistantMessage, createTempDir, createUserMessage } from "./session-test-utils.ts"; describe("InMemorySessionStorage", () => { @@ -80,7 +86,74 @@ describe("InMemorySessionStorage", () => { expect(await storage.getLabel("entry-1")).toBeUndefined(); }); - it("walks paths to root", async () => { + it("includes summary-entry usage in session stats", async () => { + const assistant: MessageEntry = { + type: "message", + id: "assistant", + parentId: null, + timestamp: "2026-01-01T00:00:00.000Z", + message: { + role: "assistant", + content: [{ type: "text", text: "reply" }], + api: "anthropic-messages", + provider: "anthropic", + model: "claude-sonnet-4-5", + usage: { + input: 10, + output: 20, + cacheRead: 30, + cacheWrite: 40, + totalTokens: 100, + cost: { input: 0.1, output: 0.2, cacheRead: 0.3, cacheWrite: 0.4, total: 1 }, + }, + stopReason: "stop", + timestamp: 0, + }, + }; + const compaction: CompactionEntry = { + type: "compaction", + id: "compaction", + parentId: "assistant", + timestamp: "2026-01-01T00:00:01.000Z", + summary: "summary", + firstKeptEntryId: "assistant", + tokensBefore: 1234, + usage: { + input: 1, + output: 2, + cacheRead: 3, + cacheWrite: 4, + totalTokens: 10, + cost: { input: 0.01, output: 0.02, cacheRead: 0.03, cacheWrite: 0.04, total: 0.1 }, + }, + }; + const branchSummary: BranchSummaryEntry = { + type: "branch_summary", + id: "branch-summary", + parentId: "compaction", + timestamp: "2026-01-01T00:00:02.000Z", + fromId: "assistant", + summary: "branch", + usage: { + input: 5, + output: 6, + cacheRead: 7, + cacheWrite: 8, + totalTokens: 26, + cost: { input: 0.05, output: 0.06, cacheRead: 0.07, cacheWrite: 0.08, total: 0.26 }, + }, + }; + const storage = new InMemorySessionStorage({ entries: [assistant, compaction, branchSummary] }); + expect(await storage.getSessionStats()).toEqual({ + messageCount: 1, + cachedTokens: 40, + uncachedTokens: 68, + totalTokens: 136, + costTotal: 1.36, + }); + }); + + it("walks paths to root or retained-tail compaction", async () => { const root: MessageEntry = { type: "message", id: "root", @@ -94,9 +167,29 @@ describe("InMemorySessionStorage", () => { parentId: "root", message: createAssistantMessage("child"), }; - const storage = new InMemorySessionStorage({ entries: [root, child] }); - expect((await storage.getPathToRoot("child")).map((entry) => entry.id)).toEqual(["root", "child"]); - expect(await storage.getPathToRoot(null)).toEqual([]); + const compaction: CompactionEntry = { + type: "compaction", + id: "compaction", + parentId: "child", + timestamp: "2026-01-01T00:00:01.000Z", + summary: "summary", + firstKeptEntryId: "child", + tokensBefore: 1234, + retainedTail: [createAssistantMessage("child")], + }; + const afterCompaction: MessageEntry = { + ...root, + id: "after-compaction", + parentId: "compaction", + message: createUserMessage("after"), + }; + const storage = new InMemorySessionStorage({ entries: [root, child, compaction, afterCompaction] }); + expect((await storage.getPathToRootOrCompaction("child")).map((entry) => entry.id)).toEqual(["root", "child"]); + expect((await storage.getPathToRootOrCompaction("after-compaction")).map((entry) => entry.id)).toEqual([ + "compaction", + "after-compaction", + ]); + expect(await storage.getPathToRootOrCompaction(null)).toEqual([]); }); }); @@ -255,7 +348,7 @@ describe("JsonlSessionStorage", () => { const reloaded = await JsonlSessionStorage.open(env, filePath); expect(await reloaded.getLeafId()).toBe("root"); expect((await reloaded.getEntries()).at(-1)).toMatchObject({ type: "leaf", targetId: "root" }); - expect((await loaded.getPathToRoot("child")).map((entry) => entry.id)).toEqual(["root", "child"]); + expect((await loaded.getPathToRootOrCompaction("child")).map((entry) => entry.id)).toEqual(["root", "child"]); }); it("finds entries by type", async () => { @@ -309,6 +402,76 @@ describe("JsonlSessionStorage", () => { expect(await loaded.getLabel("entry-1")).toBeUndefined(); }); + it("includes summary-entry usage in session stats", async () => { + const dir = createTempDir(); + const env = new NodeExecutionEnv({ cwd: dir }); + const filePath = join(dir, "session.jsonl"); + const storage = await JsonlSessionStorage.create(env, filePath, { cwd: dir, sessionId: "session-1" }); + await storage.appendEntry({ + type: "message", + id: "assistant", + parentId: null, + timestamp: "2026-01-01T00:00:00.000Z", + message: { + role: "assistant", + content: [{ type: "text", text: "reply" }], + api: "anthropic-messages", + provider: "anthropic", + model: "claude-sonnet-4-5", + usage: { + input: 10, + output: 20, + cacheRead: 30, + cacheWrite: 40, + totalTokens: 100, + cost: { input: 0.1, output: 0.2, cacheRead: 0.3, cacheWrite: 0.4, total: 1 }, + }, + stopReason: "stop", + timestamp: 0, + }, + }); + await storage.appendEntry({ + type: "compaction", + id: "compaction", + parentId: "assistant", + timestamp: "2026-01-01T00:00:01.000Z", + summary: "summary", + firstKeptEntryId: "assistant", + tokensBefore: 1234, + usage: { + input: 1, + output: 2, + cacheRead: 3, + cacheWrite: 4, + totalTokens: 10, + cost: { input: 0.01, output: 0.02, cacheRead: 0.03, cacheWrite: 0.04, total: 0.1 }, + }, + }); + await storage.appendEntry({ + type: "branch_summary", + id: "branch-summary", + parentId: "compaction", + timestamp: "2026-01-01T00:00:02.000Z", + fromId: "assistant", + summary: "branch", + usage: { + input: 5, + output: 6, + cacheRead: 7, + cacheWrite: 8, + totalTokens: 26, + cost: { input: 0.05, output: 0.06, cacheRead: 0.07, cacheWrite: 0.08, total: 0.26 }, + }, + }); + expect(await storage.getSessionStats()).toEqual({ + messageCount: 1, + cachedTokens: 40, + uncachedTokens: 68, + totalTokens: 136, + costTotal: 1.36, + }); + }); + it("reads session metadata through the line-reading filesystem operation", async () => { const dir = createTempDir(); const filePath = join(dir, "session.jsonl"); diff --git a/packages/agent/vitest.harness.config.ts b/packages/agent/vitest.harness.config.ts index 91c0d471b..8045bb1a6 100644 --- a/packages/agent/vitest.harness.config.ts +++ b/packages/agent/vitest.harness.config.ts @@ -3,6 +3,7 @@ import { defineConfig } from "vitest/config"; const aiSrcIndex = fileURLToPath(new URL("../ai/src/index.ts", import.meta.url)); const aiSrcCompat = fileURLToPath(new URL("../ai/src/compat.ts", import.meta.url)); +const agentSrcIndex = fileURLToPath(new URL("../agent/src/index.ts", import.meta.url)); export default defineConfig({ test: { @@ -21,6 +22,7 @@ export default defineConfig({ }, resolve: { alias: [ + { find: /^@earendil-works\/pi-agent-core$/, replacement: agentSrcIndex }, { find: /^@earendil-works\/pi-ai$/, replacement: aiSrcIndex }, { find: /^@earendil-works\/pi-ai\/compat$/, replacement: aiSrcCompat }, ], diff --git a/packages/coding-agent/docs/session-format.md b/packages/coding-agent/docs/session-format.md index b27b8b701..e2c49f1d0 100644 --- a/packages/coding-agent/docs/session-format.md +++ b/packages/coding-agent/docs/session-format.md @@ -232,10 +232,18 @@ Created when context is compacted. Stores a summary of earlier messages. {"type":"compaction","id":"f6g7h8i9","parentId":"e5f6g7h8","timestamp":"2024-12-03T14:10:00.000Z","summary":"User discussed X, Y, Z...","firstKeptEntryId":"c3d4e5f6","tokensBefore":50000} ``` +Newer harness-generated compactions embed the retained post-compaction context directly on the entry, instead of `firstKeptEntryId`: + +```json +{"type":"compaction","id":"f6g7h8i9","parentId":"e5f6g7h8","timestamp":"2024-12-03T14:10:00.000Z","summary":"User discussed X, Y, Z...","tokensBefore":50000,"retainedTail":[{"role":"user","content":"latest request"},{"role":"assistant","content":[{"type":"text","text":"latest reply"}],"provider":"anthropic","model":"claude-sonnet-4-5","usage":{...},"stopReason":"stop"}]} +``` + Optional fields: - `usage`: LLM usage from generating the summary; included in session token and cost totals +- `retainedTail`: Materialized `AgentMessage[]` kept after compaction. This is optional only for backward compatibility with older sessions. Newer harness-generated compactions include it so we can rebuild context from this checkpoint without walking older entries before the compaction entry. - `details`: Implementation-specific data (e.g., `{ readFiles: string[], modifiedFiles: string[] }` for default, or custom data for extensions) - `fromHook`: `true` if generated by an extension, `false`/`undefined` if pi-generated (legacy field name) +- `firstKeptEntryId`: for compatibility with old entry format. ### BranchSummaryEntry @@ -314,8 +322,9 @@ Entries form a tree: 1. Collects all entries on the path 2. If a `CompactionEntry` is on the path: - Includes the compaction entry first - - Then entries from `firstKeptEntryId` to compaction - - Then entries after compaction + - If `retainedTail` is present, it acts as a self-contained checkpoint and entries after the compaction are included + - Otherwise entries from `firstKeptEntryId` to the compaction are included + - Then entries after compaction are included 3. Preserves non-message entries in the selected range so interactive mode can render them `buildSessionContext()` builds on that entry list to produce the message list for the LLM: @@ -323,11 +332,13 @@ Entries form a tree: 1. Extracts current model and thinking level settings from the full path 2. Converts selected entries to messages: - `message` -> stored `AgentMessage` - - `compaction` -> `compactionSummary` + - `compaction` -> `compactionSummary` plus `retainedTail` when present - `branch_summary` -> `branchSummary` - `custom_message` -> `CustomMessage` - `custom` -> no context message +This makes newer compactions act like self-contained checkpoints. `retainedTail` is optional only so older sessions that only store `firstKeptEntryId` continue to load correctly. + ## Parsing Example ```typescript diff --git a/packages/coding-agent/install-lock/package-lock.json b/packages/coding-agent/install-lock/package-lock.json index f66fd9b37..5bcebab80 100644 --- a/packages/coding-agent/install-lock/package-lock.json +++ b/packages/coding-agent/install-lock/package-lock.json @@ -638,9 +638,6 @@ "cpu": [ "arm64" ], - "libc": [ - "glibc" - ], "optional": true }, "node_modules/@mariozechner/clipboard-linux-arm64-musl": { @@ -657,9 +654,6 @@ "cpu": [ "arm64" ], - "libc": [ - "musl" - ], "optional": true }, "node_modules/@mariozechner/clipboard-linux-riscv64-gnu": { @@ -676,9 +670,6 @@ "cpu": [ "riscv64" ], - "libc": [ - "glibc" - ], "optional": true }, "node_modules/@mariozechner/clipboard-linux-x64-gnu": { @@ -695,9 +686,6 @@ "cpu": [ "x64" ], - "libc": [ - "glibc" - ], "optional": true }, "node_modules/@mariozechner/clipboard-linux-x64-musl": { @@ -714,9 +702,6 @@ "cpu": [ "x64" ], - "libc": [ - "musl" - ], "optional": true }, "node_modules/@mariozechner/clipboard-win32-arm64-msvc": { @@ -790,7 +775,8 @@ "license": "Apache-2.0", "engines": { "node": ">=8.0.0" - } + }, + "peer": true }, "node_modules/@opentelemetry/semantic-conventions": { "version": "1.41.1", @@ -1834,7 +1820,8 @@ "license": "MIT", "funding": { "url": "https://github.com/sponsors/colinhacks" - } + }, + "peer": true }, "node_modules/zod-to-json-schema": { "version": "3.25.2", diff --git a/packages/coding-agent/npm-shrinkwrap.json b/packages/coding-agent/npm-shrinkwrap.json index 96c93de83..af3f2726e 100644 --- a/packages/coding-agent/npm-shrinkwrap.json +++ b/packages/coding-agent/npm-shrinkwrap.json @@ -628,9 +628,6 @@ "cpu": [ "arm64" ], - "libc": [ - "glibc" - ], "optional": true }, "node_modules/@mariozechner/clipboard-linux-arm64-musl": { @@ -647,9 +644,6 @@ "cpu": [ "arm64" ], - "libc": [ - "musl" - ], "optional": true }, "node_modules/@mariozechner/clipboard-linux-riscv64-gnu": { @@ -666,9 +660,6 @@ "cpu": [ "riscv64" ], - "libc": [ - "glibc" - ], "optional": true }, "node_modules/@mariozechner/clipboard-linux-x64-gnu": { @@ -685,9 +676,6 @@ "cpu": [ "x64" ], - "libc": [ - "glibc" - ], "optional": true }, "node_modules/@mariozechner/clipboard-linux-x64-musl": { @@ -704,9 +692,6 @@ "cpu": [ "x64" ], - "libc": [ - "musl" - ], "optional": true }, "node_modules/@mariozechner/clipboard-win32-arm64-msvc": { @@ -780,7 +765,8 @@ "license": "Apache-2.0", "engines": { "node": ">=8.0.0" - } + }, + "peer": true }, "node_modules/@opentelemetry/semantic-conventions": { "version": "1.41.1", @@ -1824,7 +1810,8 @@ "license": "MIT", "funding": { "url": "https://github.com/sponsors/colinhacks" - } + }, + "peer": true }, "node_modules/zod-to-json-schema": { "version": "3.25.2", diff --git a/packages/storage/sqlite-node/README.md b/packages/storage/sqlite-node/README.md new file mode 100644 index 000000000..849157388 --- /dev/null +++ b/packages/storage/sqlite-node/README.md @@ -0,0 +1,5 @@ +# @earendil-works/pi-agent-sqlite-node + +Node sqlite storage backend for `@earendil-works/pi-agent-core` sessions. Provides the +`node:sqlite` adapter (`SqliteDatabase` implementation) and the SQLite session +repo/storage implementation (`SqliteSessionRepo`, migrations, materialized views). diff --git a/packages/storage/sqlite-node/package.json b/packages/storage/sqlite-node/package.json new file mode 100644 index 000000000..66069bd12 --- /dev/null +++ b/packages/storage/sqlite-node/package.json @@ -0,0 +1,37 @@ +{ + "name": "@earendil-works/pi-agent-sqlite-node", + "version": "0.80.10", + "description": "Node sqlite storage backend for @earendil-works/pi-agent-core sessions", + "type": "module", + "main": "./dist/index.js", + "types": "./dist/index.d.ts", + "exports": { + ".": { + "types": "./dist/index.d.ts", + "import": "./dist/index.js" + } + }, + "files": [ + "dist", + "README.md" + ], + "scripts": { + "clean": "node -e \"import('node:fs/promises').then(fs=>fs.rm('dist',{recursive:true,force:true}))\"", + "build": "tsgo -p tsconfig.build.json && node ./scripts/prepare-dist.mjs copy-sqlite-migrations", + "prepublishOnly": "npm run clean && npm run build" + }, + "author": "Earendil Works", + "license": "MIT", + "repository": { + "type": "git", + "url": "git+https://github.com/earendil-works/pi.git", + "directory": "packages/storage/sqlite-node" + }, + "engines": { + "node": ">=22.19.0" + }, + "dependencies": { + "@earendil-works/pi-ai": "^0.80.10", + "@earendil-works/pi-agent-core": "^0.80.10" + } +} diff --git a/packages/storage/sqlite-node/scripts/prepare-dist.mjs b/packages/storage/sqlite-node/scripts/prepare-dist.mjs new file mode 100644 index 000000000..52328367f --- /dev/null +++ b/packages/storage/sqlite-node/scripts/prepare-dist.mjs @@ -0,0 +1,35 @@ +#!/usr/bin/env node + +import { cp, mkdir, rm } from "node:fs/promises"; +import { dirname, resolve } from "node:path"; +import { fileURLToPath } from "node:url"; + +const scriptDir = dirname(fileURLToPath(import.meta.url)); +const packageDir = resolve(scriptDir, ".."); +const distDir = resolve(packageDir, "dist"); +const migrationSourceDir = resolve(packageDir, "src/sqlite/migrations"); +const migrationDestDir = resolve(distDir, "sqlite/migrations"); + +async function clean() { + await rm(distDir, { recursive: true, force: true }); +} + +async function copySqliteMigrations() { + await mkdir(migrationDestDir, { recursive: true }); + await cp(migrationSourceDir, migrationDestDir, { recursive: true }); +} + +const command = process.argv[2]; + +if (command === "clean") { + await clean(); + process.exit(0); +} + +if (command === "copy-sqlite-migrations") { + await copySqliteMigrations(); + process.exit(0); +} + +console.error("Usage: node scripts/prepare-dist.mjs "); +process.exit(1); diff --git a/packages/storage/sqlite-node/src/index.ts b/packages/storage/sqlite-node/src/index.ts new file mode 100644 index 000000000..399eadb61 --- /dev/null +++ b/packages/storage/sqlite-node/src/index.ts @@ -0,0 +1,97 @@ +import type { SQLInputValue } from "node:sqlite"; +import { DatabaseSync } from "node:sqlite"; +import type { SqliteDatabase, SqliteDatabaseFactory, SqliteRunResult, SqliteStatement } from "./sqlite/types.ts"; + +function isNamedParameters(value: unknown): value is Record { + if (value === null || typeof value !== "object") return false; + if (Array.isArray(value) || ArrayBuffer.isView(value)) return false; + return true; +} + +class NodeSqliteStatement implements SqliteStatement { + private readonly statement: ReturnType; + + constructor(statement: ReturnType) { + this.statement = statement; + } + + async run(...params: unknown[]): Promise { + const [first, ...rest] = params; + const result = isNamedParameters(first) + ? this.statement.run(first, ...(rest as SQLInputValue[])) + : this.statement.run(...(params as SQLInputValue[])); + return { + changes: Number(result.changes), + lastInsertRowid: result.lastInsertRowid === undefined ? undefined : Number(result.lastInsertRowid), + }; + } + + async get(...params: unknown[]): Promise { + const [first, ...rest] = params; + return ( + isNamedParameters(first) + ? this.statement.get(first, ...(rest as SQLInputValue[])) + : this.statement.get(...(params as SQLInputValue[])) + ) as TRow | undefined; + } + + async all(...params: unknown[]): Promise { + const [first, ...rest] = params; + return ( + isNamedParameters(first) + ? this.statement.all(first, ...(rest as SQLInputValue[])) + : this.statement.all(...(params as SQLInputValue[])) + ) as TRow[]; + } +} + +class NodeSqliteDatabase implements SqliteDatabase { + private readonly db: DatabaseSync; + + constructor(db: DatabaseSync) { + this.db = db; + } + + async exec(sql: string): Promise { + this.db.exec(sql); + } + + prepare(sql: string): SqliteStatement { + return new NodeSqliteStatement(this.db.prepare(sql)); + } + + async transaction(fn: () => Promise): Promise { + this.db.exec("BEGIN"); + try { + const result = await fn(); + this.db.exec("COMMIT"); + return result; + } catch (error) { + try { + this.db.exec("ROLLBACK"); + } catch { + // Ignore rollback errors to rethrow original error. + } + throw error; + } + } + + async close(): Promise { + this.db.close(); + } +} + +export function wrapNodeSqliteDatabase(db: DatabaseSync): SqliteDatabase { + return new NodeSqliteDatabase(db); +} + +export function createNodeSqliteFactory(): SqliteDatabaseFactory { + return { + async open(path: string): Promise { + return new NodeSqliteDatabase(new DatabaseSync(path)); + }, + }; +} + +// Re-export the SQLite session storage backend and types so this package is a complete node-sqlite backend. +export * from "./sqlite/index.ts"; diff --git a/packages/storage/sqlite-node/src/sqlite/index.ts b/packages/storage/sqlite-node/src/sqlite/index.ts new file mode 100644 index 000000000..279a6083d --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/index.ts @@ -0,0 +1,4 @@ +export * from "./migrations.ts"; +export * from "./repo.ts"; +export * from "./storage/index.ts"; +export * from "./types.ts"; diff --git a/packages/storage/sqlite-node/src/sqlite/migrations.ts b/packages/storage/sqlite-node/src/sqlite/migrations.ts new file mode 100644 index 000000000..2ccd70fc3 --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/migrations.ts @@ -0,0 +1,50 @@ +import { readFile } from "node:fs/promises"; +import { fileURLToPath } from "node:url"; +import type { SqliteDatabase } from "./types.ts"; + +export interface SqliteMigration { + id: string; + order: number; + sql: string; +} + +async function loadMigrationSql(relativePath: string): Promise { + return readFile(fileURLToPath(new URL(relativePath, import.meta.url)), "utf8"); +} + +export async function loadMigrations(): Promise { + return [ + { + id: "001_initial.sql", + order: 1, + sql: await loadMigrationSql("./migrations/001_initial.sql"), + }, + ]; +} + +async function ensureMigrationsTable(db: SqliteDatabase): Promise { + await db.exec(` +CREATE TABLE IF NOT EXISTS migrations ( + id TEXT PRIMARY KEY, + applied_at TEXT NOT NULL +); +`); +} + +export async function applyMigrations(db: SqliteDatabase): Promise { + await ensureMigrationsTable(db); + const migrations = await loadMigrations(); + const appliedRows = await db.prepare("SELECT id FROM migrations ORDER BY applied_at, id").all<{ id: string }>(); + const applied = new Set(appliedRows.map((row) => row.id)); + + for (const migration of migrations) { + if (applied.has(migration.id)) continue; + await db.transaction(async () => { + await db.exec(migration.sql); + await db + .prepare("INSERT INTO migrations (id, applied_at) VALUES (?, ?)") + .run(migration.id, new Date().toISOString()); + }); + applied.add(migration.id); + } +} diff --git a/packages/storage/sqlite-node/src/sqlite/migrations/001_initial.sql b/packages/storage/sqlite-node/src/sqlite/migrations/001_initial.sql new file mode 100644 index 000000000..6e6f397d9 --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/migrations/001_initial.sql @@ -0,0 +1,59 @@ +CREATE TABLE IF NOT EXISTS sessions ( + id TEXT PRIMARY KEY, + created_at TEXT NOT NULL, + cwd TEXT NOT NULL, + parent_session_id TEXT NULL, + metadata TEXT NULL, + active_leaf_id TEXT NULL +) WITHOUT ROWID; + +CREATE INDEX IF NOT EXISTS idx_sessions_created_at ON sessions(created_at DESC); +CREATE INDEX IF NOT EXISTS idx_sessions_cwd ON sessions(cwd); +CREATE INDEX IF NOT EXISTS idx_sessions_parent ON sessions(parent_session_id); + +CREATE TABLE IF NOT EXISTS session_entries ( + session_id TEXT NOT NULL, + id TEXT NOT NULL, + entry_seq INTEGER NOT NULL, + parent_id TEXT NULL, + type TEXT NOT NULL, + timestamp TEXT NOT NULL, + payload TEXT NOT NULL, + PRIMARY KEY (session_id, id) +); + +CREATE UNIQUE INDEX IF NOT EXISTS idx_session_entries_session_seq ON session_entries(session_id, entry_seq); +CREATE INDEX IF NOT EXISTS idx_session_entries_session_parent ON session_entries(session_id, parent_id); +CREATE INDEX IF NOT EXISTS idx_session_entries_session_type ON session_entries(session_id, type); + +CREATE TABLE IF NOT EXISTS session_sequences ( + session_id TEXT PRIMARY KEY, + next_seq INTEGER NOT NULL +) WITHOUT ROWID; + +CREATE TABLE IF NOT EXISTS branch_entries ( + session_id TEXT NOT NULL, + branch_id TEXT NOT NULL, + entry_id TEXT NOT NULL, + entry_seq INTEGER NOT NULL, + PRIMARY KEY (session_id, branch_id, entry_id) +) WITHOUT ROWID; + +CREATE INDEX IF NOT EXISTS idx_branch_entries_session_branch ON branch_entries(session_id, branch_id); +CREATE INDEX IF NOT EXISTS idx_branch_entries_session_branch_seq ON branch_entries(session_id, branch_id, entry_seq); +CREATE INDEX IF NOT EXISTS idx_branch_entries_session_entry ON branch_entries(session_id, entry_id); + +CREATE TABLE IF NOT EXISTS session_materialized ( + session_id TEXT PRIMARY KEY, + payload TEXT NOT NULL +) WITHOUT ROWID; + +CREATE TABLE IF NOT EXISTS entry_materialized ( + session_id TEXT NOT NULL, + entry_seq INTEGER NOT NULL, + type TEXT NOT NULL, + payload TEXT NOT NULL, + PRIMARY KEY (session_id, entry_seq, type) +) WITHOUT ROWID; + +CREATE INDEX IF NOT EXISTS idx_entry_materialized_session_type_seq ON entry_materialized(session_id, type, entry_seq); diff --git a/packages/storage/sqlite-node/src/sqlite/repo.ts b/packages/storage/sqlite-node/src/sqlite/repo.ts new file mode 100644 index 000000000..12a4f9d02 --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/repo.ts @@ -0,0 +1,192 @@ +import type { Session, SessionStorage, SessionTreeEntry } from "@earendil-works/pi-agent-core"; +import { + createSessionId, + getEntriesToFork, + getFileSystemResultOrThrow, + SessionError, + toSession, +} from "@earendil-works/pi-agent-core"; +import { applyMigrations } from "./migrations.ts"; +import { SqliteSessionStorage } from "./storage/index.ts"; +import { rowToMetadata, type SessionRow } from "./storage/sessions.ts"; +import type { + SqliteDatabase, + SqliteDatabaseFactory, + SqliteSessionCreateOptions, + SqliteSessionListOptions, + SqliteSessionMetadata, + SqliteSessionRepoApi, + SqliteSessionRepoEnv, +} from "./types.ts"; + +function getParentPath(path: string): string { + const normalized = path.replace(/[\\/]+$/, ""); + const lastSlash = Math.max(normalized.lastIndexOf("/"), normalized.lastIndexOf("\\")); + if (lastSlash < 0) return "."; + if (lastSlash === 0) return normalized.slice(0, 1); + return normalized.slice(0, lastSlash); +} + +async function configureSqliteDatabase(db: SqliteDatabase): Promise { + await db.exec("PRAGMA journal_mode=WAL"); + await db.exec("PRAGMA synchronous=FULL"); + await db.exec("PRAGMA busy_timeout=5000"); +} + +async function cleanupSessionStorage(storage: SessionStorage): Promise { + const maybeClosable = storage as SessionStorage & { cleanup?: () => Promise }; + if (typeof maybeClosable.cleanup === "function") { + await maybeClosable.cleanup(); + } +} + +export class SqliteSessionRepo implements SqliteSessionRepoApi { + private readonly env: SqliteSessionRepoEnv; + private readonly sqlite: SqliteDatabaseFactory; + private readonly databasePathInput: string; + private databasePath: string | undefined; + + constructor(options: { env: SqliteSessionRepoEnv; sqlite: SqliteDatabaseFactory; databasePath: string }) { + this.env = options.env; + this.sqlite = options.sqlite; + this.databasePathInput = options.databasePath; + } + + private async getDatabasePath(): Promise { + if (!this.databasePath) { + this.databasePath = getFileSystemResultOrThrow( + await this.env.absolutePath(this.databasePathInput), + `Failed to resolve SQLite sessions database ${this.databasePathInput}`, + ); + } + return this.databasePath; + } + + private async ensureDatabaseDir(): Promise { + const path = await this.getDatabasePath(); + const directory = getParentPath(path); + getFileSystemResultOrThrow( + await this.env.createDir(directory, { recursive: true }), + `Failed to create SQLite sessions directory ${directory}`, + ); + } + + private async openDatabase(): Promise { + await this.ensureDatabaseDir(); + const db = await this.sqlite.open(await this.getDatabasePath()); + try { + await configureSqliteDatabase(db); + await applyMigrations(db); + return db; + } catch (error) { + await db.close(); + throw error; + } + } + + async create(options: SqliteSessionCreateOptions): Promise> { + const db = await this.openDatabase(); + try { + const id = options.id ?? createSessionId(); + const storage = await SqliteSessionStorage.create(db, await this.getDatabasePath(), { + cwd: options.cwd, + sessionId: id, + parentSessionId: options.parentSessionId, + metadata: options.metadata, + }); + return toSession(storage); + } catch (error) { + await db.close(); + throw error; + } + } + + async open(metadata: SqliteSessionMetadata): Promise> { + if ( + !getFileSystemResultOrThrow(await this.env.exists(metadata.path), `Failed to check database ${metadata.path}`) + ) { + throw new SessionError("not_found", `Session not found: ${metadata.id}`); + } + const db = await this.openDatabase(); + try { + const storage = await SqliteSessionStorage.open(db, metadata); + return toSession(storage); + } catch (error) { + await db.close(); + throw error; + } + } + + async list(options: SqliteSessionListOptions = {}): Promise { + const path = await this.getDatabasePath(); + if (!getFileSystemResultOrThrow(await this.env.exists(path), `Failed to check database ${path}`)) { + return []; + } + const db = await this.openDatabase(); + try { + const rows = options.cwd + ? await db + .prepare( + "SELECT id, created_at, metadata, cwd, parent_session_id, active_leaf_id FROM sessions WHERE cwd = ? ORDER BY created_at DESC", + ) + .all(options.cwd) + : await db + .prepare( + "SELECT id, created_at, metadata, cwd, parent_session_id, active_leaf_id FROM sessions ORDER BY created_at DESC", + ) + .all(); + return rows.map((row) => rowToMetadata(row, path)); + } finally { + await db.close(); + } + } + + async delete(metadata: SqliteSessionMetadata): Promise { + const db = await this.openDatabase(); + try { + await db.transaction(async () => { + await db.prepare("DELETE FROM branch_entries WHERE session_id = ?").run(metadata.id); + await db.prepare("DELETE FROM session_entries WHERE session_id = ?").run(metadata.id); + await db.prepare("DELETE FROM entry_materialized WHERE session_id = ?").run(metadata.id); + await db.prepare("DELETE FROM session_materialized WHERE session_id = ?").run(metadata.id); + await db.prepare("DELETE FROM session_sequences WHERE session_id = ?").run(metadata.id); + const result = await db.prepare("DELETE FROM sessions WHERE id = ?").run(metadata.id); + if (result.changes === 0) { + throw new SessionError("not_found", `Session not found: ${metadata.id}`); + } + }); + } finally { + await db.close(); + } + } + + async fork( + sourceMetadata: SqliteSessionMetadata, + options: SqliteSessionCreateOptions & { entryId?: string; position?: "before" | "at"; id?: string }, + ): Promise> { + const source = await this.open(sourceMetadata); + let forkedEntries: SessionTreeEntry[]; + try { + forkedEntries = await getEntriesToFork(source.getStorage(), options); + } finally { + await cleanupSessionStorage(source.getStorage()); + } + const db = await this.openDatabase(); + try { + const id = options.id ?? createSessionId(); + const storage = await SqliteSessionStorage.create(db, await this.getDatabasePath(), { + cwd: options.cwd, + sessionId: id, + parentSessionId: options.parentSessionId ?? sourceMetadata.id, + metadata: options.metadata ?? sourceMetadata.metadata, + }); + for (const entry of forkedEntries) { + await storage.appendEntry(entry); + } + return toSession(storage); + } catch (error) { + await db.close(); + throw error; + } + } +} diff --git a/packages/storage/sqlite-node/src/sqlite/storage/branch-entries.ts b/packages/storage/sqlite-node/src/sqlite/storage/branch-entries.ts new file mode 100644 index 000000000..0ea601561 --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/storage/branch-entries.ts @@ -0,0 +1,57 @@ +import type { SessionTreeEntry } from "@earendil-works/pi-agent-core"; +import type { SqliteDatabase } from "../types.ts"; +import { decodeEntry, type SessionEntryRow } from "./session-entries.ts"; +import { invalidSession } from "./shared.ts"; + +export interface BranchEntryRow { + entry_id: string; + entry_seq: number; +} + +export async function getMaterializedBranchPathOrCompaction( + db: SqliteDatabase, + sessionId: string, + branchId: string, + byId: Map, +): Promise { + const branchRows = await db + .prepare( + "SELECT entry_id, entry_seq FROM branch_entries WHERE session_id = ? AND branch_id = ? ORDER BY entry_seq", + ) + .all(sessionId, branchId); + if (branchRows.length === 0) { + return []; + } + const entryIds = branchRows.map((row) => row.entry_id); + const placeholders = entryIds.map(() => "?").join(", "); + const entryRows = await db + .prepare( + `SELECT session_id, id, entry_seq, parent_id, type, timestamp, payload FROM session_entries WHERE session_id = ? AND id IN (${placeholders})`, + ) + .all(sessionId, ...entryIds); + const entryRowsById = new Map(entryRows.map((row) => [row.id, row])); + const entries: SessionTreeEntry[] = []; + for (const branchRow of branchRows) { + // leaf entries are navigation markers used to mark which branch became active; + // they are not part of the model/context path reconstructed from branch_entries. + const cached = byId.get(branchRow.entry_id); + if (cached) { + if (cached.type !== "leaf") { + entries.push(cached); + } + continue; + } + const entryRow = entryRowsById.get(branchRow.entry_id); + if (!entryRow) throw invalidSession(`missing entry row for branch entry ${branchRow.entry_id}`); + try { + const entry = decodeEntry(entryRow); + byId.set(entry.id, entry); + if (entry.type !== "leaf") { + entries.push(entry); + } + } catch { + throw invalidSession(`invalid entry row for branch entry ${branchRow.entry_id}`); + } + } + return entries; +} diff --git a/packages/storage/sqlite-node/src/sqlite/storage/index.ts b/packages/storage/sqlite-node/src/sqlite/storage/index.ts new file mode 100644 index 000000000..93376852e --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/storage/index.ts @@ -0,0 +1,449 @@ +import type { + LeafEntry, + SessionEntryCursorOptions, + SessionStorage, + SessionTreeEntry, +} from "@earendil-works/pi-agent-core"; +import { SessionError } from "@earendil-works/pi-agent-core"; +import { uuidv7 } from "@earendil-works/pi-ai"; +import type { SqliteDatabase, SqliteSessionMetadata } from "../types.ts"; +import { getMaterializedBranchPathOrCompaction } from "./branch-entries.ts"; +import { decodeEntry, encodeEntry, type SessionEntryRow } from "./session-entries.ts"; +import { + applyEntryToMaterializedState, + createEmptyMaterializedState, + type EntryMaterializedRow, + entryMaterializedValues, + materializedStateFromRows, + materializedStateValues, + type SessionMaterializedRow, + type SessionMaterializedState, + serializeSummary, + sessionStatsFromMaterializedState, +} from "./session-materialized.ts"; +import { advanceSequence, getNextSequence } from "./session-sequences.ts"; +import { rowToMetadata, type SessionRow } from "./sessions.ts"; +import { generateEntryId, invalidSession, leafIdAfterEntry } from "./shared.ts"; + +async function decodeEntryRows(entryRows: SessionEntryRow[]): Promise<{ + entries: SessionTreeEntry[]; + leafId: string | null; +}> { + const entries: SessionTreeEntry[] = []; + let leafId: string | null = null; + for (const entryRow of entryRows) { + try { + const entry = decodeEntry(entryRow); + entries.push(entry); + leafId = leafIdAfterEntry(entry); + } catch { + // Keep JSONL-like permissive resume behavior: skip malformed entries. + } + } + return { entries, leafId }; +} + +async function loadEntryRowsByIds( + db: SqliteDatabase, + sessionId: string, + entryIds: string[], +): Promise> { + if (entryIds.length === 0) return new Map(); + const placeholders = entryIds.map(() => "?").join(", "); + const rows = await db + .prepare( + `SELECT session_id, id, entry_seq, parent_id, type, timestamp, payload FROM session_entries WHERE session_id = ? AND id IN (${placeholders})`, + ) + .all(sessionId, ...entryIds); + return new Map(rows.map((row) => [row.id, row])); +} + +async function loadActiveBranchId(db: SqliteDatabase, sessionId: string): Promise { + // branch_entries includes leaf navigation entries for the active branch, so the + // newest branch_entries row identifies the branch that was most recently made active. + const row = await db + .prepare( + "SELECT branch_id FROM branch_entries WHERE session_id = ? ORDER BY entry_seq DESC, branch_id DESC LIMIT 1", + ) + .get<{ branch_id: string }>(sessionId); + return row?.branch_id ?? null; +} + +async function hasExistingChild(db: SqliteDatabase, sessionId: string, parentId: string | null): Promise { + const row = + parentId === null + ? await db + .prepare("SELECT 1 AS found FROM session_entries WHERE session_id = ? AND parent_id IS NULL LIMIT 1") + .get<{ found: number }>(sessionId) + : await db + .prepare("SELECT 1 AS found FROM session_entries WHERE session_id = ? AND parent_id = ? LIMIT 1") + .get<{ found: number }>(sessionId, parentId); + return row !== undefined; +} + +async function loadSqliteStorage( + db: SqliteDatabase, + sessionId: string, +): Promise<{ + row: SessionRow; + leafId: string | null; + activeBranchId: string | null; + materializedState: SessionMaterializedState; +}> { + const row = await db + .prepare("SELECT id, created_at, metadata, cwd, parent_session_id, active_leaf_id FROM sessions WHERE id = ?") + .get(sessionId); + if (!row) throw new SessionError("not_found", `Session not found: ${sessionId}`); + + const leafId = row.active_leaf_id; + const materializedRow = await db + .prepare("SELECT session_id, payload FROM session_materialized WHERE session_id = ?") + .get(sessionId); + if (!materializedRow) throw invalidSession(`missing materialized row for session ${sessionId}`); + const entryMaterializedRows = await db + .prepare( + "SELECT session_id, entry_seq, type, payload FROM entry_materialized WHERE session_id = ? ORDER BY entry_seq, type", + ) + .all(sessionId); + return { + row, + leafId, + activeBranchId: await loadActiveBranchId(db, sessionId), + materializedState: materializedStateFromRows(materializedRow, entryMaterializedRows), + }; +} + +export class SqliteSessionStorage implements SessionStorage { + private readonly db: SqliteDatabase; + private readonly metadata: SqliteSessionMetadata; + private byId: Map; + private labelsById: Map; + private currentLeafId: string | null; + private activeBranchId: string | null; + private materializedState: SessionMaterializedState; + + private async getPathToRootOrCompactionEntries(leafId: string | null): Promise { + if (leafId === null) return []; + const path: SessionTreeEntry[] = []; + let stopAtEntryId: string | null = null; + let current = await this.getEntry(leafId); + if (!current) throw new SessionError("not_found", `Entry ${leafId} not found`); + while (current) { + path.unshift(current); + if (stopAtEntryId !== null && current.id === stopAtEntryId) break; + if (current.type === "compaction") { + if (current.retainedTail) break; + stopAtEntryId = current.firstKeptEntryId ?? null; + } + if (!current.parentId) break; + const parent = await this.getEntry(current.parentId); + if (!parent) throw new SessionError("invalid_session", `Entry ${current.parentId} not found`); + current = parent; + } + return path; + } + + private async materializeBranch(leafId: string | null): Promise { + const branchId = uuidv7(); + // Rebuild the branch path only when branch membership changes: branch switch + // (leaf navigation) or a new fork from a parent that already has a child. + // Linear appends stay cheap and extend the active branch incrementally. + const path = await this.getPathToRootOrCompactionEntries(leafId); + const entryRowsById = await loadEntryRowsByIds( + this.db, + this.metadata.id, + path.map((entry) => entry.id), + ); + for (const entry of path) { + const entryRow = entryRowsById.get(entry.id); + if (!entryRow) throw invalidSession(`missing entry row for session ${this.metadata.id} entry ${entry.id}`); + await this.db + .prepare("INSERT INTO branch_entries (session_id, branch_id, entry_id, entry_seq) VALUES (?, ?, ?, ?)") + .run(this.metadata.id, branchId, entry.id, entryRow.entry_seq); + } + this.activeBranchId = branchId; + } + + private async appendToActiveBranch(entryId: string, parentId: string | null): Promise { + if (!this.activeBranchId) { + await this.materializeBranch(parentId); + } + // After a branch is materialized/resynced, subsequent linear appends only add the + // new tip entry. We do not rebuild the full branch on every append. + if (!this.activeBranchId) { + throw invalidSession(`active branch missing for session ${this.metadata.id}`); + } + const entryRow = await this.db + .prepare("SELECT entry_seq FROM session_entries WHERE session_id = ? AND id = ?") + .get<{ entry_seq: number }>(this.metadata.id, entryId); + if (!entryRow) throw invalidSession(`missing entry row for session ${this.metadata.id} entry ${entryId}`); + await this.db + .prepare("INSERT INTO branch_entries (session_id, branch_id, entry_id, entry_seq) VALUES (?, ?, ?, ?)") + .run(this.metadata.id, this.activeBranchId, entryId, entryRow.entry_seq); + } + + private constructor( + db: SqliteDatabase, + metadata: SqliteSessionMetadata, + entries: SessionTreeEntry[] | null, + leafId: string | null, + activeBranchId: string | null, + materializedState: SessionMaterializedState, + ) { + this.db = db; + this.metadata = metadata; + this.byId = new Map((entries ?? []).map((entry) => [entry.id, entry])); + this.materializedState = materializedState; + this.labelsById = materializedState.labelsById; + this.currentLeafId = leafId; + this.activeBranchId = activeBranchId; + } + + static async open(db: SqliteDatabase, metadata: SqliteSessionMetadata): Promise { + const loaded = await loadSqliteStorage(db, metadata.id); + return new SqliteSessionStorage( + db, + rowToMetadata(loaded.row, metadata.path), + null, + loaded.leafId, + loaded.activeBranchId, + loaded.materializedState, + ); + } + + static async create( + db: SqliteDatabase, + path: string, + options: { + cwd: string; + sessionId: string; + parentSessionId?: string; + metadata?: Record; + }, + ): Promise { + const createdAt = new Date().toISOString(); + await db + .prepare( + "INSERT INTO sessions (id, created_at, metadata, cwd, parent_session_id, active_leaf_id) VALUES (?, ?, ?, ?, ?, ?)", + ) + .run( + options.sessionId, + createdAt, + options.metadata === undefined ? null : JSON.stringify(options.metadata), + options.cwd, + options.parentSessionId ?? null, + null, + ); + await db.prepare("INSERT INTO session_sequences (session_id, next_seq) VALUES (?, ?)").run(options.sessionId, 1); + await db + .prepare("INSERT INTO session_materialized (session_id, payload) VALUES (?, ?)") + .run(...materializedStateValues(options.sessionId, createEmptyMaterializedState())); + return new SqliteSessionStorage( + db, + { + id: options.sessionId, + createdAt, + cwd: options.cwd, + path, + parentSessionId: options.parentSessionId, + metadata: options.metadata, + }, + [], + null, + null, + createEmptyMaterializedState(), + ); + } + + async getMetadata(): Promise { + return this.metadata; + } + + async getLeafId(): Promise { + return this.currentLeafId; + } + + async setLeafId(leafId: string | null): Promise { + if (leafId !== null && !(await this.getEntry(leafId))) { + throw new SessionError("not_found", `Entry ${leafId} not found`); + } + const entry: LeafEntry = { + type: "leaf", + id: await this.createEntryId(), + parentId: this.currentLeafId, + timestamp: new Date().toISOString(), + targetId: leafId, + }; + await this.appendEntry(entry); + } + + async createEntryId(): Promise { + for (let i = 0; i < 100; i++) { + const id = generateEntryId(this.byId); + const existing = await this.db + .prepare("SELECT 1 AS found FROM session_entries WHERE session_id = ? AND id = ? LIMIT 1") + .get<{ found: number }>(this.metadata.id, id); + if (!existing) return id; + } + return uuidv7(); + } + + async appendEntry(entry: SessionTreeEntry): Promise { + const encoded = encodeEntry(entry); + const previousMaterializedState: SessionMaterializedState = { + ...this.materializedState, + labelsById: new Map(this.materializedState.labelsById), + modelThinkingConfigs: [...this.materializedState.modelThinkingConfigs], + currentModel: this.materializedState.currentModel ? { ...this.materializedState.currentModel } : null, + }; + const previousById = new Map(this.byId); + const previousLeafId = this.currentLeafId; + const previousActiveBranchId = this.activeBranchId; + try { + applyEntryToMaterializedState(this.materializedState, entry); + await this.db.transaction(async () => { + const parentHadExistingChild = await hasExistingChild(this.db, this.metadata.id, entry.parentId); + const nextSeq = await getNextSequence(this.db, this.metadata.id); + await this.db + .prepare( + "INSERT INTO session_entries (session_id, id, entry_seq, parent_id, type, timestamp, payload) VALUES (?, ?, ?, ?, ?, ?, ?)", + ) + .run(this.metadata.id, entry.id, nextSeq, entry.parentId, entry.type, entry.timestamp, encoded.payload); + await advanceSequence(this.db, this.metadata.id, nextSeq); + await this.db + .prepare("UPDATE session_materialized SET payload = ? WHERE session_id = ?") + .run(serializeSummary(this.materializedState), this.metadata.id); + for (const materializedEntry of entryMaterializedValues(entry)) { + await this.db + .prepare("INSERT INTO entry_materialized (session_id, entry_seq, type, payload) VALUES (?, ?, ?, ?)") + .run(this.metadata.id, nextSeq, materializedEntry.type, materializedEntry.payload); + } + this.byId.set(entry.id, entry); + this.currentLeafId = leafIdAfterEntry(entry); + await this.db + .prepare("UPDATE sessions SET active_leaf_id = ? WHERE id = ?") + .run(this.currentLeafId, this.metadata.id); + if (entry.type === "leaf") { + this.activeBranchId = null; + await this.materializeBranch(entry.targetId); + await this.appendToActiveBranch(entry.id, entry.parentId); + } else { + if (parentHadExistingChild) { + await this.materializeBranch(entry.parentId); + } + await this.appendToActiveBranch(entry.id, entry.parentId); + } + }); + } catch (error) { + this.materializedState = previousMaterializedState; + this.labelsById = previousMaterializedState.labelsById; + this.byId = previousById; + this.currentLeafId = previousLeafId; + this.activeBranchId = previousActiveBranchId; + if (error instanceof SessionError) throw error; + throw new SessionError("storage", `Failed to append SQLite session entry ${entry.id}`); + } + } + + async getEntry(id: string): Promise { + const cached = this.byId.get(id); + if (cached) return cached; + const row = await this.db + .prepare( + "SELECT session_id, id, entry_seq, parent_id, type, timestamp, payload FROM session_entries WHERE session_id = ? AND id = ?", + ) + .get(this.metadata.id, id); + if (!row) return undefined; + try { + const entry = decodeEntry(row); + this.byId.set(entry.id, entry); + return entry; + } catch { + return undefined; + } + } + + async findEntries( + type: TType, + ): Promise>> { + const rows = await this.db + .prepare( + "SELECT session_id, id, entry_seq, parent_id, type, timestamp, payload FROM session_entries WHERE session_id = ? AND type = ? ORDER BY entry_seq", + ) + .all(this.metadata.id, type); + const entries: Array> = []; + for (const row of rows) { + try { + const entry = decodeEntry(row) as Extract; + this.byId.set(entry.id, entry); + entries.push(entry); + } catch { + // Keep JSONL-like permissive resume behavior: skip malformed entries. + } + } + return entries; + } + + async getLabel(id: string): Promise { + return this.labelsById.get(id); + } + + async getSessionName(): Promise { + return this.materializedState.name; + } + + async getSessionStats() { + return sessionStatsFromMaterializedState(this.materializedState); + } + + async getPathToRootOrCompaction(leafId: string | null): Promise { + if (leafId === null) return []; + if (leafId === this.currentLeafId) { + if (!this.activeBranchId) { + throw invalidSession(`missing active branch for session ${this.metadata.id} leaf ${leafId}`); + } + return getMaterializedBranchPathOrCompaction(this.db, this.metadata.id, this.activeBranchId, this.byId); + } + return this.getPathToRootOrCompactionEntries(leafId); + } + + async getEntries(options?: SessionEntryCursorOptions): Promise { + const limit = options?.limit; + if (limit !== undefined) { + const beforeOrAtEntrySeq = + options?.afterEntrySeq ?? + ( + await this.db + .prepare("SELECT entry_seq FROM session_entries WHERE session_id = ? ORDER BY entry_seq DESC LIMIT 1") + .get<{ entry_seq: number }>(this.metadata.id) + )?.entry_seq; + if (beforeOrAtEntrySeq === undefined) { + return []; + } + const rows = await this.db + .prepare( + "SELECT session_id, id, entry_seq, parent_id, type, timestamp, payload FROM session_entries WHERE session_id = ? AND entry_seq <= ? ORDER BY entry_seq DESC LIMIT ?", + ) + .all(this.metadata.id, beforeOrAtEntrySeq, limit); + const entries = (await decodeEntryRows(rows)).entries; + for (const entry of entries) { + this.byId.set(entry.id, entry); + } + return entries.reverse(); + } + const rows = await this.db + .prepare( + "SELECT session_id, id, entry_seq, parent_id, type, timestamp, payload FROM session_entries WHERE session_id = ? ORDER BY entry_seq", + ) + .all(this.metadata.id); + const entries = (await decodeEntryRows(rows)).entries; + for (const entry of entries) { + this.byId.set(entry.id, entry); + } + return entries; + } + + async cleanup(): Promise { + await this.db.close(); + } +} diff --git a/packages/storage/sqlite-node/src/sqlite/storage/session-entries.ts b/packages/storage/sqlite-node/src/sqlite/storage/session-entries.ts new file mode 100644 index 000000000..1b1cd75a8 --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/storage/session-entries.ts @@ -0,0 +1,217 @@ +import type { SessionTreeEntry, SessionTreeEntryBase } from "@earendil-works/pi-agent-core"; +import { invalidEntry, isRecord } from "./shared.ts"; + +export interface SessionEntryRow { + session_id: string; + id: string; + entry_seq: number; + parent_id: string | null; + type: SessionTreeEntry["type"]; + timestamp: string; + payload: string; +} + +export type EncodedEntry = { + payload: string; +}; + +type EntryPayload = Omit; + +type MessagePayload = EntryPayload>; +type ThinkingLevelChangePayload = EntryPayload>; +type ModelChangePayload = EntryPayload>; +type ActiveToolsChangePayload = EntryPayload>; +type CompactionPayload = EntryPayload>; +type BranchSummaryPayload = EntryPayload>; +type CustomPayload = EntryPayload>; +type CustomMessagePayload = EntryPayload>; +type LabelPayload = EntryPayload>; +type SessionInfoPayload = EntryPayload>; +type LeafPayload = EntryPayload>; + +function parsePayload(row: SessionEntryRow): unknown { + try { + return JSON.parse(row.payload); + } catch (error) { + throw invalidEntry(`entry ${row.id} payload is not valid JSON`, error instanceof Error ? error : undefined); + } +} + +function isTextImageContentArray(value: unknown): boolean { + return ( + Array.isArray(value) && + value.every( + (item) => + isRecord(item) && typeof item.type === "string" && (item.type !== "text" || typeof item.text === "string"), + ) + ); +} + +export function validateSessionTreeEntry(entry: SessionTreeEntry): void { + if (typeof entry.id !== "string" || !entry.id) throw invalidEntry("entry is missing id"); + if (entry.parentId !== null && typeof entry.parentId !== "string") { + throw invalidEntry(`entry ${entry.id} has invalid parentId`); + } + if (typeof entry.timestamp !== "string" || !entry.timestamp) { + throw invalidEntry(`entry ${entry.id} is missing timestamp`); + } + + switch (entry.type) { + case "message": + if (!isRecord(entry.message) || typeof entry.message.role !== "string") { + throw invalidEntry(`entry ${entry.id} is missing message payload`); + } + break; + case "thinking_level_change": + if (typeof entry.thinkingLevel !== "string") throw invalidEntry(`entry ${entry.id} is missing thinkingLevel`); + break; + case "model_change": + if (typeof entry.provider !== "string" || typeof entry.modelId !== "string") { + throw invalidEntry(`entry ${entry.id} has invalid model_change payload`); + } + break; + case "active_tools_change": + if ( + !Array.isArray(entry.activeToolNames) || + entry.activeToolNames.some((value) => typeof value !== "string") + ) { + throw invalidEntry(`entry ${entry.id} has invalid active_tools_change payload`); + } + break; + case "compaction": + if ( + typeof entry.summary !== "string" || + typeof entry.firstKeptEntryId !== "string" || + typeof entry.tokensBefore !== "number" || + (entry.retainedTail !== undefined && !Array.isArray(entry.retainedTail)) + ) { + throw invalidEntry(`entry ${entry.id} has invalid compaction payload`); + } + break; + case "branch_summary": + if (typeof entry.fromId !== "string" || typeof entry.summary !== "string") { + throw invalidEntry(`entry ${entry.id} has invalid branch_summary payload`); + } + break; + case "custom": + if (typeof entry.customType !== "string") throw invalidEntry(`entry ${entry.id} has invalid custom payload`); + break; + case "custom_message": + if ( + typeof entry.customType !== "string" || + typeof entry.display !== "boolean" || + !(typeof entry.content === "string" || isTextImageContentArray(entry.content)) + ) { + throw invalidEntry(`entry ${entry.id} has invalid custom_message payload`); + } + break; + case "label": + if (typeof entry.targetId !== "string" || (entry.label !== undefined && typeof entry.label !== "string")) { + throw invalidEntry(`entry ${entry.id} has invalid label payload`); + } + break; + case "session_info": + if (entry.name !== undefined && typeof entry.name !== "string") { + throw invalidEntry(`entry ${entry.id} has invalid session_info payload`); + } + break; + case "leaf": + if (entry.targetId !== null && typeof entry.targetId !== "string") { + throw invalidEntry(`entry ${entry.id} has invalid leaf payload`); + } + break; + default: { + const exhaustive: never = entry; + throw invalidEntry(`unknown entry type ${(exhaustive as { type?: string }).type ?? "unknown"}`); + } + } +} + +function entryToPayload(entry: TEntry): EntryPayload { + const { type: _type, id: _id, parentId: _parentId, timestamp: _timestamp, ...payload } = entry; + return payload as EntryPayload; +} + +export function encodeEntry(entry: SessionTreeEntry): EncodedEntry { + validateSessionTreeEntry(entry); + return { payload: JSON.stringify(entryToPayload(entry)) }; +} + +export function decodeEntry(row: SessionEntryRow): SessionTreeEntry { + const payload = parsePayload(row); + if (!isRecord(payload)) throw invalidEntry(`entry ${row.id} payload is not an object`); + const base = { + id: row.id, + parentId: row.parent_id, + timestamp: row.timestamp, + }; + + switch (row.type) { + case "message": { + if (!("message" in payload)) throw invalidEntry(`entry ${row.id} is missing message payload`); + const messagePayload = payload as MessagePayload; + return { ...base, type: "message", ...messagePayload }; + } + case "thinking_level_change": + if (typeof payload.thinkingLevel !== "string") throw invalidEntry(`entry ${row.id} is missing thinkingLevel`); + return { ...base, type: "thinking_level_change", ...(payload as ThinkingLevelChangePayload) }; + case "model_change": + if (typeof payload.provider !== "string" || typeof payload.modelId !== "string") { + throw invalidEntry(`entry ${row.id} has invalid model_change payload`); + } + return { ...base, type: "model_change", ...(payload as ModelChangePayload) }; + case "active_tools_change": + if ( + !Array.isArray(payload.activeToolNames) || + payload.activeToolNames.some((value) => typeof value !== "string") + ) { + throw invalidEntry(`entry ${row.id} has invalid active_tools_change payload`); + } + return { ...base, type: "active_tools_change", ...(payload as ActiveToolsChangePayload) }; + case "compaction": + if ( + typeof payload.summary !== "string" || + typeof payload.firstKeptEntryId !== "string" || + typeof payload.tokensBefore !== "number" || + (payload.retainedTail !== undefined && !Array.isArray(payload.retainedTail)) + ) { + throw invalidEntry(`entry ${row.id} has invalid compaction payload`); + } + return { ...base, type: "compaction", ...(payload as CompactionPayload) }; + case "branch_summary": + if (typeof payload.fromId !== "string" || typeof payload.summary !== "string") { + throw invalidEntry(`entry ${row.id} has invalid branch_summary payload`); + } + return { ...base, type: "branch_summary", ...(payload as BranchSummaryPayload) }; + case "custom": + if (typeof payload.customType !== "string") throw invalidEntry(`entry ${row.id} has invalid custom payload`); + return { ...base, type: "custom", ...(payload as CustomPayload) }; + case "custom_message": + if ( + typeof payload.customType !== "string" || + typeof payload.display !== "boolean" || + !("content" in payload) + ) { + throw invalidEntry(`entry ${row.id} has invalid custom_message payload`); + } + return { ...base, type: "custom_message", ...(payload as CustomMessagePayload) }; + case "label": + if (typeof payload.targetId !== "string") throw invalidEntry(`entry ${row.id} has invalid label payload`); + if (payload.label !== undefined && typeof payload.label !== "string") { + throw invalidEntry(`entry ${row.id} has invalid label payload`); + } + return { ...base, type: "label", ...(payload as LabelPayload) }; + case "session_info": + if (payload.name !== undefined && typeof payload.name !== "string") { + throw invalidEntry(`entry ${row.id} has invalid session_info payload`); + } + return { ...base, type: "session_info", ...(payload as SessionInfoPayload) }; + case "leaf": + if (payload.targetId !== null && typeof payload.targetId !== "string") { + throw invalidEntry(`entry ${row.id} has invalid leaf payload`); + } + return { ...base, type: "leaf", ...(payload as LeafPayload) }; + default: + throw invalidEntry(`unknown entry type ${row.type}`); + } +} diff --git a/packages/storage/sqlite-node/src/sqlite/storage/session-materialized.ts b/packages/storage/sqlite-node/src/sqlite/storage/session-materialized.ts new file mode 100644 index 000000000..bba1d9913 --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/storage/session-materialized.ts @@ -0,0 +1,368 @@ +import type { SessionStats, SessionTreeEntry, ThinkingLevel } from "@earendil-works/pi-agent-core"; +import { invalidSession, isRecord } from "./shared.ts"; + +export interface SessionMaterializedRow { + session_id: string; + payload: string; +} + +export interface EntryMaterializedRow { + session_id: string; + entry_seq: number; + type: string; + payload: string; +} + +export interface ModelThinkingConfig { + provider: string; + modelId: string; + thinkingLevel: ThinkingLevel; +} + +export interface SessionMaterializedState { + name: string | undefined; + messageCount: number; + cachedTokens: number; + uncachedTokens: number; + totalTokens: number; + costTotal: number; + labelsById: Map; + modelThinkingConfigs: ModelThinkingConfig[]; + currentModel: { provider: string; modelId: string } | null; + currentThinkingLevel: ThinkingLevel | null; +} + +interface SessionMaterializedSummary { + name?: string; + messageCount: number; + cachedTokens: number; + uncachedTokens: number; + totalTokens: number; + costTotal: number; + currentModel?: { provider: string; modelId: string } | null; + currentThinkingLevel?: ThinkingLevel | null; +} + +function compareModelThinkingConfig(left: ModelThinkingConfig, right: ModelThinkingConfig): number { + return ( + left.provider.localeCompare(right.provider) || + left.modelId.localeCompare(right.modelId) || + left.thinkingLevel.localeCompare(right.thinkingLevel) + ); +} + +function normalizeModelThinkingConfigs(configs: readonly ModelThinkingConfig[]): ModelThinkingConfig[] { + const unique = new Map(); + for (const config of configs) { + unique.set(`${config.provider}\u0000${config.modelId}\u0000${config.thinkingLevel}`, config); + } + return [...unique.values()].sort(compareModelThinkingConfig); +} + +function addModelThinkingConfig( + state: SessionMaterializedState, + provider: string, + modelId: string, + thinkingLevel: ThinkingLevel, +): void { + state.modelThinkingConfigs = normalizeModelThinkingConfigs([ + ...state.modelThinkingConfigs, + { provider, modelId, thinkingLevel }, + ]); +} + +export function isThinkingLevel(value: unknown): value is ThinkingLevel { + return ( + value === "off" || + value === "minimal" || + value === "low" || + value === "medium" || + value === "high" || + value === "xhigh" + ); +} + +function getAssistantUsage(message: unknown): + | { + provider: string; + modelId: string; + input: number; + output: number; + cacheRead: number; + cacheWrite: number; + costTotal: number; + } + | undefined { + if (!isRecord(message) || message.role !== "assistant") return undefined; + if (typeof message.provider !== "string" || typeof message.model !== "string") return undefined; + if (!isRecord(message.usage) || !isRecord(message.usage.cost)) return undefined; + const { input, output, cacheRead, cacheWrite } = message.usage; + const costTotal = message.usage.cost.total; + if ( + typeof input !== "number" || + typeof output !== "number" || + typeof cacheRead !== "number" || + typeof cacheWrite !== "number" || + typeof costTotal !== "number" + ) { + return undefined; + } + return { + provider: message.provider, + modelId: message.model, + input, + output, + cacheRead, + cacheWrite, + costTotal, + }; +} + +export function createEmptyMaterializedState(): SessionMaterializedState { + return { + name: undefined, + messageCount: 0, + cachedTokens: 0, + uncachedTokens: 0, + totalTokens: 0, + costTotal: 0, + labelsById: new Map(), + modelThinkingConfigs: [], + currentModel: null, + currentThinkingLevel: null, + }; +} + +export function applyEntryToMaterializedState(state: SessionMaterializedState, entry: SessionTreeEntry): void { + switch (entry.type) { + case "session_info": + state.name = entry.name?.trim() || undefined; + break; + case "label": { + const label = entry.label?.trim(); + if (label) { + state.labelsById.set(entry.targetId, label); + } else { + state.labelsById.delete(entry.targetId); + } + break; + } + case "model_change": + state.currentModel = { provider: entry.provider, modelId: entry.modelId }; + if (state.currentThinkingLevel) { + addModelThinkingConfig(state, entry.provider, entry.modelId, state.currentThinkingLevel); + } + break; + case "thinking_level_change": + if (!isThinkingLevel(entry.thinkingLevel)) break; + state.currentThinkingLevel = entry.thinkingLevel; + if (state.currentModel) { + addModelThinkingConfig(state, state.currentModel.provider, state.currentModel.modelId, entry.thinkingLevel); + } + break; + case "message": { + state.messageCount += 1; + const usage = getAssistantUsage(entry.message); + if (!usage) break; + state.cachedTokens += usage.cacheRead; + state.uncachedTokens += usage.input + usage.cacheWrite; + state.totalTokens += usage.input + usage.output + usage.cacheRead + usage.cacheWrite; + state.costTotal += usage.costTotal; + state.currentModel = { provider: usage.provider, modelId: usage.modelId }; + if (state.currentThinkingLevel) { + addModelThinkingConfig(state, usage.provider, usage.modelId, state.currentThinkingLevel); + } + break; + } + case "compaction": + case "branch_summary": { + const usage = entry.usage; + if ( + !isRecord(usage) || + !isRecord(usage.cost) || + typeof usage.input !== "number" || + typeof usage.output !== "number" || + typeof usage.cacheRead !== "number" || + typeof usage.cacheWrite !== "number" || + typeof usage.cost.total !== "number" + ) { + break; + } + state.cachedTokens += usage.cacheRead; + state.uncachedTokens += usage.input + usage.cacheWrite; + state.totalTokens += usage.input + usage.output + usage.cacheRead + usage.cacheWrite; + state.costTotal += usage.cost.total; + break; + } + case "active_tools_change": + case "custom": + case "custom_message": + case "leaf": + break; + default: { + const exhaustive: never = entry; + void exhaustive; + break; + } + } +} + +export function serializeSummary(state: SessionMaterializedState): string { + const summary: SessionMaterializedSummary = { + name: state.name, + messageCount: state.messageCount, + cachedTokens: state.cachedTokens, + uncachedTokens: state.uncachedTokens, + totalTokens: state.totalTokens, + costTotal: state.costTotal, + currentModel: state.currentModel, + currentThinkingLevel: state.currentThinkingLevel, + }; + return JSON.stringify(summary); +} + +function parseSummary(json: string): SessionMaterializedSummary { + let parsed: unknown; + try { + parsed = JSON.parse(json); + } catch (error) { + throw invalidSession( + `materialized session summary is not valid JSON`, + error instanceof Error ? error : undefined, + ); + } + if (!isRecord(parsed) || Array.isArray(parsed)) { + throw invalidSession("materialized session summary is not an object"); + } + const currentModel = parsed.currentModel; + const currentThinkingLevel = parsed.currentThinkingLevel; + if ( + (parsed.name !== undefined && typeof parsed.name !== "string") || + typeof parsed.messageCount !== "number" || + typeof parsed.cachedTokens !== "number" || + typeof parsed.uncachedTokens !== "number" || + typeof parsed.totalTokens !== "number" || + typeof parsed.costTotal !== "number" || + (currentModel !== undefined && + currentModel !== null && + (!isRecord(currentModel) || + typeof currentModel.provider !== "string" || + typeof currentModel.modelId !== "string")) || + (currentThinkingLevel !== undefined && currentThinkingLevel !== null && !isThinkingLevel(currentThinkingLevel)) + ) { + throw invalidSession("materialized session summary has invalid fields"); + } + return { + name: parsed.name?.trim() || undefined, + messageCount: parsed.messageCount, + cachedTokens: parsed.cachedTokens, + uncachedTokens: parsed.uncachedTokens, + totalTokens: parsed.totalTokens, + costTotal: parsed.costTotal, + currentModel: + currentModel && isRecord(currentModel) + ? { provider: currentModel.provider as string, modelId: currentModel.modelId as string } + : (currentModel ?? undefined), + currentThinkingLevel: (currentThinkingLevel as ThinkingLevel | null | undefined) ?? undefined, + }; +} + +function parseEntryMaterializedPayload(row: EntryMaterializedRow): unknown { + try { + return JSON.parse(row.payload); + } catch (error) { + throw invalidSession( + `materialized entry row ${row.entry_seq} is not valid JSON`, + error instanceof Error ? error : undefined, + ); + } +} + +export function materializedStateFromRows( + summaryRow: SessionMaterializedRow, + entryRows: EntryMaterializedRow[], +): SessionMaterializedState { + const summary = parseSummary(summaryRow.payload); + const state: SessionMaterializedState = { + name: summary.name, + messageCount: summary.messageCount, + cachedTokens: summary.cachedTokens, + uncachedTokens: summary.uncachedTokens, + totalTokens: summary.totalTokens, + costTotal: summary.costTotal, + labelsById: new Map(), + modelThinkingConfigs: [], + currentModel: summary.currentModel ?? null, + currentThinkingLevel: summary.currentThinkingLevel ?? null, + }; + for (const row of entryRows) { + const payload = parseEntryMaterializedPayload(row); + if (!isRecord(payload)) throw invalidSession(`materialized entry row ${row.entry_seq} is not an object`); + if (row.type === "label") { + if (typeof payload.targetId !== "string") { + throw invalidSession(`materialized label row ${row.entry_seq} is missing targetId`); + } + if (payload.label !== null && payload.label !== undefined && typeof payload.label !== "string") { + throw invalidSession(`materialized label row ${row.entry_seq} has invalid label`); + } + const label = typeof payload.label === "string" ? payload.label.trim() : ""; + if (label) { + state.labelsById.set(payload.targetId, label); + } else { + state.labelsById.delete(payload.targetId); + } + continue; + } + if (row.type !== "label") { + } + } + return state; +} + +export function sessionStatsFromMaterializedState(state: SessionMaterializedState): SessionStats { + return { + messageCount: state.messageCount, + cachedTokens: state.cachedTokens, + uncachedTokens: state.uncachedTokens, + totalTokens: state.totalTokens, + costTotal: state.costTotal, + }; +} + +export function materializedStateValues( + sessionId: string, + state: SessionMaterializedState, +): [sessionId: string, payload: string] { + return [sessionId, serializeSummary(state)]; +} + +export function entryMaterializedValues( + entry: SessionTreeEntry, +): Array<{ type: EntryMaterializedRow["type"]; payload: string }> { + switch (entry.type) { + case "label": + return [ + { + type: "label", + payload: JSON.stringify({ targetId: entry.targetId, label: entry.label ?? null }), + }, + ]; + case "model_change": + case "thinking_level_change": + case "message": + return []; + case "active_tools_change": + case "branch_summary": + case "compaction": + case "custom": + case "custom_message": + case "leaf": + case "session_info": + return []; + default: { + const exhaustive: never = entry; + void exhaustive; + return []; + } + } +} diff --git a/packages/storage/sqlite-node/src/sqlite/storage/session-sequences.ts b/packages/storage/sqlite-node/src/sqlite/storage/session-sequences.ts new file mode 100644 index 000000000..a38bb668f --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/storage/session-sequences.ts @@ -0,0 +1,16 @@ +import type { SqliteDatabase } from "../types.ts"; +import { invalidSession } from "./shared.ts"; + +export async function getNextSequence(db: SqliteDatabase, sessionId: string): Promise { + const sequenceRow = await db + .prepare("SELECT next_seq FROM session_sequences WHERE session_id = ?") + .get<{ next_seq: number }>(sessionId); + if (!sequenceRow) { + throw invalidSession(`missing sequence row for session ${sessionId}`); + } + return sequenceRow.next_seq; +} + +export async function advanceSequence(db: SqliteDatabase, sessionId: string, nextSeq: number): Promise { + await db.prepare("UPDATE session_sequences SET next_seq = ? WHERE session_id = ?").run(nextSeq + 1, sessionId); +} diff --git a/packages/storage/sqlite-node/src/sqlite/storage/sessions.ts b/packages/storage/sqlite-node/src/sqlite/storage/sessions.ts new file mode 100644 index 000000000..2cb4d4959 --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/storage/sessions.ts @@ -0,0 +1,40 @@ +import { SessionError } from "@earendil-works/pi-agent-core"; +import type { SqliteSessionMetadata } from "../types.ts"; + +export interface SessionRow { + id: string; + created_at: string; + metadata: string | null; + cwd: string; + parent_session_id: string | null; + active_leaf_id: string | null; +} + +function parseMetadata(metadata: string | null, sessionId: string): Record | undefined { + if (metadata === null) return undefined; + let parsed: unknown; + try { + parsed = JSON.parse(metadata); + } catch (error) { + throw new SessionError( + "invalid_session", + `Invalid SQLite session ${sessionId}: metadata is not valid JSON`, + error instanceof Error ? error : undefined, + ); + } + if (typeof parsed !== "object" || parsed === null || Array.isArray(parsed)) { + throw new SessionError("invalid_session", `Invalid SQLite session ${sessionId}: metadata must be an object`); + } + return parsed as Record; +} + +export function rowToMetadata(row: SessionRow, path: string): SqliteSessionMetadata { + return { + id: row.id, + createdAt: row.created_at, + cwd: row.cwd, + path, + parentSessionId: row.parent_session_id ?? undefined, + metadata: parseMetadata(row.metadata, row.id), + }; +} diff --git a/packages/storage/sqlite-node/src/sqlite/storage/shared.ts b/packages/storage/sqlite-node/src/sqlite/storage/shared.ts new file mode 100644 index 000000000..5fc94fcb2 --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/storage/shared.ts @@ -0,0 +1,29 @@ +import type { SessionTreeEntry } from "@earendil-works/pi-agent-core"; +import { SessionError } from "@earendil-works/pi-agent-core"; +import { uuidv7 } from "@earendil-works/pi-ai"; + +export function generateEntryId(byId: { has(id: string): boolean }): string { + for (let i = 0; i < 100; i++) { + // The uuidv7 prefix is timestamp-derived and nearly constant between calls, + // so short ids must come from the random tail. + const id = uuidv7().slice(0, 8); + if (!byId.has(id)) return id; + } + return uuidv7(); +} + +export function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null; +} + +export function invalidSession(message: string, cause?: Error): SessionError { + return new SessionError("invalid_session", `Invalid SQLite session: ${message}`, cause); +} + +export function invalidEntry(message: string, cause?: Error): SessionError { + return new SessionError("invalid_entry", `Invalid SQLite session entry: ${message}`, cause); +} + +export function leafIdAfterEntry(entry: SessionTreeEntry): string | null { + return entry.type === "leaf" ? entry.targetId : entry.id; +} diff --git a/packages/storage/sqlite-node/src/sqlite/types.ts b/packages/storage/sqlite-node/src/sqlite/types.ts new file mode 100644 index 000000000..4e97ffd2f --- /dev/null +++ b/packages/storage/sqlite-node/src/sqlite/types.ts @@ -0,0 +1,55 @@ +import type { FileSystem, SessionCreateOptions, SessionMetadata, SessionRepo } from "@earendil-works/pi-agent-core"; + +/** Result of a prepared SQLite statement execution. */ +export interface SqliteRunResult { + /** Number of rows changed by the statement. */ + changes: number; + /** Inserted row id when the backend exposes one. */ + lastInsertRowid?: number; +} + +/** Prepared SQLite statement capability used by the SQLite session backend. */ +export interface SqliteStatement { + run(...params: unknown[]): Promise; + get(...params: unknown[]): Promise; + all(...params: unknown[]): Promise; +} + +/** SQLite database capability used by the SQLite session backend. */ +export interface SqliteDatabase { + exec(sql: string): Promise; + prepare(sql: string): SqliteStatement; + transaction(fn: () => Promise): Promise; + close(): Promise; +} + +export interface SqliteDatabaseFactory { + open(path: string): Promise; +} + +export interface SqliteSessionMetadata extends SessionMetadata { + cwd: string; + path: string; + parentSessionId?: string; + metadata?: Record; +} + +export interface SqliteSessionCreateOptions extends SessionCreateOptions { + cwd: string; + parentSessionId?: string; + metadata?: Record; +} + +export interface SqliteSessionListOptions { + cwd?: string; +} + +export interface SqliteSessionBackendOptions { + kind: "sqlite"; + databasePath: string; +} + +export interface SqliteSessionRepoApi + extends SessionRepo {} + +export type SqliteSessionRepoEnv = Pick; diff --git a/packages/storage/sqlite-node/tsconfig.build.json b/packages/storage/sqlite-node/tsconfig.build.json new file mode 100644 index 000000000..7cdb9e03c --- /dev/null +++ b/packages/storage/sqlite-node/tsconfig.build.json @@ -0,0 +1,11 @@ +{ + "extends": "../../../tsconfig.base.json", + "compilerOptions": { + "rootDir": "src", + "outDir": "dist", + "declaration": true, + "emitDeclarationOnly": false, + "noEmit": false + }, + "include": ["src/**/*.ts"] +} diff --git a/tsconfig.json b/tsconfig.json index e144626b0..626d6bbb4 100644 --- a/tsconfig.json +++ b/tsconfig.json @@ -12,6 +12,7 @@ "@earendil-works/pi-ai/dist/*": ["./packages/ai/src/*"], "@earendil-works/pi-agent-core": ["./packages/agent/src/index.ts"], "@earendil-works/pi-agent-core/*": ["./packages/agent/src/*"], + "@earendil-works/pi-agent-sqlite-node": ["./packages/storage/sqlite-node/src/index.ts"], "@earendil-works/pi-coding-agent": ["./packages/coding-agent/src/index.ts"], "@earendil-works/pi-coding-agent/hooks": ["./packages/coding-agent/src/core/hooks/index.ts"], "@earendil-works/pi-coding-agent/*": ["./packages/coding-agent/src/*"], @@ -24,6 +25,6 @@ "@earendil-works/pi-agent-old/*": ["./packages/agent-old/src/*"] } }, - "include": ["packages/*/src/**/*", "packages/*/test/**/*", "packages/coding-agent/examples/**/*"], + "include": ["packages/*/src/**/*", "packages/*/test/**/*", "packages/storage/*/src/**/*", "packages/storage/*/test/**/*", "packages/coding-agent/examples/**/*"], "exclude": ["**/dist/**", "packages/coding-agent/examples/extensions/gondolin/**"] } From 8e53e0e49cf3ecec20c840c4cb03e692e96c50f1 Mon Sep 17 00:00:00 2001 From: David Brailovsky Date: Tue, 21 Jul 2026 10:07:20 +0000 Subject: [PATCH 35/62] compaction & branch summarization follow retry policy fixes #6647 compaction (auto & manual) and branch summarization retry on transient failures. use the same retry policy from settings. emit events for the tui to show indication of retries --- packages/ai/src/utils/retry.ts | 110 ++++++++++ packages/ai/test/retry.test.ts | 122 ++++++++++- .../coding-agent/src/core/agent-session.ts | 54 ++++- .../core/compaction/branch-summarization.ts | 17 +- .../src/core/compaction/compaction.ts | 41 +++- .../src/modes/interactive/interactive-mode.ts | 26 +++ ...tion-retries-transient-stream-drop.test.ts | 192 ++++++++++++++++++ 7 files changed, 546 insertions(+), 16 deletions(-) create mode 100644 packages/coding-agent/test/suite/regressions/6647-compaction-retries-transient-stream-drop.test.ts diff --git a/packages/ai/src/utils/retry.ts b/packages/ai/src/utils/retry.ts index 6332ff55a..0a433439c 100644 --- a/packages/ai/src/utils/retry.ts +++ b/packages/ai/src/utils/retry.ts @@ -85,6 +85,116 @@ const RETRYABLE_PROVIDER_ERROR_PATTERN = buildProviderErrorPattern([ "ResourceExhausted", ]); +/** + * Retry policy: bounded attempts with exponential backoff (`baseDelayMs * 2^(attempt-1)`). + * Matches `settings.retry` (`enabled`, `maxRetries`, `baseDelayMs`) in coding-agent; kept + * here so the classifier and the policy-driven retry loop live together and stay reusable + * by the SDK and other callers. + */ +export interface RetryPolicy { + enabled: boolean; + /** Max retry attempts (0 = no retries). The initial call never counts as a retry. */ + maxRetries: number; + /** Base delay in ms. Per-attempt delay is `baseDelayMs * 2^(attempt-1)` before jitter. */ + baseDelayMs: number; +} + +/** Optional callbacks emitted by {@link retryAssistantCall} around each retry. */ +export interface RetryCallbacks { + /** Emitted before the backoff sleep of each retry attempt (1-indexed). */ + onRetry?: (attempt: number, maxAttempts: number, delayMs: number, errorMessage: string) => void; + /** Emitted after the backoff sleep, immediately before the retried call starts. */ + onRetryAttemptStart?: () => void; + /** Emitted once when the loop ends: success if a later call returned a non-error message. */ + onRetryEnd?: (success: boolean, attempt: number, finalError?: string) => void; +} + +class RetrySleepAbortError extends Error { + constructor() { + super("Aborted"); + } +} + +function sleep(ms: number, signal?: AbortSignal): Promise { + return new Promise((resolve, reject) => { + if (signal?.aborted) { + reject(new RetrySleepAbortError()); + return; + } + const timeout = setTimeout(resolve, ms); + signal?.addEventListener( + "abort", + () => { + clearTimeout(timeout); + reject(new RetrySleepAbortError()); + }, + { once: true }, + ); + }); +} + +/** + * Run a single assistant-producing call with bounded retry on transient errors. + * + * Behavior: + * - A non-error response (success or aborted) is returned immediately; aborts are + * never retried. Aborts during the backoff sleep are normalized to an aborted + * `AssistantMessage` too, so callers do not need to care when cancellation happened. + * - A non-retryable error (per {@link isRetryableAssistantError}, including quota/ + * billing exhaustion) is returned immediately so deterministic errors fail fast. + * - Otherwise retries up to `maxRetries` times with exponential backoff, emitting + * `onRetry` before each sleep, `onRetryAttemptStart` after each sleep before the + * retried call starts, and `onRetryEnd` once at the end (whether the loop ends in + * success, exhausted retries, or an aborted backoff). + * + * When `policy` is undefined or disabled, the first response is returned unchanged + * (equivalent to calling `produce()` directly). + */ +export async function retryAssistantCall( + produce: () => Promise, + policy: RetryPolicy | undefined, + signal: AbortSignal | undefined, + callbacks?: RetryCallbacks, +): Promise { + const maxAttempts = policy?.enabled ? policy.maxRetries : 0; + + let attempt = 0; + let lastRetry: { attempt: number; errorMessage: string } | undefined; + for (;;) { + const response = await produce(); + + // Success or abort: never retry an aborted message; non-error returns as-is. + if (response.stopReason !== "error") { + if (lastRetry) callbacks?.onRetryEnd?.(true, lastRetry.attempt); + return response; + } + + // Non-retryable, or budget exhausted: return the final error message. + if (attempt >= maxAttempts || !isRetryableAssistantError(response)) { + if (lastRetry) callbacks?.onRetryEnd?.(false, lastRetry.attempt, response.errorMessage); + return response; + } + + attempt++; + lastRetry = { attempt, errorMessage: response.errorMessage || "Unknown error" }; + const delayMs = policy!.baseDelayMs * 2 ** (attempt - 1); + callbacks?.onRetry?.(attempt, maxAttempts, delayMs, lastRetry.errorMessage); + + // Normalize aborts during retry backoff to the same AssistantMessage shape as + // provider stream aborts, so callers do not need to care when cancellation happened. + try { + await sleep(delayMs, signal); + } catch (error) { + callbacks?.onRetryEnd?.(false, attempt, lastRetry.errorMessage); + if (error instanceof RetrySleepAbortError) { + return { ...response, stopReason: "aborted", errorMessage: undefined }; + } + throw error; + } + callbacks?.onRetryAttemptStart?.(); + } +} + /** * Classifies whether a failed assistant message looks like a transient provider * or transport error, so callers can decide if the last assistant turn should be diff --git a/packages/ai/test/retry.test.ts b/packages/ai/test/retry.test.ts index abdc8716a..71743f671 100644 --- a/packages/ai/test/retry.test.ts +++ b/packages/ai/test/retry.test.ts @@ -1,6 +1,6 @@ -import { describe, expect, it } from "vitest"; +import { describe, expect, it, vi } from "vitest"; import { fauxAssistantMessage } from "../src/providers/faux.ts"; -import { isRetryableAssistantError } from "../src/utils/retry.ts"; +import { isRetryableAssistantError, type RetryPolicy, retryAssistantCall } from "../src/utils/retry.ts"; const openAIExplicitRetryMessage = "An error occurred while processing your request. You can retry your request, or contact us through our help center at help.openai.com if the error persists. Please include the request ID req_******** in your message."; @@ -66,3 +66,121 @@ describe("provider retry classification", () => { expect(isRetryableAssistantError(fauxAssistantMessage("not an error"))).toBe(false); }); }); + +describe("retryAssistantCall", () => { + const disabled: RetryPolicy = { enabled: false, maxRetries: 3, baseDelayMs: 0 }; + const enabled: RetryPolicy = { enabled: true, maxRetries: 3, baseDelayMs: 0 }; + + it("returns a successful response immediately without retrying", async () => { + const produce = vi.fn(async () => fauxAssistantMessage("ok")); + const res = await retryAssistantCall(produce, enabled, undefined); + expect(res.content).toEqual([{ type: "text", text: "ok" }]); + expect(produce).toHaveBeenCalledTimes(1); + }); + + it("does not retry an aborted message", async () => { + const produce = vi.fn(async () => fauxAssistantMessage("", { stopReason: "aborted" })); + const onRetry = vi.fn(); + const res = await retryAssistantCall(produce, enabled, undefined, { onRetry }); + expect(res.stopReason).toBe("aborted"); + expect(produce).toHaveBeenCalledTimes(1); + expect(onRetry).not.toHaveBeenCalled(); + }); + + it("does not retry a non-retryable error (quota/billing)", async () => { + const produce = vi.fn(async () => + fauxAssistantMessage("", { stopReason: "error", errorMessage: "insufficient_quota" }), + ); + const onRetry = vi.fn(); + const onRetryEnd = vi.fn(); + const res = await retryAssistantCall(produce, enabled, undefined, { onRetry, onRetryEnd }); + expect(res.stopReason).toBe("error"); + expect(produce).toHaveBeenCalledTimes(1); + expect(onRetry).not.toHaveBeenCalled(); + expect(onRetryEnd).not.toHaveBeenCalled(); + }); + + it("retries a transient error up to maxRetries then returns the final error", async () => { + const produce = vi.fn(async () => fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" })); + const onRetry = vi.fn(); + const onRetryEnd = vi.fn(); + const res = await retryAssistantCall(produce, enabled, undefined, { onRetry, onRetryEnd }); + expect(res.stopReason).toBe("error"); + expect(produce).toHaveBeenCalledTimes(4); // 1 initial + 3 retries + expect(onRetry).toHaveBeenCalledTimes(3); + expect(onRetryEnd).toHaveBeenCalledWith(false, 3, "terminated"); + }); + + it("stops retrying once a call succeeds", async () => { + let n = 0; + const produce = vi.fn(async () => { + n++; + return n < 3 + ? fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" }) + : fauxAssistantMessage("recovered"); + }); + const onRetryEnd = vi.fn(); + const res = await retryAssistantCall(produce, enabled, undefined, { onRetryEnd }); + expect(res.content).toEqual([{ type: "text", text: "recovered" }]); + expect(produce).toHaveBeenCalledTimes(3); + expect(onRetryEnd).toHaveBeenCalledWith(true, 2); + }); + + it("does not retry when policy is disabled", async () => { + const produce = vi.fn(async () => fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" })); + const onRetry = vi.fn(); + const onRetryEnd = vi.fn(); + const res = await retryAssistantCall(produce, disabled, undefined, { onRetry, onRetryEnd }); + expect(res.stopReason).toBe("error"); + expect(produce).toHaveBeenCalledTimes(1); + expect(onRetry).not.toHaveBeenCalled(); + expect(onRetryEnd).not.toHaveBeenCalled(); + }); + + it("emits onRetryAttemptStart after backoff before each retried call", async () => { + const events: string[] = []; + let n = 0; + const produce = vi.fn(async () => { + events.push(`produce:${n}`); + n++; + return n < 3 + ? fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" }) + : fauxAssistantMessage("recovered"); + }); + const onRetry = vi.fn((attempt: number) => { + events.push(`retry:${attempt}`); + }); + const onRetryAttemptStart = vi.fn(() => { + events.push("attempt-start"); + }); + const res = await retryAssistantCall(produce, enabled, undefined, { onRetry, onRetryAttemptStart }); + expect(res.content).toEqual([{ type: "text", text: "recovered" }]); + expect(onRetry).toHaveBeenCalledTimes(2); + expect(onRetryAttemptStart).toHaveBeenCalledTimes(2); + expect(events).toEqual([ + "produce:0", + "retry:1", + "attempt-start", + "produce:1", + "retry:2", + "attempt-start", + "produce:2", + ]); + }); + + it("aborts backoff sleep via signal, returns an aborted message, and emits onRetryEnd(false)", async () => { + const controller = new AbortController(); + const produce = vi.fn(async () => fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" })); + const policy: RetryPolicy = { enabled: true, maxRetries: 5, baseDelayMs: 10_000 }; + const onRetryEnd = vi.fn(); + const p = retryAssistantCall(produce, policy, controller.signal, { onRetryEnd }); + // Let one error call resolve and the first backoff sleep start, then abort. + await vi.waitFor(() => expect(produce).toHaveBeenCalled()); + controller.abort(); + const res = await p; + expect(res.stopReason).toBe("aborted"); + expect(res.errorMessage).toBeUndefined(); + expect(produce).toHaveBeenCalledTimes(1); + expect(onRetryEnd).toHaveBeenCalledWith(false, 1, "terminated"); + }); +}); diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index b0fdbb076..19adf6dd1 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -41,6 +41,7 @@ import { isContextOverflow, isRetryableAssistantError, modelsAreEqual, + type RetryCallbacks, resetApiProviders, streamSimple, } from "@earendil-works/pi-ai/compat"; @@ -161,7 +162,21 @@ export type AgentSessionEvent = errorMessage?: string; } | { type: "auto_retry_start"; attempt: number; maxAttempts: number; delayMs: number; errorMessage: string } - | { type: "auto_retry_end"; success: boolean; attempt: number; finalError?: string }; + | { type: "auto_retry_end"; success: boolean; attempt: number; finalError?: string } + | { + type: "summarization_retry_start"; + attempt: number; + maxAttempts: number; + delayMs: number; + errorMessage: string; + } + | { type: "summarization_retry_attempt_start"; source: "branchSummary" } + | { + type: "summarization_retry_attempt_start"; + source: "compaction"; + reason: "manual" | "threshold" | "overflow"; + } + | { type: "summarization_retry_end" }; /** Listener function for agent session events */ export type AgentSessionEventListener = (event: AgentSessionEvent) => void; @@ -1838,6 +1853,8 @@ export class AgentSession { this.thinkingLevel, this.agent.streamFunction, env, + this.settingsManager.getRetrySettings(), + this._summarizationRetryCallbacks({ source: "compaction", reason: "manual" }), ); summary = result.summary; firstKeptEntryId = result.firstKeptEntryId; @@ -2114,6 +2131,8 @@ export class AgentSession { this.thinkingLevel, this.agent.streamFunction, env, + this.settingsManager.getRetrySettings(), + this._summarizationRetryCallbacks({ source: "compaction", reason }), ); summary = compactResult.summary; firstKeptEntryId = compactResult.firstKeptEntryId; @@ -2620,6 +2639,37 @@ export class AgentSession { return isRetryableAssistantError(message); } + /** + * Retry policy + callbacks shared by compaction and branch-summary summarization calls. + * Uses the same `settings.retry` budget/backoff as agent-turn retries so a single transient + * stream drop no longer fails the whole operation. `source` carries the context + * the TUI needs to render the retry and recreate the underlying indicator. + */ + private _summarizationRetryCallbacks( + source: { source: "branchSummary" } | { source: "compaction"; reason: "manual" | "threshold" | "overflow" }, + ): RetryCallbacks { + return { + onRetry: (attempt, maxAttempts, delayMs, errorMessage) => { + this._emit({ + type: "summarization_retry_start", + attempt, + maxAttempts, + delayMs, + errorMessage, + }); + }, + onRetryAttemptStart: () => { + this._emit({ + type: "summarization_retry_attempt_start", + ...source, + }); + }, + onRetryEnd: () => { + this._emit({ type: "summarization_retry_end" }); + }, + }; + } + /** * Prepare a retryable error for continuation with exponential backoff. * @returns true if the caller should continue the agent, false otherwise @@ -2934,6 +2984,8 @@ export class AgentSession { replaceInstructions, reserveTokens: branchSummarySettings.reserveTokens, streamFn: this.agent.streamFunction, + retry: this.settingsManager.getRetrySettings(), + callbacks: this._summarizationRetryCallbacks({ source: "branchSummary" }), }); if (result.aborted) { return { cancelled: true, aborted: true }; diff --git a/packages/coding-agent/src/core/compaction/branch-summarization.ts b/packages/coding-agent/src/core/compaction/branch-summarization.ts index 3366f06a7..dbb1217e5 100644 --- a/packages/coding-agent/src/core/compaction/branch-summarization.ts +++ b/packages/coding-agent/src/core/compaction/branch-summarization.ts @@ -6,9 +6,9 @@ */ import type { AgentMessage, StreamFn } from "@earendil-works/pi-agent-core"; +import type { RetryCallbacks, RetryPolicy } from "@earendil-works/pi-ai"; import { contentText } from "@earendil-works/pi-ai"; import type { Model, SimpleStreamOptions, Usage } from "@earendil-works/pi-ai/compat"; -import { completeSimple } from "@earendil-works/pi-ai/compat"; import { convertToLlm, createBranchSummaryMessage, @@ -16,7 +16,7 @@ import { createCustomMessage, } from "../messages.ts"; import type { ReadonlySessionManager, SessionEntry } from "../session-manager.ts"; -import { estimateTokens } from "./compaction.ts"; +import { completeSummarization, estimateTokens } from "./compaction.ts"; import { computeFileLists, createFileOps, @@ -83,6 +83,10 @@ export interface GenerateBranchSummaryOptions { reserveTokens?: number; /** Optional session stream function. Used to preserve SDK request behavior without mutating agent state. */ streamFn?: StreamFn; + /** Retry policy for transient summarization errors. Reuses coding-agent's `settings.retry`. */ + retry?: RetryPolicy; + /** Optional callbacks for retry reporting (e.g. TUI retry indicators). */ + callbacks?: RetryCallbacks; } // ============================================================================ @@ -300,6 +304,8 @@ export async function generateBranchSummary( replaceInstructions, reserveTokens = 16384, streamFn, + retry, + callbacks, } = options; // Token budget = context window minus reserved space for prompt + response @@ -338,12 +344,11 @@ export async function generateBranchSummary( // Call LLM for summarization. Prefer the session stream function so SDK // request behavior (timeouts, retries, attribution headers) stays consistent - // without running through agent state/events. + // without running through agent state/events. Retried via completeSummarization + // so transient stream drops reuse the configured retry policy. const context = { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages }; const requestOptions: SimpleStreamOptions = { apiKey, headers, env, signal, maxTokens: 2048 }; - const response = streamFn - ? await (await streamFn(model, context, requestOptions)).result() - : await completeSimple(model, context, requestOptions); + const response = await completeSummarization(model, context, requestOptions, streamFn, retry, callbacks); // Check if aborted or errored if (response.stopReason === "aborted") { diff --git a/packages/coding-agent/src/core/compaction/compaction.ts b/packages/coding-agent/src/core/compaction/compaction.ts index ffad75e85..0ae6deeb5 100644 --- a/packages/coding-agent/src/core/compaction/compaction.ts +++ b/packages/coding-agent/src/core/compaction/compaction.ts @@ -6,7 +6,7 @@ */ import type { AgentMessage, StreamFn, ThinkingLevel } from "@earendil-works/pi-agent-core"; -import { contentText } from "@earendil-works/pi-ai"; +import { contentText, type RetryCallbacks, type RetryPolicy, retryAssistantCall } from "@earendil-works/pi-ai"; import type { AssistantMessage, Context, Model, SimpleStreamOptions, Usage } from "@earendil-works/pi-ai/compat"; import { completeSimple } from "@earendil-works/pi-ai/compat"; import { convertToLlm } from "../messages.ts"; @@ -552,17 +552,24 @@ function createSummarizationOptions( return options; } -async function completeSummarization( +/** + * Shared choke point for every compaction/branch-summary summarization call. Wraps the + * single LLM call in {@link retryAssistantCall} so transient stream drops (e.g. + * `terminated`, socket close) honor the configured retry policy instead of failing + * the whole compaction on the first attempt. Deterministic errors and aborts return + * immediately (see {@link retryAssistantCall}). + */ +export async function completeSummarization( model: Model, context: Context, options: SimpleStreamOptions, streamFn?: StreamFn, + retry?: RetryPolicy, + callbacks?: RetryCallbacks, ): Promise { - if (!streamFn) { - return completeSimple(model, context, options); - } - const stream = await streamFn(model, context, options); - return stream.result(); + const produce = async (): Promise => + streamFn ? (await streamFn(model, context, options)).result() : completeSimple(model, context, options); + return retryAssistantCall(produce, retry, options.signal, callbacks); } /** @@ -581,6 +588,8 @@ export async function generateSummary( thinkingLevel?: ThinkingLevel, streamFn?: StreamFn, env?: Record, + retry?: RetryPolicy, + callbacks?: RetryCallbacks, ): Promise { return ( await generateSummaryWithUsage( @@ -595,6 +604,8 @@ export async function generateSummary( thinkingLevel, streamFn, env, + retry, + callbacks, ) ).text; } @@ -612,6 +623,8 @@ export async function generateSummaryWithUsage( thinkingLevel?: ThinkingLevel, streamFn?: StreamFn, env?: Record, + retry?: RetryPolicy, + callbacks?: RetryCallbacks, ): Promise<{ text: string; usage: Usage }> { const maxTokens = Math.min( Math.floor(0.8 * reserveTokens), @@ -651,6 +664,8 @@ export async function generateSummaryWithUsage( { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages }, completionOptions, streamFn, + retry, + callbacks, ); if (response.stopReason === "error") { @@ -801,6 +816,8 @@ export async function compact( thinkingLevel?: ThinkingLevel, streamFn?: StreamFn, env?: Record, + retry?: RetryPolicy, + callbacks?: RetryCallbacks, ): Promise { const { firstKeptEntryId, @@ -833,6 +850,8 @@ export async function compact( thinkingLevel, streamFn, env, + retry, + callbacks, ); historyText = historyResult.text; historyUsage = historyResult.usage; @@ -847,6 +866,8 @@ export async function compact( signal, thinkingLevel, streamFn, + retry, + callbacks, ); // Merge into single summary summary = `${historyText}\n\n---\n\n**Turn Context (split turn):**\n\n${turnPrefixResult.text}`; @@ -865,6 +886,8 @@ export async function compact( thinkingLevel, streamFn, env, + retry, + callbacks, ); summary = result.text; summaryUsage = result.usage; @@ -900,6 +923,8 @@ async function generateTurnPrefixSummary( signal?: AbortSignal, thinkingLevel?: ThinkingLevel, streamFn?: StreamFn, + retry?: RetryPolicy, + callbacks?: RetryCallbacks, ): Promise<{ text: string; usage: Usage }> { const maxTokens = Math.min( Math.floor(0.5 * reserveTokens), @@ -921,6 +946,8 @@ async function generateTurnPrefixSummary( { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages }, createSummarizationOptions(model, maxTokens, apiKey, headers, env, signal, thinkingLevel), streamFn, + retry, + callbacks, ); if (response.stopReason === "error") { diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index 54415407c..89e79c615 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -3110,6 +3110,32 @@ export class InteractiveMode { this.ui.requestRender(); break; } + + case "summarization_retry_start": { + this.showError(event.errorMessage); + this.showStatusIndicator( + new RetryStatusIndicator(this.ui, event.attempt, event.maxAttempts, event.delayMs), + ); + this.ui.requestRender(); + break; + } + + case "summarization_retry_attempt_start": { + this.clearStatusIndicator("retry"); + if (event.source === "branchSummary") { + this.showStatusIndicator(new BranchSummaryStatusIndicator(this.ui)); + } else { + this.showStatusIndicator(new CompactionStatusIndicator(this.ui, event.reason)); + } + this.ui.requestRender(); + break; + } + + case "summarization_retry_end": { + this.clearStatusIndicator("retry"); + this.ui.requestRender(); + break; + } } } diff --git a/packages/coding-agent/test/suite/regressions/6647-compaction-retries-transient-stream-drop.test.ts b/packages/coding-agent/test/suite/regressions/6647-compaction-retries-transient-stream-drop.test.ts new file mode 100644 index 000000000..210b0a3b6 --- /dev/null +++ b/packages/coding-agent/test/suite/regressions/6647-compaction-retries-transient-stream-drop.test.ts @@ -0,0 +1,192 @@ +import type { StreamFn } from "@earendil-works/pi-agent-core"; +import { type AssistantMessage, createAssistantMessageEventStream, fauxAssistantMessage } from "@earendil-works/pi-ai"; +import { afterEach, describe, expect, it } from "vitest"; +import { createHarness, type Harness } from "../harness.ts"; + +/** + * Regression for #6647: compaction runs a single non-retried summarization call, so a + * transient mid-stream socket death (`terminated`) failed the whole compaction. + * Verifies that summarization now reuses `settings.retry` (bounded retries with + * exponential backoff gated on isRetryableAssistantError), emits + * `summarization_retry_*` events, and that aborts / non-retryable errors are not retried. + */ +describe("#6647 compaction retries transient summarization failures", () => { + const harnesses: Harness[] = []; + + afterEach(() => { + while (harnesses.length > 0) { + harnesses.pop()?.cleanup(); + } + }); + + function createUsage(totalTokens: number) { + return { + input: totalTokens, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }; + } + + function seedCompactableSession(harness: Harness): void { + harness.settingsManager.applyOverrides({ compaction: { keepRecentTokens: 1 } }); + const now = Date.now(); + harness.sessionManager.appendMessage({ + role: "user", + content: [{ type: "text", text: "message to compact" }], + timestamp: now - 1000, + }); + const model = harness.getModel(); + const assistant: AssistantMessage = { + ...fauxAssistantMessage("", { stopReason: "stop", timestamp: now - 500 }), + api: model.api, + provider: model.provider, + model: model.id, + usage: createUsage(100), + }; + assistant.content = [{ type: "text", text: "assistant response to compact" }]; + harness.sessionManager.appendMessage(assistant); + harness.session.agent.state.messages = harness.sessionManager.buildSessionContext().messages; + } + + /** streamFn that responds with the given sequence of assistant messages across calls. */ + function useScriptedStreamFn(harness: Harness, script: AssistantMessage[]): () => number { + let callCount = 0; + const streamFunction: StreamFn = (model) => { + const message = script[callCount] ?? script[script.length - 1]!; + callCount++; + const stream = createAssistantMessageEventStream(); + queueMicrotask(() => { + if (message.stopReason === "error" || message.stopReason === "aborted") { + stream.push({ + type: "error", + reason: message.stopReason, + error: { ...message, api: model.api, provider: model.provider, model: model.id }, + }); + } else { + stream.push({ + type: "done", + reason: message.stopReason, + message: { ...message, api: model.api, provider: model.provider, model: model.id }, + }); + } + }); + return stream; + }; + harness.session.agent.streamFunction = streamFunction; + return () => callCount; + } + + it("retries a transient `terminated` summarization error and compacts successfully", async () => { + const harness = await createHarness({ withConfiguredAuth: false }); + harnesses.push(harness); + seedCompactableSession(harness); + harness.settingsManager.applyOverrides({ retry: { enabled: true, maxRetries: 3, baseDelayMs: 0 } }); + + const model = harness.getModel(); + const error = (errorMessage: string): AssistantMessage => ({ + ...fauxAssistantMessage("", { stopReason: "error", errorMessage }), + usage: createUsage(10), + }); + const success: AssistantMessage = { + ...fauxAssistantMessage("recovered summary"), + usage: createUsage(10), + }; + const getCallCount = useScriptedStreamFn(harness, [error("terminated"), error("terminated"), success]); + + const result = await harness.session.compact(); + + expect(result.summary).toContain("recovered summary"); + expect(getCallCount()).toBe(3); // 1 initial + 2 retries + const starts = harness.eventsOfType("summarization_retry_start"); + const ends = harness.eventsOfType("summarization_retry_end"); + expect(starts).toHaveLength(2); + expect(ends).toHaveLength(1); + expect(starts[0]).toMatchObject({ attempt: 1, maxAttempts: 3, errorMessage: "terminated" }); + expect(starts[1]).toMatchObject({ attempt: 2, maxAttempts: 3 }); + expect(ends[0]).toMatchObject({ type: "summarization_retry_end" }); + // model.* referenced to keep imports honest + expect(model.id).toBeTruthy(); + }); + + it("does not retry a non-retryable error (insufficient_quota)", async () => { + const harness = await createHarness({ withConfiguredAuth: false }); + harnesses.push(harness); + seedCompactableSession(harness); + harness.settingsManager.applyOverrides({ retry: { enabled: true, maxRetries: 3, baseDelayMs: 0 } }); + + const error: AssistantMessage = { + ...fauxAssistantMessage("", { stopReason: "error", errorMessage: "insufficient_quota" }), + usage: createUsage(10), + }; + const getCallCount = useScriptedStreamFn(harness, [error]); + + await expect(harness.session.compact()).rejects.toThrow("insufficient_quota"); + expect(getCallCount()).toBe(1); + expect(harness.eventsOfType("summarization_retry_start")).toHaveLength(0); + }); + + it("does not retry when retry is disabled", async () => { + const harness = await createHarness({ withConfiguredAuth: false }); + harnesses.push(harness); + seedCompactableSession(harness); + harness.settingsManager.applyOverrides({ retry: { enabled: false, maxRetries: 3, baseDelayMs: 0 } }); + + const error: AssistantMessage = { + ...fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" }), + usage: createUsage(10), + }; + const getCallCount = useScriptedStreamFn(harness, [error]); + + await expect(harness.session.compact()).rejects.toThrow("terminated"); + expect(getCallCount()).toBe(1); + expect(harness.eventsOfType("summarization_retry_start")).toHaveLength(0); + }); + + it("stops retrying after maxRetries and reports failure", async () => { + const harness = await createHarness({ withConfiguredAuth: false }); + harnesses.push(harness); + seedCompactableSession(harness); + harness.settingsManager.applyOverrides({ retry: { enabled: true, maxRetries: 2, baseDelayMs: 0 } }); + + const error: AssistantMessage = { + ...fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" }), + usage: createUsage(10), + }; + const getCallCount = useScriptedStreamFn(harness, [error, error, error]); + + await expect(harness.session.compact()).rejects.toThrow("terminated"); + expect(getCallCount()).toBe(3); // 1 initial + 2 retries + const starts = harness.eventsOfType("summarization_retry_start"); + const ends = harness.eventsOfType("summarization_retry_end"); + expect(starts).toHaveLength(2); + expect(ends).toHaveLength(1); + expect(ends[0]).toMatchObject({ type: "summarization_retry_end" }); + }); + + it("aborts an in-flight retry backoff via abortCompaction", async () => { + const harness = await createHarness({ withConfiguredAuth: false }); + harnesses.push(harness); + seedCompactableSession(harness); + harness.settingsManager.applyOverrides({ retry: { enabled: true, maxRetries: 5, baseDelayMs: 30_000 } }); + + const error: AssistantMessage = { + ...fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" }), + usage: createUsage(10), + }; + useScriptedStreamFn(harness, [error, error, error]); + + const compactPromise = harness.session.compact(); + // Let the first error resolve and the retry backoff sleep start. + await new Promise((resolve) => setTimeout(resolve, 0)); + harness.session.abortCompaction(); + + // The aborted retry backoff rejects with an AbortError (matching the real SDK + // abort path), which compaction classifies as aborted. + await expect(compactPromise).rejects.toThrow(); + const compactionEnd = harness.eventsOfType("compaction_end").at(-1); + expect(compactionEnd).toMatchObject({ aborted: true }); + }); +}); From 8495f9d0d6407d4ec94e16a685df70740335dd29 Mon Sep 17 00:00:00 2001 From: Cristina Poncela Cubeiro <140309543+cristinaponcela@users.noreply.github.com> Date: Tue, 21 Jul 2026 13:32:42 +0200 Subject: [PATCH 36/62] chore: rename orchestrator to server (#6898) --- package-lock.json | 22 ++++++------------- package.json | 4 ++-- .../{orchestrator => server}/CHANGELOG.md | 4 ++++ packages/{orchestrator => server}/README.md | 6 ++--- .../{orchestrator => server}/package.json | 11 ++++++---- packages/{orchestrator => server}/src/cli.ts | 10 ++++----- .../{orchestrator => server}/src/config.ts | 16 +++++++------- .../{orchestrator => server}/src/handler.ts | 8 +++---- .../{orchestrator => server}/src/index.ts | 0 .../src/ipc/client.ts | 8 +++---- .../src/ipc/protocol.ts | 16 +++++++------- .../src/ipc/server.ts | 8 +++---- .../{orchestrator => server}/src/radius.ts | 16 +++++++------- .../src/rpc-process.ts | 2 +- .../{orchestrator => server}/src/serve.ts | 6 ++--- .../{orchestrator => server}/src/storage.ts | 14 ++++++------ .../src/supervisor.ts | 4 ++-- .../{orchestrator => server}/src/types.ts | 0 .../tsconfig.build.json | 0 tsconfig.json | 4 ++-- 20 files changed, 79 insertions(+), 80 deletions(-) rename packages/{orchestrator => server}/CHANGELOG.md (70%) rename packages/{orchestrator => server}/README.md (67%) rename packages/{orchestrator => server}/package.json (82%) rename packages/{orchestrator => server}/src/cli.ts (84%) rename packages/{orchestrator => server}/src/config.ts (79%) rename packages/{orchestrator => server}/src/handler.ts (93%) rename packages/{orchestrator => server}/src/index.ts (100%) rename packages/{orchestrator => server}/src/ipc/client.ts (73%) rename packages/{orchestrator => server}/src/ipc/protocol.ts (80%) rename packages/{orchestrator => server}/src/ipc/server.ts (96%) rename packages/{orchestrator => server}/src/radius.ts (96%) rename packages/{orchestrator => server}/src/rpc-process.ts (98%) rename packages/{orchestrator => server}/src/serve.ts (91%) rename packages/{orchestrator => server}/src/storage.ts (84%) rename packages/{orchestrator => server}/src/supervisor.ts (99%) rename packages/{orchestrator => server}/src/types.ts (100%) rename packages/{orchestrator => server}/tsconfig.build.json (100%) diff --git a/package-lock.json b/package-lock.json index 2799f6d7c..5b9e30244 100644 --- a/package-lock.json +++ b/package-lock.json @@ -798,25 +798,14 @@ "resolved": "packages/coding-agent", "link": true }, - "node_modules/@earendil-works/pi-orchestrator": { - "resolved": "packages/orchestrator", + "node_modules/@earendil-works/pi-server": { + "resolved": "packages/server", "link": true }, "node_modules/@earendil-works/pi-tui": { "resolved": "packages/tui", "link": true }, - "node_modules/@emnapi/wasi-threads": { - "version": "1.2.1", - "resolved": "https://registry.npmjs.org/@emnapi/wasi-threads/-/wasi-threads-1.2.1.tgz", - "integrity": "sha512-uTII7OYF+/Mes/MrcIOYp5yOtSMLBWSIoLPpcgwipoiKbli6k322tcoFsxoIIxPDqW01SQGAgko4EzZi2BNv2w==", - "dev": true, - "license": "MIT", - "optional": true, - "dependencies": { - "tslib": "^2.4.0" - } - }, "node_modules/@esbuild/aix-ppc64": { "version": "0.28.1", "resolved": "https://registry.npmjs.org/@esbuild/aix-ppc64/-/aix-ppc64-0.28.1.tgz", @@ -6086,13 +6075,16 @@ } } }, - "packages/orchestrator": { - "name": "@earendil-works/pi-orchestrator", + "packages/server": { + "name": "@earendil-works/pi-server", "version": "0.80.10", "license": "MIT", "dependencies": { "@earendil-works/pi-coding-agent": "^0.80.10" }, + "bin": { + "server": "dist/cli.js" + }, "devDependencies": { "shx": "0.4.0" }, diff --git a/package.json b/package.json index 8cdca8557..4f710af98 100644 --- a/package.json +++ b/package.json @@ -13,8 +13,8 @@ ], "scripts": { "clean": "npm run clean --workspaces", - "build": "cd packages/tui && npm run build && cd ../ai && npm run build && cd ../agent && npm run build && cd ../storage/sqlite-node && npm run build && cd ../../coding-agent && npm run build && cd ../orchestrator && npm run build", - "build:offline": "cd packages/tui && npm run build && cd ../ai && npm run build:offline && cd ../agent && npm run build && cd ../storage/sqlite-node && npm run build && cd ../../coding-agent && npm run build && cd ../orchestrator && npm run build", + "build": "cd packages/tui && npm run build && cd ../ai && npm run build && cd ../agent && npm run build && cd ../storage/sqlite-node && npm run build && cd ../../coding-agent && npm run build && cd ../server && npm run build", + "build:offline": "cd packages/tui && npm run build && cd ../ai && npm run build:offline && cd ../agent && npm run build && cd ../storage/sqlite-node && npm run build && cd ../../coding-agent && npm run build && cd ../server && npm run build", "check": "biome check --write --error-on-warnings . && npm run check:pinned-deps && npm run check:ts-imports && npm run check:shrinkwrap && npm run check:install-lock:coding-agent && tsgo --noEmit && npm run check:browser-smoke", "check:browser-smoke": "node scripts/check-browser-smoke.mjs", "check:pinned-deps": "node scripts/check-pinned-deps.mjs", diff --git a/packages/orchestrator/CHANGELOG.md b/packages/server/CHANGELOG.md similarity index 70% rename from packages/orchestrator/CHANGELOG.md rename to packages/server/CHANGELOG.md index 9762afa6d..2fdaaac04 100644 --- a/packages/orchestrator/CHANGELOG.md +++ b/packages/server/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Changed + +- Renamed the orchestrator workspace package and internal server references to server. + ## [0.80.10] - 2026-07-16 ## [0.80.9] - 2026-07-16 diff --git a/packages/orchestrator/README.md b/packages/server/README.md similarity index 67% rename from packages/orchestrator/README.md rename to packages/server/README.md index aabd72436..2e8e877b7 100644 --- a/packages/orchestrator/README.md +++ b/packages/server/README.md @@ -1,11 +1,11 @@ -# @earendil-works/pi-orchestrator +# @earendil-works/pi-server Experimental. This package is under active development and may change or be removed without notice. Its CLI, APIs, and behavior are not yet stable. -Orchestrator package for pi. +Server package for pi. ## CLI ```bash -orchestrator --help +server --help ``` diff --git a/packages/orchestrator/package.json b/packages/server/package.json similarity index 82% rename from packages/orchestrator/package.json rename to packages/server/package.json index c1d7e95af..7c039af37 100644 --- a/packages/orchestrator/package.json +++ b/packages/server/package.json @@ -1,7 +1,7 @@ { - "name": "@earendil-works/pi-orchestrator", + "name": "@earendil-works/pi-server", "version": "0.80.10", - "description": "experimental orchestrator package for pi", + "description": "experimental server package for pi", "type": "module", "main": "./dist/index.js", "types": "./dist/index.d.ts", @@ -11,6 +11,9 @@ "import": "./dist/index.js" } }, + "bin": { + "server": "./dist/cli.js" + }, "files": [ "dist", "README.md", @@ -24,14 +27,14 @@ }, "keywords": [ "pi", - "orchestrator" + "server" ], "author": "Earendil Works", "license": "MIT", "repository": { "type": "git", "url": "git+https://github.com/earendil-works/pi.git", - "directory": "packages/orchestrator" + "directory": "packages/server" }, "engines": { "node": ">=22.19.0" diff --git a/packages/orchestrator/src/cli.ts b/packages/server/src/cli.ts similarity index 84% rename from packages/orchestrator/src/cli.ts rename to packages/server/src/cli.ts index 856dd6584..738d416ed 100644 --- a/packages/orchestrator/src/cli.ts +++ b/packages/server/src/cli.ts @@ -18,7 +18,7 @@ const packageJson = JSON.parse(readFileSync(join(__dirname, "../package.json"), function printHelp(): void { console.log( - `orchestrator v${packageJson.version}\n\nUsage:\n orchestrator serve\n orchestrator list\n orchestrator spawn [--cwd ] [--label