diff --git a/index.ts b/index.ts index 36fa7bd..d263138 100644 --- a/index.ts +++ b/index.ts @@ -5,6 +5,8 @@ * https://api.commandcode.ai/provider/v1 */ +import { AssistantMessageEventStream } from "@earendil-works/pi-ai" +import { streamSimple as streamNativeProvider } from "@earendil-works/pi-ai/compat" import { getAgentDir, type ExtensionAPI, @@ -14,6 +16,8 @@ import { import { join } from "node:path" import { getConfiguredApiKey } from "./src/api-key.ts" +import { createStreamCommandCode } from "./src/core.ts" +import { calculateCommandCodeCost } from "./src/cost.ts" import { baseUrlForModel, DEFAULT_MODELS_URL, @@ -25,8 +29,10 @@ import { type CommandCodeModel, } from "./src/models.ts" import { getApiKey as getOAuthApiKey, login, refreshToken } from "./src/oauth.ts" +import { normalizeCommandCodeMessage } from "./src/overflow.ts" import { MODEL_COSTS, ZERO_MODEL_COST } from "./src/pricing.ts" import { createCommandCodeRuntime } from "./src/runtime.ts" +import { createCommandCodeTransportRouter } from "./src/transport.ts" function commandCodeHeaders(): Record | undefined { if (process.env.COMMANDCODE_ZDR === "1") { @@ -38,6 +44,7 @@ function commandCodeHeaders(): Record | undefined { function createProviderConfig( models: readonly CommandCodeModel[], apiBase: string, + streamCommandCode: ProviderConfig["streamSimple"], ): ProviderConfig { const headers = commandCodeHeaders() return { @@ -45,6 +52,7 @@ function createProviderConfig( baseUrl: apiBase, apiKey: getConfiguredApiKey() ?? "$COMMANDCODE_API_KEY", api: "openai-completions", + streamSimple: streamCommandCode, headers, oauth: { name: "Command Code", @@ -82,12 +90,32 @@ function createProviderConfig( } } +function legacyApiBase(providerApiBase: string): string { + return providerApiBase.replace(/\/provider\/v1\/?$/, "") +} + export default async function (pi: ExtensionAPI) { const apiBase = process.env.COMMANDCODE_API_BASE ?? DEFAULT_PROVIDER_API_BASE const modelsUrl = process.env.COMMANDCODE_MODELS_URL ?? DEFAULT_MODELS_URL const modelsTimeoutMs = getModelsTimeoutMs() const modelsCachePath = process.env.COMMANDCODE_MODELS_CACHE ?? join(getAgentDir(), "commandcode-models.json") + const streamGenerate = createStreamCommandCode({ + createStream: () => new AssistantMessageEventStream(), + calculateCost: calculateCommandCodeCost, + apiBase: legacyApiBase(apiBase), + }) + const transport = createCommandCodeTransportRouter({ + createStream: () => new AssistantMessageEventStream(), + streamProvider: streamNativeProvider, + streamGenerate, + }) + + pi.on("message_end", async (event, ctx) => { + if (event.message.role !== "assistant") return + const normalized = normalizeCommandCodeMessage(event.message, ctx.model?.provider) + return normalized ? { message: normalized.message } : undefined + }) const runtime = createCommandCodeRuntime(pi, { endpoint: modelsUrl, @@ -98,7 +126,7 @@ export default async function (pi: ExtensionAPI) { cachePath: modelsCachePath, timeoutMs: modelsTimeoutMs, }), - createProviderConfig: (models) => createProviderConfig(models, apiBase), + createProviderConfig: (models) => createProviderConfig(models, apiBase, transport.stream), }) await runtime.initialize() diff --git a/package.json b/package.json index 9af81b6..c99a0aa 100644 --- a/package.json +++ b/package.json @@ -29,18 +29,23 @@ "LICENSE" ], "scripts": { - "test": "npm run typecheck && tsx tests/test-package-manifest.ts && tsx tests/test-api-key.ts && tsx tests/test-models.ts && tsx tests/test-runtime.ts && tsx tests/test-pricing.ts && tsx tests/test-oauth.ts && node tests/test-pi-isolated.mjs && node tests/test-pi-authenticated.mjs && node tests/test-pi-local.mjs && node tests/test-omp-compat.mjs", + "test": "npm run typecheck && tsx tests/test-package-manifest.ts && tsx tests/test-api-key.ts && tsx tests/test-pure-functions.ts && tsx tests/test-models.ts && tsx tests/test-runtime.ts && tsx tests/test-pricing.ts && tsx tests/test-cost.ts && tsx tests/test-oauth.ts && tsx tests/test-abort.ts && tsx tests/test-overflow.ts && tsx tests/test-stream.ts && tsx tests/test-retry.ts && tsx tests/test-transport.ts && node tests/test-pi-isolated.mjs && node tests/test-pi-authenticated.mjs && node tests/test-pi-local.mjs && node tests/test-omp-compat.mjs", "typecheck": "tsc --noEmit", "format:check": "prettier --check '**/*.{ts,mjs,json,md}'", "format": "prettier --write '**/*.{ts,mjs,json,md}'", "pi:isolated": "node scripts/pi-isolated.mjs", "pi:authenticated": "node scripts/pi-authenticated.mjs", - "test:unit": "tsx tests/test-api-key.ts && tsx tests/test-models.ts && tsx tests/test-runtime.ts && tsx tests/test-pricing.ts && tsx tests/test-oauth.ts", + "test:unit": "tsx tests/test-api-key.ts && tsx tests/test-pure-functions.ts && tsx tests/test-models.ts && tsx tests/test-runtime.ts && tsx tests/test-pricing.ts && tsx tests/test-cost.ts && tsx tests/test-oauth.ts && tsx tests/test-abort.ts && tsx tests/test-overflow.ts && tsx tests/test-stream.ts && tsx tests/test-retry.ts && tsx tests/test-transport.ts", "test:api-key": "tsx tests/test-api-key.ts", "test:models": "tsx tests/test-models.ts", "test:runtime": "tsx tests/test-runtime.ts", "test:pricing": "tsx tests/test-pricing.ts", "test:oauth": "tsx tests/test-oauth.ts", + "test:abort": "tsx tests/test-abort.ts", + "test:overflow": "tsx tests/test-overflow.ts", + "test:stream": "tsx tests/test-stream.ts", + "test:retry": "tsx tests/test-retry.ts", + "test:transport": "tsx tests/test-transport.ts", "test:pi-isolated": "node tests/test-pi-isolated.mjs", "test:pi-authenticated": "node tests/test-pi-authenticated.mjs", "test:pi-local": "node tests/test-pi-local.mjs", diff --git a/src/converters.ts b/src/converters.ts new file mode 100644 index 0000000..4012fbb --- /dev/null +++ b/src/converters.ts @@ -0,0 +1,296 @@ +import { existsSync, readFileSync } from "node:fs" +import { homedir } from "node:os" +import { join } from "node:path" + +import type { MessageLike, StopReason, ToolLike } from "./types.ts" +import { toJsonSchema } from "./json-schema.ts" + +export { toJsonSchema } from "./json-schema.ts" + +export function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value) +} + +export function stringValue(value: unknown): string | undefined { + return typeof value === "string" ? value : undefined +} + +export function recordArray(value: unknown): readonly Record[] { + if (!Array.isArray(value)) return [] + return value.filter(isRecord) +} + +export function recordOrEmpty(value: unknown): Record { + if (isRecord(value)) return value + if (typeof value === "string") { + try { + const parsed: unknown = JSON.parse(value) + if (isRecord(parsed)) return parsed + } catch { + // Some providers stream incomplete JSON argument fragments. + } + } + return {} +} + +export function numberValue(value: unknown): number | undefined { + return typeof value === "number" && Number.isFinite(value) ? value : undefined +} + +function defaultAuthPaths(home: string): string[] { + return [ + join(home, ".commandcode", "auth.json"), + join(home, ".omp", "agent", "auth.json"), + join(home, ".pi", "agent", "auth.json"), + ] +} + +function apiKeyFromCredentialRecord(value: unknown): string | undefined { + if (!isRecord(value)) return undefined + + const type = stringValue(value.type) + if (type === "api") return stringValue(value.key) + if (type === "oauth") return stringValue(value.access) + + return stringValue(value.key) ?? stringValue(value.access) +} + +function imageParts(value: unknown): readonly Record[] { + if (isRecord(value)) return value.type === "image" ? [value] : [] + return recordArray(value).filter((part) => part.type === "image") +} + +function imageContentError(role: string): Error { + return new Error(`Selected Command Code model does not support image content in ${role}`) +} + +export function assertTextOnlyMessages(messages?: readonly MessageLike[]): void { + for (const message of messages ?? []) { + if (imageParts(message.content).length > 0) { + const role = message.role === "toolResult" ? "tool results" : `${message.role} messages` + throw imageContentError(role) + } + } +} + +function imageToCommandCode(part: Record): Record { + const data = stringValue(part.data) + const mimeType = stringValue(part.mimeType) + if (!data || !mimeType) + throw new Error("Invalid image content: expected base64 data and mimeType") + + return { + type: "image", + image: `data:${mimeType};base64,${data}`, + mimeType, + } +} + +function userContentToCommandCode(content: unknown, allowImages: boolean): unknown { + if (typeof content === "string") return content + + return recordArray(content).flatMap((part) => { + if (part.type === "text") return [{ type: "text", text: stringValue(part.text) ?? "" }] + if (part.type === "image") { + if (!allowImages) throw imageContentError("user messages") + return [imageToCommandCode(part)] + } + return [] + }) +} + +export function getApiKey( + options: { + env?: NodeJS.ProcessEnv + authPaths?: readonly string[] + homeDir?: () => string + } = {}, +): string | undefined { + const env = options.env ?? process.env + if (env.COMMANDCODE_API_KEY) return env.COMMANDCODE_API_KEY + + const home = options.homeDir?.() ?? homedir() + const authPaths = options.authPaths ?? defaultAuthPaths(home) + + for (const authPath of authPaths) { + try { + if (!existsSync(authPath)) continue + const parsed: unknown = JSON.parse(readFileSync(authPath, "utf-8")) + if (!isRecord(parsed)) continue + + // Legacy: direct apiKey or commandcode field. + const apiKey = stringValue(parsed.apiKey) + if (apiKey) return apiKey + const commandcode = stringValue(parsed.commandcode) + if (commandcode) return commandcode + + // pi stores OAuth credentials as {"commandcode": {"type":"oauth","access":"..."}}. + // The official Command Code CLI stores API credentials under "command-code". + const providerKey = + apiKeyFromCredentialRecord(parsed.commandcode) ?? + apiKeyFromCredentialRecord(parsed["command-code"]) + if (providerKey) return providerKey + } catch { + // Ignore malformed or unreadable auth files. + } + } + + return undefined +} + +export function textContent(message: { content?: unknown }): string { + return recordArray(message.content) + .filter((part) => part.type === "text") + .map((part) => stringValue(part.text) ?? "") + .join("\n") +} + +export function getEnvironmentInfo(): string { + return `${process.platform}-${process.arch}, Node.js ${process.version}` +} + +export function toolsToJson(tools?: readonly ToolLike[]): unknown[] { + if (!tools) return [] + return tools.map((tool) => ({ + type: "function", + name: tool.name, + description: tool.description, + input_schema: tool.parameters ? toJsonSchema(tool.parameters) : {}, + })) +} + +function completeToolCallIds(messages?: readonly MessageLike[]): Set { + const callIds = new Set() + const resultIds = new Set() + + for (const message of messages ?? []) { + if (message.role === "assistant") { + for (const content of recordArray(message.content)) { + if (content.type === "toolCall") { + const id = stringValue(content.id) + if (id) callIds.add(id) + } + } + } else if (message.role === "toolResult") { + if (message.toolCallId) resultIds.add(message.toolCallId) + } + } + + return new Set([...callIds].filter((id) => resultIds.has(id))) +} + +export function messagesToCC( + messages?: readonly MessageLike[], + options: { allowImages?: boolean } = {}, +): unknown[] { + const allowImages = options.allowImages ?? false + if (!allowImages) assertTextOnlyMessages(messages) + + const out: unknown[] = [] + const pairedToolCallIds = completeToolCallIds(messages) + + for (const message of messages ?? []) { + if (message.role === "user") { + out.push({ + role: "user", + content: userContentToCommandCode(message.content, allowImages), + }) + } else if (message.role === "assistant") { + const parts: unknown[] = [] + for (const content of recordArray(message.content)) { + if (content.type === "text") { + parts.push({ type: "text", text: stringValue(content.text) ?? "" }) + } else if (content.type === "toolCall") { + const toolCallId = stringValue(content.id) ?? "" + if (!pairedToolCallIds.has(toolCallId)) continue + parts.push({ + type: "tool-call", + toolCallId, + toolName: stringValue(content.name) ?? "", + input: recordOrEmpty(content.arguments), + }) + } + } + if (parts.length > 0) out.push({ role: "assistant", content: parts }) + } else if (message.role === "toolResult") { + if (!message.toolCallId || !pairedToolCallIds.has(message.toolCallId)) continue + out.push({ + role: "tool", + content: [ + { + type: "tool-result", + toolCallId: message.toolCallId, + toolName: message.toolName, + output: message.isError + ? { type: "error-text", value: textContent(message) } + : { type: "text", value: textContent(message) }, + }, + ], + }) + + const images = imageParts(message.content) + if (images.length > 0) { + if (!allowImages) throw imageContentError("tool results") + out.push({ + role: "user", + content: images.map(imageToCommandCode), + }) + } + } + } + return out +} + +export function parseStreamEventLine(line: string): unknown | undefined { + let trimmed = line.trim() + if (!trimmed || trimmed.startsWith(":") || trimmed.startsWith("event:")) return undefined + if (trimmed.startsWith("data:")) trimmed = trimmed.slice(5).trim() + if (!trimmed || trimmed === "[DONE]") return undefined + + try { + const parsed: unknown = JSON.parse(trimmed) + return parsed + } catch { + return undefined + } +} + +export function mapFinishReason(reason: unknown): StopReason { + if (reason === "tool-calls") return "toolUse" + if ( + reason === "length" || + reason === "max_tokens" || + reason === "max-tokens" || + reason === "max_output_tokens" + ) { + return "length" + } + return "stop" +} + +function promptPartToText(value: unknown, depth = 0): string { + if (depth > 10) return "" + if (typeof value === "string") return value + if (Array.isArray(value)) + return value + .map((v) => promptPartToText(v, depth + 1)) + .filter(Boolean) + .join("\n") + if (!isRecord(value)) return "" + const text = stringValue(value.text) + if (text) return text + const content = promptPartToText(value.content, depth + 1) + if (content) return content + return "" +} + +export function systemPromptToText(value: unknown): string { + if (value === undefined || value === null) return "" + if (typeof value === "string") return value + if (Array.isArray(value)) + return value + .map((v) => promptPartToText(v, 0)) + .filter(Boolean) + .join("\n\n") + return promptPartToText(value, 0) +} diff --git a/src/core.ts b/src/core.ts new file mode 100644 index 0000000..f1898dd --- /dev/null +++ b/src/core.ts @@ -0,0 +1,741 @@ +/** + * Testable Command Code provider core. + * + * The runtime imports live in index.ts; this module takes injected stream/cost + * dependencies so tests can exercise the real serialization and stream parser. + */ + +import { randomUUID } from "node:crypto" + +import { commandCodeErrorMessage, redactCommandCodeErrorText } from "./overflow.ts" +import { modelSupportsImageInput } from "./models.ts" +import { + getApiKey, + getEnvironmentInfo, + isRecord, + assertTextOnlyMessages, + mapFinishReason, + messagesToCC, + numberValue, + parseStreamEventLine, + recordOrEmpty, + stringValue, + toolsToJson, + systemPromptToText, +} from "./converters.ts" +import type { + AssistantMessageEventStreamLike, + AssistantMessageLike, + ContextLike, + CoreDependencies, + ErrorReason, + ModelLike, + StopReason, + StreamOptions, + TerminalReason, + TextContent, + ToolCallContent, + Usage, +} from "./types.ts" + +export * from "./converters.ts" +export * from "./overflow.ts" +export * from "./types.ts" + +export const DEFAULT_API_BASE = "https://api.commandcode.ai" +export const COMMAND_CODE_CLI_VERSION = "1.15.1" + +const DEFAULT_GENERATE_MAX_TOKENS = 64_000 +const DEFAULT_MAX_RETRIES = 0 +const DEFAULT_MAX_RETRY_DELAY_MS = 60_000 +const BASE_RETRY_DELAY_MS = 500 + +function isRetryableStatus(status: number): boolean { + return status === 429 || (status >= 500 && status < 600) +} + +function parseRetryAfterSeconds(value: string | null): number | undefined { + if (!value) return undefined + const seconds = Number(value) + if (Number.isFinite(seconds) && seconds >= 0) return seconds + const date = Date.parse(value) + if (!Number.isNaN(date)) return Math.max(0, (date - Date.now()) / 1000) + return undefined +} + +function effectiveMaxRetryDelayMs(value: number | undefined): number { + if (value === undefined) return DEFAULT_MAX_RETRY_DELAY_MS + if (value === 0) return Number.POSITIVE_INFINITY + return value +} + +function retryDelayMs( + attempt: number, + retryAfterHeader: string | null, + maxDelayMs: number, +): number { + const retryAfterMs = parseRetryAfterSeconds(retryAfterHeader) + if (retryAfterMs !== undefined) { + if (retryAfterMs * 1000 > maxDelayMs) return -1 + return retryAfterMs * 1000 + } + const exponential = BASE_RETRY_DELAY_MS * 2 ** attempt + const jitter = exponential * 0.2 * Math.random() + return Math.min(exponential + jitter, maxDelayMs) +} + +function defaultUsage(): Usage { + return { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + } +} + +function commandCodeUsage(event: Record): Record | undefined { + return isRecord(event.totalUsage) ? event.totalUsage : undefined +} + +function commandCodeInputTokenDetails( + usage: Record, +): Record | undefined { + return isRecord(usage.inputTokenDetails) ? usage.inputTokenDetails : undefined +} + +function headersToRecord(headers: Headers): Record { + const out: Record = {} + headers.forEach((value, key) => { + out[key] = value + }) + return out +} + +function abortError(message = "The operation was aborted"): DOMException { + return new DOMException(message, "AbortError") +} + +function timeoutError(timeoutMs: number | undefined): Error { + return new Error( + timeoutMs === undefined + ? "Command Code API request timed out" + : `Command Code API request timed out after ${timeoutMs}ms`, + ) +} + +function successStopReason(reason: TerminalReason): StopReason { + if (reason === "length" || reason === "toolUse") return reason + return "stop" +} + +function generateMaxTokens(model: ModelLike, options?: StreamOptions): number { + return Math.min( + options?.maxTokens ?? model.maxTokens, + model.maxTokens, + DEFAULT_GENERATE_MAX_TOKENS, + ) +} + +function mappedReasoningEffort(model: ModelLike, options?: StreamOptions): string | undefined { + const level = options?.reasoning + if (!level || level === "off" || !model.reasoning) return undefined + + const effortMap = model.thinking?.effortMap ?? model.thinkingLevelMap + const mapped = effortMap?.[level] + return typeof mapped === "string" && mapped !== "off" ? mapped : undefined +} + +export function projectSlugFromPath(pathName: string): string { + const slug = pathName + .toLowerCase() + .replace(/^[a-z]:/i, "") + .replace(/[^a-z0-9]+/g, "-") + .replace(/^-+|-+$/g, "") + return slug || "project" +} + +export function createStreamCommandCode(deps: CoreDependencies) { + const apiBase = deps.apiBase ?? DEFAULT_API_BASE + const fetchImpl = deps.fetchImpl ?? fetch + const cwd = deps.cwd ?? (() => process.cwd()) + const now = deps.now ?? (() => Date.now()) + const uuid = deps.uuid ?? (() => randomUUID()) + const delay = + deps.delay ?? + ((ms: number, signal: AbortSignal) => { + if (signal.aborted) return Promise.reject(abortError()) + return new Promise((resolve, reject) => { + const id = setTimeout(() => { + signal.removeEventListener("abort", onAbort) + resolve() + }, ms) + const onAbort = () => { + clearTimeout(id) + reject(abortError()) + } + signal.addEventListener("abort", onAbort, { once: true }) + }) + }) + + function raceAbort(promise: Promise, signal: AbortSignal): Promise { + if (signal.aborted) return Promise.reject(abortError()) + + return new Promise((resolve, reject) => { + const onAbort = () => reject(abortError()) + signal.addEventListener("abort", onAbort, { once: true }) + promise.then( + (value) => { + signal.removeEventListener("abort", onAbort) + resolve(value) + }, + (error: unknown) => { + signal.removeEventListener("abort", onAbort) + reject(error) + }, + ) + }) + } + + function raceAbortWithTimeout( + promise: Promise, + controller: AbortController, + timeoutMs: number | undefined, + ): Promise { + if (timeoutMs === undefined) return raceAbort(promise, controller.signal) + + return new Promise((resolve, reject) => { + const timer = setTimeout(() => { + controller.abort() + reject(timeoutError(timeoutMs)) + }, timeoutMs) + raceAbort(promise, controller.signal).then( + (value) => { + clearTimeout(timer) + resolve(value) + }, + (error: unknown) => { + clearTimeout(timer) + reject(error) + }, + ) + }) + } + + return function streamCommandCode( + model: ModelLike, + context: ContextLike, + options?: StreamOptions, + ): AssistantMessageEventStreamLike { + const stream = deps.createStream() + + async function run() { + // OMP may pass the legacy env-var name "COMMANDCODE_API_KEY" (old pi) + // or "$COMMANDCODE_API_KEY" (new pi) as the apiKey value instead of + // resolving it. Filter out these specific strings. + const LEGACY_API_KEY_REF = "$COMMANDCODE_API_KEY" + const OLD_API_KEY_REF = "COMMANDCODE_API_KEY" + const hostKey = + options?.apiKey && + options.apiKey !== LEGACY_API_KEY_REF && + options.apiKey !== OLD_API_KEY_REF + ? options.apiKey + : undefined + + const apiKey = + hostKey ?? + getApiKey({ + env: deps.env, + authPaths: deps.authPaths, + homeDir: deps.homeDir, + }) + + if (!apiKey) { + const msg: AssistantMessageLike = { + role: "assistant", + content: [], + api: model.api, + provider: model.provider, + model: model.id, + usage: defaultUsage(), + stopReason: "error", + errorMessage: + "No Command Code API key. Run /login and select Command Code, set the COMMANDCODE_API_KEY env var, or configure ~/.commandcode/auth.json, ~/.pi/agent/auth.json or ~/.omp/agent/auth.json", + timestamp: now(), + } + stream.push({ type: "error", reason: "error", error: msg }) + stream.end() + return + } + + const output: AssistantMessageLike = { + role: "assistant", + content: [], + api: model.api, + provider: model.provider, + model: model.id, + usage: defaultUsage(), + stopReason: "stop", + timestamp: now(), + } + + const controller = new AbortController() + let reader: ReadableStreamDefaultReader | undefined + let textBlock: TextContent | undefined + let currentTextIdx = -1 + let thinkingIdx = -1 + let finished = false + + const abortUpstream = () => { + if (!controller.signal.aborted) controller.abort() + try { + reader?.cancel().catch(() => undefined) + } catch { + // Reader cancellation is best-effort. + } + } + + if (options?.signal?.aborted) { + abortUpstream() + } else { + options?.signal?.addEventListener("abort", abortUpstream, { + once: true, + }) + } + + const endTextBlock = () => { + if (!textBlock) return + stream.push({ + type: "text_end", + contentIndex: currentTextIdx, + content: textBlock.text, + partial: output, + }) + textBlock = undefined + currentTextIdx = -1 + } + + const endThinking = () => { + if (thinkingIdx < 0) return + const tc = output.content[thinkingIdx] + if (tc && tc.type === "thinking") { + stream.push({ + type: "thinking_end", + contentIndex: thinkingIdx, + content: (tc as { thinking: string }).thinking, + partial: output, + }) + } + thinkingIdx = -1 + } + + const handleEvent = (event: unknown) => { + if (!isRecord(event)) return + + switch (event.type) { + case "text-delta": { + endThinking() + if (!textBlock) { + textBlock = { type: "text", text: "" } + output.content.push(textBlock) + currentTextIdx = output.content.length - 1 + stream.push({ + type: "text_start", + contentIndex: currentTextIdx, + partial: output, + }) + } + const delta = stringValue(event.text) ?? "" + textBlock.text += delta + stream.push({ + type: "text_delta", + contentIndex: currentTextIdx, + delta, + partial: output, + }) + break + } + + case "reasoning-start": { + endTextBlock() + break + } + + case "reasoning-delta": { + endTextBlock() + const delta = stringValue(event.text) ?? "" + if (thinkingIdx < 0) { + output.content.push({ type: "thinking", thinking: delta }) + thinkingIdx = output.content.length - 1 + stream.push({ + type: "thinking_start", + contentIndex: thinkingIdx, + partial: output, + }) + } else { + const tc = output.content[thinkingIdx] + if (tc && tc.type === "thinking") { + ;(tc as { thinking: string }).thinking += delta + } + } + stream.push({ + type: "thinking_delta", + contentIndex: thinkingIdx, + delta, + partial: output, + }) + break + } + + case "reasoning-end": { + endThinking() + break + } + + case "tool-result": { + break + } + + case "tool-call": { + endTextBlock() + endThinking() + const toolCall: ToolCallContent = { + type: "toolCall", + id: stringValue(event.toolCallId) ?? "", + name: stringValue(event.toolName) ?? "", + arguments: recordOrEmpty(event.input ?? event.args ?? event.arguments), + } + output.content.push(toolCall) + const idx = output.content.length - 1 + stream.push({ + type: "toolcall_start", + contentIndex: idx, + partial: output, + }) + stream.push({ + type: "toolcall_end", + contentIndex: idx, + toolCall, + partial: output, + }) + break + } + + case "finish": { + const usage = commandCodeUsage(event) + if (usage) { + const details = commandCodeInputTokenDetails(usage) + const totalInput = numberValue(usage.inputTokens) ?? 0 + const input = numberValue(details?.noCacheTokens) + const cacheRead = numberValue(details?.cacheReadTokens) ?? 0 + const cacheWrite = numberValue(details?.cacheWriteTokens) ?? 0 + output.usage.input = input ?? Math.max(0, totalInput - cacheRead - cacheWrite) + output.usage.output = numberValue(usage.outputTokens) ?? 0 + output.usage.cacheRead = cacheRead + output.usage.cacheWrite = cacheWrite + output.usage.totalTokens = + output.usage.input + + output.usage.output + + output.usage.cacheRead + + output.usage.cacheWrite + deps.calculateCost(model, output.usage) + } + output.stopReason = mapFinishReason(event.finishReason) + finished = true + break + } + + case "error": { + const message = + commandCodeErrorMessage(event.error) ?? + commandCodeErrorMessage(event.message) ?? + "Stream error" + output.stopReason = "error" + output.errorMessage = message + throw new Error(message) + } + } + } + + try { + stream.push({ type: "start", partial: output }) + if (controller.signal.aborted) throw abortError("Aborted") + + const workingDir = cwd() + const threadId = uuid() + const reasoningEffort = mappedReasoningEffort(model, options) + const timeoutMs = options?.timeoutMs + + const allowImages = modelSupportsImageInput(model.id) + if (!allowImages) assertTextOnlyMessages(context.messages) + + let body: unknown = { + config: { + workingDir, + date: new Date(now()).toISOString().split("T")[0], + environment: getEnvironmentInfo(), + structure: [], + isGitRepo: false, + currentBranch: "", + mainBranch: "", + gitStatus: "", + recentCommits: [], + }, + memory: null, + taste: null, + skills: null, + params: { + model: model.id, + messages: messagesToCC(context.messages, { allowImages }), + tools: toolsToJson(context.tools), + system: systemPromptToText(context.systemPrompt), + max_tokens: generateMaxTokens(model, options), + temperature: 0.3, + stream: true, + ...(reasoningEffort ? { reasoning_effort: reasoningEffort } : {}), + }, + threadId, + } + + const payloadController = new AbortController() + const onPayloadAbort = () => payloadController.abort() + controller.signal.addEventListener("abort", onPayloadAbort, { once: true }) + let nextBody: unknown + try { + nextBody = await raceAbortWithTimeout( + Promise.resolve(options?.onPayload?.(body, model)), + payloadController, + timeoutMs, + ) + } finally { + controller.signal.removeEventListener("abort", onPayloadAbort) + } + if (nextBody !== undefined) body = nextBody + + const maxRetries = options?.maxRetries ?? DEFAULT_MAX_RETRIES + const maxRetryDelayMs = effectiveMaxRetryDelayMs(options?.maxRetryDelayMs) + const requestHeaders = { + "Content-Type": "application/json", + Authorization: `Bearer ${apiKey}`, + "x-command-code-version": COMMAND_CODE_CLI_VERSION, + "x-cli-environment": "production", + "x-project-slug": projectSlugFromPath(workingDir), + "x-taste-learning": "true", + "x-co-flag": "false", + ...options?.headers, + } + const bodyStr = JSON.stringify(body) + + let response!: Response + retryLoop: for (let attempt = 0; ; attempt++) { + const attemptController = new AbortController() + let attemptTimedOut = false + let attemptTimeoutId: ReturnType | undefined + + const clearAttemptTimeout = () => { + if (attemptTimeoutId !== undefined) { + clearTimeout(attemptTimeoutId) + attemptTimeoutId = undefined + } + } + + if (timeoutMs !== undefined) { + attemptTimeoutId = setTimeout(() => { + attemptTimedOut = true + attemptController.abort() + }, timeoutMs) + } + const onOuterAbort = () => attemptController.abort() + controller.signal.addEventListener("abort", onOuterAbort, { once: true }) + const raceAttempt = (promise: Promise): Promise => + raceAbort(promise, attemptController.signal).catch((error: unknown) => { + if (attemptTimedOut) throw timeoutError(timeoutMs) + throw error + }) + + try { + try { + response = await fetchImpl(`${apiBase}/alpha/generate`, { + method: "POST", + headers: requestHeaders, + body: bodyStr, + signal: attemptController.signal, + }) + } catch (fetchError: unknown) { + if (controller.signal.aborted) throw abortError("Aborted") + if (attemptTimedOut) { + if (attempt < maxRetries) continue retryLoop + throw timeoutError(timeoutMs) + } + throw fetchError + } + + // --- HTTP-level retry --- + if (!response.ok && isRetryableStatus(response.status)) { + const retryAfter = response.headers.get("retry-after") + const waitMs = retryDelayMs(attempt, retryAfter, maxRetryDelayMs) + if (waitMs < 0) { + const requestedSeconds = parseRetryAfterSeconds(retryAfter) ?? 0 + const capLabel = + maxRetryDelayMs === Number.POSITIVE_INFINITY ? "disabled" : `${maxRetryDelayMs}ms` + throw new Error(`Retry-After delay ${requestedSeconds}s exceeds max ${capLabel}`) + } + if (attempt < maxRetries) { + await response.text().catch(() => "") + if (waitMs > 0) await delay(waitMs, controller.signal) + continue retryLoop + } + } + + try { + await raceAttempt( + Promise.resolve( + options?.onResponse?.( + { + status: response.status, + headers: headersToRecord(response.headers), + }, + model, + ), + ), + ) + } catch (error: unknown) { + if (attemptTimedOut && attempt < maxRetries) continue retryLoop + throw error + } + + if (!response.ok) { + const errBody = await raceAttempt(response.text().catch(() => "")) + let errorDetail: string | undefined + try { + const parsedBody: unknown = JSON.parse(errBody) + errorDetail = commandCodeErrorMessage(parsedBody) + } catch { + // Preserve useful plain-text provider errors only after secret + // redaction; upstream/proxy bodies may echo credentials. + } + const safeBody = redactCommandCodeErrorText(errBody).slice(0, 500) + const detail = redactCommandCodeErrorText( + errorDetail ?? (safeBody || "Provider returned an error"), + ) + throw new Error(`Command Code API error ${response.status}: ${detail}`) + } + + // --- Read response stream --- + reader = response.body?.getReader() + if (!reader) throw new Error("No response body") + + const decoder = new TextDecoder() + let buffer = "" + + try { + readLoop: for (;;) { + if (controller.signal.aborted) throw abortError("Aborted") + const { done, value } = await raceAbort(reader.read(), attemptController.signal) + if (done) { + if (buffer.trim()) handleEvent(parseStreamEventLine(buffer)) + break + } + if (controller.signal.aborted) throw abortError("Aborted") + + buffer += decoder.decode(value, { stream: true }) + const lines = buffer.split("\n") + buffer = lines.pop() ?? "" + + for (const line of lines) { + if (controller.signal.aborted) throw abortError("Aborted") + handleEvent(parseStreamEventLine(line)) + if (finished) break readLoop + } + } + } catch (streamError: unknown) { + // Stream-level error (e.g. API returned 200 OK but sent an error event) + // or per-attempt timeout during stream reading. + await reader.cancel().catch(() => {}) + try { + reader.releaseLock() + } catch {} + reader = undefined + + if (controller.signal.aborted) throw streamError + + // Never retry after visible content was emitted (including timeout mid-stream). + const canRetry = output.content.length === 0 && attempt < maxRetries + if (canRetry) { + output.content.length = 0 + textBlock = undefined + currentTextIdx = -1 + thinkingIdx = -1 + output.stopReason = "stop" + output.errorMessage = undefined + finished = false + const waitMs = attemptTimedOut ? 0 : retryDelayMs(attempt, null, maxRetryDelayMs) + if (waitMs > 0) await delay(waitMs, controller.signal) + continue retryLoop + } + if (attemptTimedOut) throw timeoutError(timeoutMs) + throw streamError + } + + // Stream completed successfully. + endTextBlock() + endThinking() + + stream.push({ + type: "done", + reason: successStopReason(output.stopReason), + message: output, + }) + stream.end() + break retryLoop + } finally { + controller.signal.removeEventListener("abort", onOuterAbort) + clearAttemptTimeout() + } + } + } catch (error: unknown) { + const reason: ErrorReason = controller.signal.aborted ? "aborted" : "error" + output.stopReason = reason + output.errorMessage = + reason === "aborted" + ? "Request aborted" + : redactCommandCodeErrorText(error instanceof Error ? error.message : String(error)) + stream.push({ type: "error", reason, error: output }) + stream.end() + } finally { + options?.signal?.removeEventListener("abort", abortUpstream) + try { + await reader?.cancel() + } catch { + // Reader may already be closed/cancelled. + } + try { + reader?.releaseLock() + } catch { + // Reader may already be released/cancelled by the abort path. + } + } + } + + run().catch((error: unknown) => { + const msg: AssistantMessageLike = { + role: "assistant", + content: [], + api: model.api, + provider: model.provider, + model: model.id, + usage: defaultUsage(), + stopReason: "error", + errorMessage: redactCommandCodeErrorText( + error instanceof Error ? error.message : String(error), + ), + timestamp: now(), + } + stream.push({ type: "error", reason: "error", error: msg }) + stream.end() + }) + + return stream + } +} diff --git a/src/cost.ts b/src/cost.ts new file mode 100644 index 0000000..55c1bd0 --- /dev/null +++ b/src/cost.ts @@ -0,0 +1,31 @@ +/** + * Local cost calculation for Command Code usage. + * + * Mirrors pi-ai's `calculateCost` arithmetic exactly. The provider ships its + * own copy because Oh My Pi's legacy pi-ai shim does not export + * `calculateCost`, which broke extension installation there (issue #24). + * `tests/test-cost.ts` locks this implementation to the pi-ai original. + */ + +import type { ModelLike, Usage } from "./types.ts" + +export function calculateCommandCodeCost(model: ModelLike, usage: Usage): void { + const inputTokens = usage.input + usage.cacheRead + usage.cacheWrite + let rates = model.cost + let matchedThreshold = -1 + for (const tier of model.cost.tiers ?? []) { + if (inputTokens > tier.inputTokensAbove && tier.inputTokensAbove > matchedThreshold) { + rates = tier + matchedThreshold = tier.inputTokensAbove + } + } + + const longWrite = usage.cacheWrite1h ?? 0 + const shortWrite = usage.cacheWrite - longWrite + usage.cost.input = (rates.input / 1_000_000) * usage.input + usage.cost.output = (rates.output / 1_000_000) * usage.output + usage.cost.cacheRead = (rates.cacheRead / 1_000_000) * usage.cacheRead + usage.cost.cacheWrite = (rates.cacheWrite * shortWrite + rates.input * 2 * longWrite) / 1_000_000 + usage.cost.total = + usage.cost.input + usage.cost.output + usage.cost.cacheRead + usage.cost.cacheWrite +} diff --git a/src/json-schema.ts b/src/json-schema.ts new file mode 100644 index 0000000..a2ffc02 --- /dev/null +++ b/src/json-schema.ts @@ -0,0 +1,382 @@ +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value) +} + +function stringValue(value: unknown): string | undefined { + return typeof value === "string" ? value : undefined +} + +function booleanValue(value: unknown): boolean | undefined { + return typeof value === "boolean" ? value : undefined +} + +type JsonSchemaValue = boolean | Record + +const JSON_SCHEMA_TYPES = new Set([ + "array", + "boolean", + "integer", + "null", + "number", + "object", + "string", +]) + +const LEGACY_KINDS = new Set([ + "any", + "array", + "boolean", + "enum", + "integer", + "intersect", + "intersection", + "literal", + "never", + "null", + "nullable", + "number", + "object", + "optional", + "string", + "undefined", + "union", + "unknown", +]) + +const LEGACY_FIELDS = new Set([ + "element", + "kind", + "inner", + "optional", + "value", + "values", + "variants", + "wrapped", +]) + +const SCHEMA_MAP_FIELDS = new Set([ + "$defs", + "definitions", + "dependentSchemas", + "patternProperties", + "properties", +]) + +const SCHEMA_ARRAY_FIELDS = new Set(["allOf", "anyOf", "oneOf", "prefixItems"]) + +const SCHEMA_VALUE_FIELDS = new Set([ + "additionalItems", + "additionalProperties", + "contains", + "contentSchema", + "else", + "if", + "items", + "not", + "propertyNames", + "then", + "unevaluatedItems", + "unevaluatedProperties", +]) + +const SCHEMA_KEYWORDS = new Set([ + "$anchor", + "$comment", + "$defs", + "$dynamicAnchor", + "$dynamicRef", + "$id", + "$ref", + "$schema", + "$vocabulary", + "additionalItems", + "additionalProperties", + "allOf", + "anyOf", + "const", + "contains", + "contentEncoding", + "contentMediaType", + "contentSchema", + "default", + "definitions", + "dependentRequired", + "dependentSchemas", + "description", + "else", + "enum", + "examples", + "exclusiveMaximum", + "exclusiveMinimum", + "format", + "if", + "items", + "maxContains", + "maxItems", + "maxLength", + "maxProperties", + "maximum", + "minContains", + "minItems", + "minLength", + "minProperties", + "minimum", + "multipleOf", + "not", + "oneOf", + "pattern", + "patternProperties", + "prefixItems", + "properties", + "propertyNames", + "readOnly", + "required", + "title", + "type", + "unevaluatedItems", + "unevaluatedProperties", + "uniqueItems", + "writeOnly", +]) + +function stringArray(value: unknown): string[] | undefined { + if (!Array.isArray(value)) return undefined + const values = value.filter((item): item is string => typeof item === "string") + return values.length === value.length ? values : undefined +} + +function validSchemaType(value: unknown): boolean { + if (typeof value === "string") return JSON_SCHEMA_TYPES.has(value) + if (!Array.isArray(value) || value.length === 0) return false + return value.every((item) => typeof item === "string" && JSON_SCHEMA_TYPES.has(item)) +} + +function legacyKind(schema: Record): string | undefined { + const explicitKind = stringValue(schema.kind)?.toLowerCase() + if (explicitKind && LEGACY_KINDS.has(explicitKind)) return explicitKind + + const type = stringValue(schema.type) + const normalized = type?.toLowerCase() + if (!normalized || !LEGACY_KINDS.has(normalized)) return undefined + if (!validSchemaType(type) || Object.keys(schema).some((key) => LEGACY_FIELDS.has(key))) { + return normalized + } + return undefined +} + +function looksLikeJsonSchema(schema: Record): boolean { + if (Object.keys(schema).length === 0) return true + if (schema.type !== undefined && !validSchemaType(schema.type)) return false + return Object.keys(schema).some((key) => SCHEMA_KEYWORDS.has(key)) +} + +function isOptionalSchema(schema: unknown): boolean { + if (!isRecord(schema)) return false + if (booleanValue(schema.optional) === true) return true + + const kind = legacyKind(schema) + if (kind === "optional") return true + if (kind !== "union") return false + + const variants = Array.isArray(schema.variants) + ? schema.variants + : Array.isArray(schema.anyOf) + ? schema.anyOf + : [] + return variants.some((variant) => legacyKind(isRecord(variant) ? variant : {}) === "undefined") +} + +function schemaValue(value: unknown, seen: WeakSet): JsonSchemaValue { + if (typeof value === "boolean") return value + if (!isRecord(value)) return {} + return convertSchema(value, seen) +} + +function setSchemaProperty(target: Record, key: string, value: unknown): void { + Object.defineProperty(target, key, { + configurable: true, + enumerable: true, + value, + writable: true, + }) +} + +function schemaMap(value: unknown, seen: WeakSet): Record { + if (!isRecord(value)) return {} + const out: Record = {} + for (const [key, item] of Object.entries(value)) { + setSchemaProperty(out, key, schemaValue(item, seen)) + } + return out +} + +function schemaArray(value: unknown, seen: WeakSet): unknown[] { + if (!Array.isArray(value)) return [] + return value.map((item) => schemaValue(item, seen)) +} + +function isSchemaValue(value: unknown): value is JsonSchemaValue { + return typeof value === "boolean" || isRecord(value) +} + +function copySchemaObject( + source: Record, + seen: WeakSet, + legacy: boolean, + forcedType?: string, +): JsonSchemaValue { + const out: Record = {} + + for (const [key, value] of Object.entries(source)) { + if (legacy && LEGACY_FIELDS.has(key)) continue + if (key === "nullable" || (forcedType !== undefined && key === "type")) continue + + if (key === "required") { + const required = stringArray(value) + if (required) out.required = required + } else if (SCHEMA_MAP_FIELDS.has(key)) { + out[key] = schemaMap(value, seen) + } else if (SCHEMA_ARRAY_FIELDS.has(key)) { + out[key] = schemaArray(value, seen) + } else if (SCHEMA_VALUE_FIELDS.has(key)) { + out[key] = + Array.isArray(value) && key === "items" + ? schemaArray(value, seen) + : schemaValue(value, seen) + } else { + out[key] = value + } + } + + if (forcedType !== undefined) out.type = forcedType + if (booleanValue(source.nullable) === true) return makeNullable(out) + return out +} + +function makeNullable(schema: Record): Record { + const type = schema.type + if (typeof type === "string") { + if (type === "null") return schema + return { ...schema, type: [type, "null"] } + } + if (Array.isArray(type) && !type.includes("null")) { + return { ...schema, type: [...type, "null"] } + } + if (Array.isArray(schema.anyOf)) { + return { ...schema, anyOf: [...schema.anyOf, { type: "null" }] } + } + return { anyOf: [schema, { type: "null" }] } +} + +function legacyVariants(schema: Record): unknown[] { + if (Array.isArray(schema.variants)) return schema.variants + if (Array.isArray(schema.anyOf)) return schema.anyOf + return [] +} + +function convertLegacySchema( + source: Record, + kind: string, + seen: WeakSet, +): JsonSchemaValue { + if (kind === "optional") return schemaValue(source.wrapped ?? source.inner, seen) + if (kind === "nullable") { + const wrapped = schemaValue(source.wrapped ?? source.inner, seen) + return typeof wrapped === "boolean" ? wrapped : makeNullable(wrapped) + } + if (kind === "undefined" || kind === "never" || kind === "any" || kind === "unknown") return {} + + if (kind === "union" || kind === "intersect" || kind === "intersection") { + const variants = legacyVariants(source) + .map((variant) => schemaValue(variant, seen)) + .filter( + (variant) => + isSchemaValue(variant) && + (typeof variant === "boolean" || Object.keys(variant).length > 0), + ) + if (variants.length === 0) return copySchemaObject(source, seen, true) + if (variants.length === 1) return variants[0] ?? {} + + const out = copySchemaObject(source, seen, true) + if (typeof out !== "boolean") out[kind === "union" ? "anyOf" : "allOf"] = variants + return out + } + + if (kind === "object") { + const converted = copySchemaObject(source, seen, true, "object") + if (typeof converted === "boolean") return converted + const out = converted + const sourceProperties = isRecord(source.properties) ? source.properties : undefined + if (!sourceProperties) return out + + const properties: Record = {} + const optional = stringArray(source.optional) ?? [] + for (const [key, value] of Object.entries(sourceProperties)) { + setSchemaProperty(properties, key, schemaValue(value, seen)) + } + out.properties = properties + + const explicitRequired = stringArray(source.required) + const required = + explicitRequired ?? + Object.entries(sourceProperties) + .filter(([key, value]) => !optional.includes(key) && !isOptionalSchema(value)) + .map(([key]) => key) + if (required.length > 0) out.required = required + else delete out.required + return out + } + + if (kind === "array") { + const converted = copySchemaObject(source, seen, true, "array") + if (typeof converted === "boolean") return converted + const out = converted + if (!("items" in source) && "element" in source) out.items = schemaValue(source.element, seen) + return out + } + + if (kind === "enum") { + const converted = copySchemaObject(source, seen, true) + if (typeof converted === "boolean") return converted + const out = converted + if (!("enum" in out) && Array.isArray(source.values)) out.enum = source.values + return out + } + + if (kind === "literal") { + const converted = copySchemaObject(source, seen, true) + if (typeof converted === "boolean") return converted + const out = converted + if (!("const" in out) && "value" in source) out.const = source.value + return out + } + + const scalarType = + kind === "string" || + kind === "number" || + kind === "boolean" || + kind === "integer" || + kind === "null" + ? kind + : undefined + return scalarType ? copySchemaObject(source, seen, true, scalarType) : {} +} + +function convertSchema(source: Record, seen: WeakSet): JsonSchemaValue { + if (seen.has(source)) return {} + seen.add(source) + try { + const kind = legacyKind(source) + if (kind) return convertLegacySchema(source, kind, seen) + if (!looksLikeJsonSchema(source)) return {} + return copySchemaObject(source, seen, false) + } finally { + seen.delete(source) + } +} + +export function toJsonSchema(schema: unknown): unknown { + if (typeof schema === "boolean") return schema + if (!isRecord(schema)) return {} + return convertSchema(schema, new WeakSet()) +} diff --git a/src/overflow.ts b/src/overflow.ts new file mode 100644 index 0000000..3e3e106 --- /dev/null +++ b/src/overflow.ts @@ -0,0 +1,120 @@ +const COMMAND_CODE_PROVIDER = "commandcode" +const CONTEXT_OVERFLOW_PREFIX = "context_length_exceeded:" + +const COMMAND_CODE_OVERFLOW_PATTERNS = [ + /\b(?:context[_\s-]*(?:length|window)|model[_\s-]*context[_\s-]*window)[_\s-]*(?:exceeded|overflow(?:ed)?|too[_\s-]*(?:large|long))\b/i, + /\b(?:context|prompt|input)[_\s-]*(?:length|window|size|tokens?|limit|maximum)\b[\s\S]{0,120}\b(?:exceed(?:ed|s)?|overflow(?:ed|s)?|too\s+(?:large|long)|(?:maximum|limit)\s+(?:reached|exceeded|hit))\b/i, + /\b(?:exceed(?:ed|s)?|overflow(?:ed|s)?|too\s+(?:large|long))\b[\s\S]{0,120}\b(?:context|prompt|input)[_\s-]*(?:length|window|size|tokens?|limit|maximum)\b/i, + /\b(?:prompt|input|context)\b[\s\S]{0,32}\btoo\s+(?:large|long)\b/i, + /\b(?:prompt|input)[_\s-]*too[_\s-]*(?:large|long)\b/i, + /\b(?:prompt|input)[_\s-]*tokens?[_\s-]*(?:limit|maximum|max)[_\s-]*(?:exceeded|reached)\b/i, + /\b(?:prompt|input)[_\s-]*(?:tokens?|length|size)\b[\s\S]{0,120}\b(?:limit|maximum)\b[\s\S]{0,40}\b(?:exceed(?:ed|s)?|reached|hit)\b/i, + /\b(?:maximum|limit)[_\s-]+(?:allowed[_\s-]+)?(?:context|prompt|input)[_\s-]*(?:length|window|size|tokens?)\b/i, +] + +const NON_OVERFLOW_PATTERNS = [ + /\brate[_\s-]*limit\b/i, + /\btoo\s+many\s+requests\b/i, + /\b(?:capacity|quota|throttl(?:e|ed|ing)?|concurren(?:cy|t)|overloaded)\b/i, + /\b(?:service|temporarily)\s+unavailable\b/i, + /\bstatus(?:[_\s-]*code)?\s*[:=]\s*429\b/i, +] + +const CONTEXT_OVERFLOW_PREFIX_PATTERN = /context_length_exceeded/i + +const HTTP_RATE_LIMIT_STATUS_PATTERNS = [ + /\b(?:api\s+error|http|status(?:[_\s-]*code)?|status[_\s-]*code)\s*[:(]?\s*429\b/i, + /["']?(?:status|status[_\s-]*code)["']?\s*:\s*429\b/i, +] + +const BEARER_PATTERN = /\bBearer\s+[A-Za-z0-9._~+/=-]+/gi +const CREDENTIAL_PATTERN = + /\b(?:api[-_ ]?key|apikey|access[-_ ]?token|refresh[-_ ]?token|token|secret|password|authorization)\s*[=:]\s*[^\s,;)]+/gi +const USER_TOKEN_PATTERN = /\b(?:user|cc)_[A-Za-z0-9_-]{8,}\b/gi +const QUERY_SECRET_PATTERN = + /([?&](?:api[-_ ]?key|apikey|access_token|refresh_token|token|secret|password)=)[^&#\s]+/gi +const STANDALONE_SECRET_PATTERN = + /\b(?:sk|rk|ghp|github_pat|xox[baprs])[-_A-Za-z0-9]{16,}\b|\beyJ[A-Za-z0-9_-]{20,}\.[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]{10,}\b/g + +export function redactCommandCodeErrorText(value: string): string { + return value + .replace(BEARER_PATTERN, "Bearer [redacted]") + .replace(CREDENTIAL_PATTERN, (match) => { + const separatorIndex = match.search(/[=:]/) + return separatorIndex < 0 ? "[redacted]" : `${match.slice(0, separatorIndex + 1)}[redacted]` + }) + .replace(USER_TOKEN_PATTERN, "[redacted]") + .replace(QUERY_SECRET_PATTERN, "$1[redacted]") + .replace(STANDALONE_SECRET_PATTERN, "[redacted]") +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null +} + +export interface CommandCodeMessageLike { + role: string + provider: string + stopReason: string + errorMessage?: string +} + +export function commandCodeErrorMessage(value: unknown): string | undefined { + if (typeof value === "string") return value + if (!isRecord(value)) return undefined + + const record = value + const parts: string[] = [] + for (const key of [ + "message", + "errorMessage", + "error", + "detail", + "details", + "code", + "type", + "reason", + ]) { + const part = commandCodeErrorMessage(record[key]) + if (part && !parts.includes(part)) parts.push(part) + } + + for (const key of ["status", "statusCode", "httpStatus"]) { + const status = record[key] + if (typeof status === "string" || typeof status === "number") { + const statusPart = `status: ${status}` + if (!parts.includes(statusPart)) parts.push(statusPart) + } + } + + return parts.length > 0 ? redactCommandCodeErrorText(parts.join(": ")) : undefined +} + +export function normalizeCommandCodeErrorMessage( + errorMessage: string | undefined, +): string | undefined { + if (!errorMessage) return undefined + if (CONTEXT_OVERFLOW_PREFIX_PATTERN.test(errorMessage)) return undefined + if (NON_OVERFLOW_PATTERNS.some((pattern) => pattern.test(errorMessage))) return undefined + if (HTTP_RATE_LIMIT_STATUS_PATTERNS.some((pattern) => pattern.test(errorMessage))) + return undefined + if (!COMMAND_CODE_OVERFLOW_PATTERNS.some((pattern) => pattern.test(errorMessage))) + return undefined + + return `${CONTEXT_OVERFLOW_PREFIX} ${errorMessage}` +} + +export function normalizeCommandCodeMessage( + message: T, + modelProvider?: string, +): { message: T & { errorMessage: string } } | undefined { + if (message.role !== "assistant" || message.stopReason !== "error") return undefined + if (message.provider !== COMMAND_CODE_PROVIDER && modelProvider !== COMMAND_CODE_PROVIDER) { + return undefined + } + + const errorMessage = normalizeCommandCodeErrorMessage(message.errorMessage) + if (!errorMessage) return undefined + + return { message: { ...message, errorMessage } } +} diff --git a/src/transport.ts b/src/transport.ts new file mode 100644 index 0000000..9d83d5d --- /dev/null +++ b/src/transport.ts @@ -0,0 +1,139 @@ +import type { + AssistantMessageEvent, + AssistantMessageEventStreamLike, + ContextLike, + ModelLike, + StreamOptions, +} from "./types.ts" + +export type CommandCodeTransport = "unknown" | "provider" | "generate" + +interface TransportDependencies { + createStream: () => AssistantMessageEventStreamLike + streamProvider: ( + model: ModelLike, + context: ContextLike, + options?: StreamOptions, + ) => AssistantMessageEventStreamLike + streamGenerate: ( + model: ModelLike, + context: ContextLike, + options?: StreamOptions, + ) => AssistantMessageEventStreamLike +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value) +} + +async function isUpgradeRequired(response: Response): Promise { + if (response.status !== 403) return false + + try { + const body: unknown = await response.clone().json() + if (!isRecord(body)) return false + const error = isRecord(body.error) ? body.error : body + return error.code === "upgrade_required" + } catch { + return false + } +} + +export function createCommandCodeTransportRouter(deps: TransportDependencies) { + let transport: CommandCodeTransport = "unknown" + let apiKey: string | undefined + + function pipe( + source: AssistantMessageEventStreamLike, + target: AssistantMessageEventStreamLike, + ): Promise { + return (async () => { + for await (const event of source) target.push(event) + })() + } + + return { + getTransport(): CommandCodeTransport { + return transport + }, + + reset(): void { + transport = "unknown" + apiKey = undefined + }, + + stream( + model: ModelLike, + context: ContextLike, + options?: StreamOptions, + ): AssistantMessageEventStreamLike { + if (options?.apiKey !== apiKey) { + apiKey = options?.apiKey + transport = "unknown" + } + if (transport === "generate") return deps.streamGenerate(model, context, options) + + const output = deps.createStream() + let upgradeRequired = false + const fetchImpl = options?.fetch ?? fetch + const providerOptions: StreamOptions = { + ...options, + fetch: async (input, init) => { + const response = await fetchImpl(input, init) + if (await isUpgradeRequired(response)) upgradeRequired = true + return response + }, + onResponse: async (response, responseModel) => { + if (upgradeRequired) return + await options?.onResponse?.(response, responseModel) + }, + } + + const run = async () => { + const providerStream = deps.streamProvider(model, context, providerOptions) + + for await (const event of providerStream) { + if (!upgradeRequired) { + transport = "provider" + output.push(event) + } + } + + if (upgradeRequired) { + transport = "generate" + await pipe(deps.streamGenerate(model, context, options), output) + } + output.end() + } + + run().catch((error: unknown) => { + const message = error instanceof Error ? error.message : String(error) + output.push({ + type: "error", + reason: "error", + error: { + role: "assistant", + content: [], + api: model.api, + provider: model.provider, + model: model.id, + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "error", + errorMessage: message, + timestamp: Date.now(), + }, + }) + output.end() + }) + + return output + }, + } +} diff --git a/src/types.ts b/src/types.ts new file mode 100644 index 0000000..dc680ea --- /dev/null +++ b/src/types.ts @@ -0,0 +1,202 @@ +export type StopReason = "stop" | "length" | "toolUse" +export type ErrorReason = "error" | "aborted" +export type TerminalReason = StopReason | ErrorReason + +export interface UsageCost { + input: number + output: number + cacheRead: number + cacheWrite: number + total: number +} + +export interface Usage { + input: number + output: number + cacheRead: number + cacheWrite: number + cacheWrite1h?: number + totalTokens: number + cost: UsageCost +} + +export interface TextContent { + type: "text" + text: string +} + +export interface ThinkingContent { + type: "thinking" + thinking: string +} + +export interface ToolCallContent { + type: "toolCall" + id: string + name: string + arguments: Record +} + +export type AssistantContent = TextContent | ThinkingContent | ToolCallContent + +export interface AssistantMessageLike { + role: "assistant" + content: AssistantContent[] + api: unknown + provider: string + model: string + usage: Usage + stopReason: TerminalReason + errorMessage?: string + timestamp: number +} + +export interface ModelCostRates { + input: number + output: number + cacheRead: number + cacheWrite: number +} + +export interface ModelCostTier extends ModelCostRates { + inputTokensAbove: number +} + +export interface ModelCost extends ModelCostRates { + tiers?: readonly ModelCostTier[] +} + +export interface ModelLike { + id: string + api: unknown + provider: string + maxTokens: number + cost: ModelCost + reasoning?: boolean + thinkingLevelMap?: Partial> + thinking?: { + mode?: "effort" + effortMap?: Partial> + efforts?: readonly string[] + } +} + +export interface MessageLike { + role: string + content?: unknown + toolCallId?: string + toolName?: string + isError?: boolean +} + +export interface ToolLike { + name: string + description?: string + parameters?: unknown +} + +export interface ContextLike { + systemPrompt?: string + messages?: readonly MessageLike[] + tools?: readonly ToolLike[] +} + +export interface ProviderResponseInfo { + status: number + headers: Record +} + +export interface StreamOptions { + apiKey?: string + signal?: AbortSignal + headers?: Record + fetch?: typeof fetch + maxTokens?: number + /** Resolved pi thinking level; forwarded only through the model's map. */ + reasoning?: string + onPayload?: (payload: unknown, model: ModelLike) => unknown | Promise + onResponse?: (response: ProviderResponseInfo, model: ModelLike) => void | Promise + /** + * HTTP request timeout in milliseconds. + * Applied per-attempt; on timeout the request is retried if retries remain. + */ + timeoutMs?: number + /** + * Maximum retry attempts for transient HTTP errors (429, 5xx). + * Default: 0 (pi agent-level retry handles visible retries when unset). + */ + maxRetries?: number + /** + * Maximum delay in milliseconds to wait for a retry when the server requests + * a long wait via Retry-After. If the server's requested delay exceeds this + * value, the request fails immediately. Default: 60000 (60 seconds). + * Set to 0 to disable the cap. + */ + maxRetryDelayMs?: number +} + +export type AssistantMessageEvent = + | { type: "start"; partial: AssistantMessageLike } + | { type: "text_start"; contentIndex: number; partial: AssistantMessageLike } + | { + type: "text_delta" + contentIndex: number + delta: string + partial: AssistantMessageLike + } + | { + type: "text_end" + contentIndex: number + content: string + partial: AssistantMessageLike + } + | { + type: "thinking_start" + contentIndex: number + partial: AssistantMessageLike + } + | { + type: "thinking_delta" + contentIndex: number + delta: string + partial: AssistantMessageLike + } + | { + type: "thinking_end" + contentIndex: number + content: string + partial: AssistantMessageLike + } + | { + type: "toolcall_start" + contentIndex: number + partial: AssistantMessageLike + } + | { + type: "toolcall_end" + contentIndex: number + toolCall: ToolCallContent + partial: AssistantMessageLike + } + | { type: "done"; reason: StopReason; message: AssistantMessageLike } + | { type: "error"; reason: ErrorReason; error: AssistantMessageLike } + +export interface AssistantMessageEventStreamLike extends AsyncIterable { + push(event: AssistantMessageEvent): void + end(): void +} + +export interface CoreDependencies { + createStream: () => AssistantMessageEventStreamLike + calculateCost: (model: ModelLike, usage: Usage) => void + apiBase?: string + fetchImpl?: typeof fetch + authPaths?: readonly string[] + env?: NodeJS.ProcessEnv + cwd?: () => string + now?: () => number + uuid?: () => string + homeDir?: () => string + /** Injectable delay for retry backoff. Defaults to setTimeout. */ + delay?: (ms: number, signal: AbortSignal) => Promise +} diff --git a/tests/helpers.ts b/tests/helpers.ts new file mode 100644 index 0000000..f4b3f65 --- /dev/null +++ b/tests/helpers.ts @@ -0,0 +1,289 @@ +import { createServer, type IncomingHttpHeaders, type Server } from "node:http" + +import { + createStreamCommandCode, + type AssistantMessageEvent, + type AssistantMessageEventStreamLike, + type ContextLike, + type CoreDependencies, + type ModelLike, + type Usage, +} from "../src/core.ts" + +export function createTestEventStream(): AssistantMessageEventStreamLike { + const events: AssistantMessageEvent[] = [] + const waiters: Array<() => void> = [] + let ended = false + + const wake = () => { + const waiter = waiters.shift() + if (waiter) waiter() + } + + return { + push(event: AssistantMessageEvent) { + events.push(event) + wake() + }, + end() { + ended = true + while (waiters.length > 0) wake() + }, + [Symbol.asyncIterator]() { + let index = 0 + return { + async next(): Promise> { + while (index >= events.length && !ended) { + await new Promise((resolve) => waiters.push(resolve)) + } + if (index < events.length) { + const value = events[index] + index += 1 + return { done: false, value } + } + return { done: true, value: undefined } + }, + } + }, + } +} + +export async function collectEvents( + stream: AssistantMessageEventStreamLike, + timeoutMs = 2_000, +): Promise { + const events: AssistantMessageEvent[] = [] + + const collect = async () => { + for await (const event of stream) { + events.push(event) + if (event.type === "done" || event.type === "error") break + } + return events + } + + return await Promise.race([ + collect(), + new Promise((_, reject) => { + setTimeout( + () => reject(new Error(`Timed out collecting stream events after ${timeoutMs}ms`)), + timeoutMs, + ) + }), + ]) +} + +export function makeModel(overrides: Partial = {}): ModelLike { + return { + id: "deepseek/deepseek-v4-flash", + api: "commandcode-custom", + provider: "commandcode", + maxTokens: 384_000, + cost: { input: 0.14, output: 0.28, cacheRead: 0.0028, cacheWrite: 0 }, + ...overrides, + } +} + +export function makeContext(overrides: Partial = {}): ContextLike { + return { + systemPrompt: "You are a test assistant.", + messages: [{ role: "user", content: "hello" }], + tools: [], + ...overrides, + } +} + +export interface TestDepsResult { + streamCommandCode: ReturnType + calculatedUsages: Usage[] +} + +export function createTestDeps(overrides: Partial = {}): TestDepsResult { + const calculatedUsages: Usage[] = [] + const streamCommandCode = createStreamCommandCode({ + createStream: createTestEventStream, + calculateCost: (_model, usage) => { + calculatedUsages.push({ + ...usage, + cost: { ...usage.cost }, + }) + }, + env: {}, + authPaths: [], + now: () => new Date("2026-05-05T12:00:00Z").getTime(), + uuid: () => "00000000-0000-4000-8000-000000000000", + cwd: () => "/repo", + delay: async () => {}, + ...overrides, + }) + return { streamCommandCode, calculatedUsages } +} + +type SuccessPlan = { + type: "success" + status?: number + events?: string[] + chunks?: string[] + delays?: number[] + hangAfterLast?: boolean + /** Delay in ms before the server starts sending the response. */ + responseDelay?: number +} + +type ErrorPlan = { + type: "error" + status: number + body: string + headers?: Record +} + +export type ResponsePlan = SuccessPlan | ErrorPlan + +function headersToRecord(headers: IncomingHttpHeaders): Record { + const out: Record = {} + for (const [key, value] of Object.entries(headers)) { + if (typeof value === "string") out[key] = value + else if (Array.isArray(value)) out[key] = value.join(", ") + } + return out +} + +export interface MockCommandCodeServer { + baseUrl(): string + mockResponse(plan: ResponsePlan): void + mockResponseQueue(plans: ResponsePlan[]): void + reset(): void + close(): Promise + lastRequestBody(): unknown + lastRequestHeaders(): Record + requestCount(): number + responseClosedBeforeEnd(): boolean +} + +export async function startMockCommandCodeServer(): Promise { + let planQueue: ResponsePlan[] = [{ type: "success", events: [] }] + let lastBody: unknown + let lastHeaders: Record = {} + let requests = 0 + let closedBeforeEnd = false + let port = 0 + + const server: Server = createServer((req, res) => { + if (req.method !== "POST" || req.url !== "/alpha/generate") { + res.writeHead(404) + res.end("Not found") + return + } + + requests += 1 + lastHeaders = headersToRecord(req.headers) + let body = "" + req.on("data", (chunk: Buffer) => { + body += chunk.toString("utf-8") + }) + req.on("end", () => { + try { + const parsed: unknown = JSON.parse(body) + lastBody = parsed + } catch { + lastBody = undefined + } + + // Pop the first plan from the queue; keep the last one as fallback. + const plan = planQueue.length > 1 ? planQueue.shift()! : planQueue[0] + + if (plan.type === "error") { + const headers: Record = { "Content-Type": "text/plain", ...plan.headers } + res.writeHead(plan.status, headers) + res.end(plan.body) + return + } + + res.writeHead(plan.status ?? 200, { + "Content-Type": "text/plain; charset=utf-8", + "Transfer-Encoding": "chunked", + }) + + let ended = false + res.on("close", () => { + if (!ended) closedBeforeEnd = true + }) + + const chunks = plan.chunks ?? (plan.events ?? []).map((event) => `${event}\n`) + const delays = plan.delays ?? chunks.map(() => 0) + let index = 0 + + const sendNext = () => { + if (index >= chunks.length) { + if (!plan.hangAfterLast) { + ended = true + res.end() + } + return + } + + res.write(chunks[index]) + index += 1 + if (index < chunks.length) { + setTimeout(sendNext, delays[index] ?? 0) + } else if (!plan.hangAfterLast) { + ended = true + res.end() + } + } + + if (plan.responseDelay) { + setTimeout(sendNext, plan.responseDelay) + } else { + sendNext() + } + }) + }) + + await new Promise((resolve) => { + server.listen(0, () => { + const address = server.address() + if (typeof address === "object" && address) port = address.port + resolve() + }) + }) + + return { + baseUrl: () => `http://127.0.0.1:${port}`, + mockResponse(plan: ResponsePlan) { + planQueue = [plan] + }, + mockResponseQueue(plans: ResponsePlan[]) { + planQueue = [...plans] + }, + reset() { + planQueue = [{ type: "success", events: [] }] + lastBody = undefined + lastHeaders = {} + requests = 0 + closedBeforeEnd = false + }, + close() { + return new Promise((resolve) => server.close(() => resolve())) + }, + lastRequestBody: () => lastBody, + lastRequestHeaders: () => lastHeaders, + requestCount: () => requests, + responseClosedBeforeEnd: () => closedBeforeEnd, + } +} + +export function objectAt(value: unknown, path: readonly string[]): unknown { + let current = value + for (const key of path) { + if (Array.isArray(current)) { + const index = Number(key) + if (!Number.isInteger(index)) return undefined + current = current[index] + continue + } + if (typeof current !== "object" || current === null) return undefined + current = Object.getOwnPropertyDescriptor(current, key)?.value + } + return current +} diff --git a/tests/test-abort.ts b/tests/test-abort.ts new file mode 100644 index 0000000..8db0ad5 --- /dev/null +++ b/tests/test-abort.ts @@ -0,0 +1,85 @@ +/** + * Abort tests against the real streamCommandCode core. + */ + +import assert from "node:assert/strict" +import { after, before, beforeEach, describe, it } from "node:test" + +import { + collectEvents, + createTestDeps, + makeContext, + makeModel, + startMockCommandCodeServer, + type MockCommandCodeServer, +} from "./helpers.ts" + +let server: MockCommandCodeServer + +before(async () => { + server = await startMockCommandCodeServer() +}) + +after(async () => { + await server.close() +}) + +beforeEach(() => { + server.reset() +}) + +describe("streamCommandCode — abort behavior", () => { + it("emits aborted error when signal is already aborted", async () => { + const controller = new AbortController() + controller.abort() + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { + apiKey: "mock-key", + signal: controller.signal, + }), + ) + + assert.deepEqual( + events.map((event) => event.type), + ["start", "error"], + ) + const error = events.at(-1) + assert.equal(error?.type, "error") + if (error?.type !== "error") throw new Error("expected error") + assert.equal(error.reason, "aborted") + assert.equal(error.error.stopReason, "aborted") + assert.equal(server.requestCount(), 0) + }) + + it("emits aborted error and cancels the response reader mid-stream", async () => { + server.mockResponse({ + type: "success", + events: [JSON.stringify({ type: "text-delta", text: "first" })], + hangAfterLast: true, + }) + const controller = new AbortController() + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + const stream = streamCommandCode(makeModel(), makeContext(), { + apiKey: "mock-key", + signal: controller.signal, + }) + + setTimeout(() => controller.abort(), 50) + const events = await collectEvents(stream, 2_000) + + assert.ok( + events.some((event) => event.type === "text_delta"), + "stream should process data before abort", + ) + const error = events.at(-1) + assert.equal(error?.type, "error") + if (error?.type !== "error") throw new Error("expected error") + assert.equal(error.reason, "aborted") + assert.equal(error.error.errorMessage, "Request aborted") + await new Promise((resolve) => setTimeout(resolve, 50)) + assert.ok(server.responseClosedBeforeEnd(), "abort should close the hanging upstream response") + }) +}) diff --git a/tests/test-cost.ts b/tests/test-cost.ts new file mode 100644 index 0000000..ae7f730 --- /dev/null +++ b/tests/test-cost.ts @@ -0,0 +1,183 @@ +/** + * Regression test for the local cost calculation. + * + * The provider ships its own cost function because Oh My Pi's legacy pi-ai + * shim does not export `calculateCost` (see issue #24). This test locks the + * local implementation to pi-ai's documented per-million-token arithmetic + * without installing another pi-ai runtime next to the extension. + */ + +import assert from "node:assert/strict" +import { describe, it } from "node:test" + +import { calculateCommandCodeCost } from "../src/cost.ts" +import type { Usage } from "../src/types.ts" + +interface CostRates { + input: number + output: number + cacheRead: number + cacheWrite: number +} + +interface CostTable extends CostRates { + tiers?: Array +} + +const COST_FIXTURES: Record = { + "zero-cost-model": { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + "claude-sonnet-4-6": { input: 3, output: 15, cacheRead: 0.3, cacheWrite: 3.75 }, + "deepseek/deepseek-v4-pro": { + input: 0.435, + output: 0.87, + cacheRead: 0.003625, + cacheWrite: 0, + }, + "Qwen/Qwen3.7-Max": { input: 2.5, output: 7.5, cacheRead: 0.5, cacheWrite: 3.13 }, + "Qwen/Qwen3.7-Flash": { + input: 0.03, + output: 0.13, + cacheRead: 0.006, + cacheWrite: 0.038, + tiers: [ + { inputTokensAbove: 32_000, input: 0.1, output: 0.4, cacheRead: 0.02, cacheWrite: 0.125 }, + { inputTokensAbove: 256_000, input: 0.2, output: 0.8, cacheRead: 0.04, cacheWrite: 0.25 }, + ], + }, + "gpt-5.5": { input: 5, output: 30, cacheRead: 0.5, cacheWrite: 0 }, +} + +const USAGE_CASES = [ + { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + { input: 1, output: 1, cacheRead: 1, cacheWrite: 1 }, + { input: 812, output: 187, cacheRead: 52_000, cacheWrite: 3_100 }, + { input: 1_000_000, output: 65_536, cacheRead: 998_877, cacheWrite: 123_456 }, + { input: 7, output: 999_999_999, cacheRead: 0.5, cacheWrite: 42 }, +] + +function commandCodeModel(id: string, cost: CostTable) { + return { + id, + api: "commandcode-custom", + provider: "commandcode", + cost, + maxTokens: 65_536, + } +} + +function assertClose(actual: number, expected: number) { + assert.ok( + Math.abs(actual - expected) <= + Number.EPSILON * Math.max(1, Math.abs(actual), Math.abs(expected)), + `expected ${actual} to be close to ${expected}`, + ) +} + +function freshUsage(tokens: (typeof USAGE_CASES)[number]): Usage { + return { + ...tokens, + totalTokens: tokens.input + tokens.output + tokens.cacheRead + tokens.cacheWrite, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + } +} + +function expectedCost(cost: CostTable, tokens: (typeof USAGE_CASES)[number]): Usage["cost"] { + const inputTokens = tokens.input + tokens.cacheRead + tokens.cacheWrite + let rates: CostRates = cost + let matchedThreshold = -1 + for (const tier of cost.tiers ?? []) { + if (inputTokens > tier.inputTokensAbove && tier.inputTokensAbove > matchedThreshold) { + rates = tier + matchedThreshold = tier.inputTokensAbove + } + } + + const input = (rates.input / 1_000_000) * tokens.input + const output = (rates.output / 1_000_000) * tokens.output + const cacheRead = (rates.cacheRead / 1_000_000) * tokens.cacheRead + const cacheWrite = (rates.cacheWrite * tokens.cacheWrite) / 1_000_000 + return { + input, + output, + cacheRead, + cacheWrite, + total: input + output + cacheRead + cacheWrite, + } +} + +describe("calculateCommandCodeCost()", () => { + it("applies per-million-token rates to all cost fields", () => { + for (const [id, cost] of Object.entries(COST_FIXTURES)) { + const model = commandCodeModel(id, cost) + + for (const tokens of USAGE_CASES) { + const usage = freshUsage(tokens) + calculateCommandCodeCost(model, usage) + + assert.deepEqual( + usage.cost, + expectedCost(cost, tokens), + `${id} cost for tokens=${JSON.stringify(tokens)}`, + ) + } + } + }) + + it("applies the highest request-wide input tier above its threshold", () => { + const model = commandCodeModel("Qwen/Qwen3.7-Flash", COST_FIXTURES["Qwen/Qwen3.7-Flash"]) + + const atThreshold = freshUsage({ + input: 32_000, + output: 1_000, + cacheRead: 0, + cacheWrite: 0, + }) + calculateCommandCodeCost(model, atThreshold) + assertClose(atThreshold.cost.input, (0.03 * 32_000) / 1_000_000) + + const aboveFirstTier = freshUsage({ + input: 30_000, + output: 1_000, + cacheRead: 2_001, + cacheWrite: 0, + }) + calculateCommandCodeCost(model, aboveFirstTier) + assertClose(aboveFirstTier.cost.input, (0.1 * 30_000) / 1_000_000) + assertClose(aboveFirstTier.cost.cacheRead, (0.02 * 2_001) / 1_000_000) + + const aboveHighestTier = freshUsage({ + input: 100_000, + output: 1_000, + cacheRead: 156_001, + cacheWrite: 0, + }) + calculateCommandCodeCost(model, aboveHighestTier) + assertClose(aboveHighestTier.cost.input, (0.2 * 100_000) / 1_000_000) + assertClose(aboveHighestTier.cost.output, (0.8 * 1_000) / 1_000_000) + }) + + it("prices one-hour cache writes at twice the active input rate", () => { + const model = commandCodeModel("claude-sonnet-4-6", COST_FIXTURES["claude-sonnet-4-6"]) + const usage = freshUsage({ input: 0, output: 0, cacheRead: 0, cacheWrite: 1_000 }) + usage.cacheWrite1h = 400 + + calculateCommandCodeCost(model, usage) + + const expectedShortWrite = (3.75 * 600) / 1_000_000 + const expectedLongWrite = (3 * 2 * 400) / 1_000_000 + assertClose(usage.cost.cacheWrite, expectedShortWrite + expectedLongWrite) + }) + + it("writes the total as the sum of all cost components", () => { + const model = commandCodeModel("claude-sonnet-4-6", COST_FIXTURES["claude-sonnet-4-6"]) + const usage = freshUsage({ input: 1_000, output: 500, cacheRead: 10_000, cacheWrite: 2_000 }) + + calculateCommandCodeCost(model, usage) + + assert.equal( + usage.cost.total, + usage.cost.input + usage.cost.output + usage.cost.cacheRead + usage.cost.cacheWrite, + ) + assert.ok(usage.cost.total > 0) + }) +}) diff --git a/tests/test-overflow.ts b/tests/test-overflow.ts new file mode 100644 index 0000000..dc28888 --- /dev/null +++ b/tests/test-overflow.ts @@ -0,0 +1,196 @@ +import assert from "node:assert/strict" +import { after, before, beforeEach, describe, it } from "node:test" + +import { + commandCodeErrorMessage, + normalizeCommandCodeErrorMessage, + normalizeCommandCodeMessage, +} from "../src/overflow.ts" +import { + collectEvents, + createTestDeps, + makeContext, + makeModel, + startMockCommandCodeServer, + type MockCommandCodeServer, +} from "./helpers.ts" + +let server: MockCommandCodeServer + +before(async () => { + server = await startMockCommandCodeServer() +}) + +after(async () => { + await server.close() +}) + +beforeEach(() => { + server.reset() +}) + +describe("Command Code overflow normalization", () => { + it("normalizes Command Code context errors to pi's generic overflow marker", () => { + const normalized = normalizeCommandCodeErrorMessage("Prompt token limit exceeded") + + assert.equal(normalized, "context_length_exceeded: Prompt token limit exceeded") + }) + + it("is idempotent and leaves unrelated, rate-limit, and capacity errors unchanged", () => { + assert.equal( + normalizeCommandCodeErrorMessage("context_length_exceeded: Prompt token limit exceeded"), + undefined, + ) + assert.equal(normalizeCommandCodeErrorMessage("OpenAI request failed"), undefined) + assert.equal( + normalizeCommandCodeErrorMessage("Prompt token limit exceeded due to rate limit"), + undefined, + ) + assert.equal( + normalizeCommandCodeErrorMessage("Command Code API error 429: context window exceeded"), + undefined, + ) + assert.equal( + normalizeCommandCodeErrorMessage("context window exceeded: status: 429"), + undefined, + ) + assert.equal( + normalizeCommandCodeErrorMessage("The input is too long"), + "context_length_exceeded: The input is too long", + ) + assert.equal( + normalizeCommandCodeErrorMessage("Input exceeds context limit"), + "context_length_exceeded: Input exceeds context limit", + ) + assert.equal( + normalizeCommandCodeErrorMessage("Context window exceeded: provider capacity reached"), + undefined, + ) + }) + + it("scopes finalized message normalization to Command Code", () => { + const message = { + role: "assistant" as const, + provider: "commandcode", + stopReason: "error" as const, + errorMessage: "model context window exceeded", + } + + const normalized = normalizeCommandCodeMessage(message) + assert.equal( + normalized?.message.errorMessage, + "context_length_exceeded: model context window exceeded", + ) + assert.equal(normalizeCommandCodeMessage({ ...message, provider: "openai" }), undefined) + assert.equal(normalizeCommandCodeMessage({ ...message, stopReason: "stop" }), undefined) + }) + + it("extracts nested stream error messages without exposing credentials", () => { + assert.equal( + commandCodeErrorMessage({ + error: { details: { errorMessage: "context window exceeded" } }, + }), + "context window exceeded", + ) + }) + + it("redacts secrets from finalized provider errors", async () => { + server.mockResponse({ + type: "error", + status: 400, + body: "api_key=user_secret_value", + }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { apiKey: "mock-key" }), + ) + const error = events.at(-1) + assert.equal(error?.type, "error") + if (error?.type !== "error") throw new Error("expected error") + assert.doesNotMatch(error.error.errorMessage ?? "", /user_secret_value/) + assert.match(error.error.errorMessage ?? "", /api_key=\[redacted\]/) + }) + + it("normalizes HTTP error bodies containing nested context errors", async () => { + server.mockResponse({ + type: "error", + status: 400, + body: JSON.stringify({ error: { message: "Prompt token limit exceeded" } }), + }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { apiKey: "mock-key" }), + ) + const error = events.at(-1) + + assert.equal(error?.type, "error") + if (error?.type !== "error") throw new Error("expected error") + const normalized = normalizeCommandCodeMessage(error.error) + assert.match(normalized?.message.errorMessage ?? "", /^context_length_exceeded:/) + }) + + it("does not normalize an HTTP rate-limit response that mentions context", async () => { + server.mockResponse({ + type: "error", + status: 429, + body: JSON.stringify({ error: { message: "context window exceeded" } }), + }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { apiKey: "mock-key" }), + ) + const error = events.at(-1) + + assert.equal(error?.type, "error") + if (error?.type !== "error") throw new Error("expected error") + assert.equal(normalizeCommandCodeMessage(error.error), undefined) + }) + + it("normalizes nested stream error events", async () => { + server.mockResponse({ + type: "success", + events: [ + JSON.stringify({ + type: "error", + error: { details: { message: "model context window exceeded" } }, + }), + ], + }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { apiKey: "mock-key" }), + ) + const error = events.at(-1) + + assert.equal(error?.type, "error") + if (error?.type !== "error") throw new Error("expected error") + const normalized = normalizeCommandCodeMessage(error.error) + assert.equal( + normalized?.message.errorMessage, + "context_length_exceeded: model context window exceeded", + ) + + server.mockResponse({ + type: "success", + events: [ + JSON.stringify({ + type: "error", + error: { message: "context window exceeded", status: 429 }, + }), + ], + }) + const retryEvents = await collectEvents( + createTestDeps({ apiBase: server.baseUrl() }).streamCommandCode(makeModel(), makeContext(), { + apiKey: "mock-key", + }), + ) + const retryError = retryEvents.at(-1) + assert.equal(retryError?.type, "error") + if (retryError?.type !== "error") throw new Error("expected error") + assert.equal(normalizeCommandCodeMessage(retryError.error), undefined) + }) +}) diff --git a/tests/test-pure-functions.ts b/tests/test-pure-functions.ts new file mode 100644 index 0000000..136025f --- /dev/null +++ b/tests/test-pure-functions.ts @@ -0,0 +1,658 @@ +/** + * Unit tests for the real pure helpers exported by src/core.ts. + * These are hermetic: no pi runtime and no network. + */ + +import assert from "node:assert/strict" +import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from "node:fs" +import { tmpdir } from "node:os" +import { join } from "node:path" +import { describe, it } from "node:test" + +import { + assertTextOnlyMessages, + getApiKey, + getEnvironmentInfo, + mapFinishReason, + messagesToCC, + parseStreamEventLine, + projectSlugFromPath, + textContent, + toJsonSchema, + toolsToJson, +} from "../src/core.ts" +import { redactCommandCodeErrorText } from "../src/overflow.ts" + +import { objectAt } from "./helpers.ts" + +describe("getApiKey()", () => { + it("uses COMMANDCODE_API_KEY from provided env", () => { + assert.equal(getApiKey({ env: { COMMANDCODE_API_KEY: "env-key" }, authPaths: [] }), "env-key") + }) + + it("reads apiKey, commandcode, pi OAuth, and official CLI credential fields", () => { + const dir = mkdtempSync(join(tmpdir(), "cc-auth-")) + try { + const first = join(dir, "first.json") + const second = join(dir, "second.json") + const oauth = join(dir, "oauth.json") + const official = join(dir, "official.json") + writeFileSync(first, JSON.stringify({ apiKey: "file-key" })) + writeFileSync(second, JSON.stringify({ commandcode: "fallback-key" })) + writeFileSync( + oauth, + JSON.stringify({ + commandcode: { + type: "oauth", + access: "oauth-access-key", + refresh: "oauth-refresh-key", + expires: Date.now() + 3600000, + }, + }), + ) + writeFileSync( + official, + JSON.stringify({ + "command-code": { + type: "api", + key: "official-cli-key", + }, + }), + ) + assert.equal(getApiKey({ env: {}, authPaths: [first, second] }), "file-key") + assert.equal(getApiKey({ env: {}, authPaths: [second] }), "fallback-key") + assert.equal(getApiKey({ env: {}, authPaths: [oauth] }), "oauth-access-key") + assert.equal(getApiKey({ env: {}, authPaths: [official] }), "official-cli-key") + } finally { + rmSync(dir, { recursive: true, force: true }) + } + }) + + it("ignores malformed auth files", () => { + const dir = mkdtempSync(join(tmpdir(), "cc-auth-bad-")) + try { + const bad = join(dir, "bad.json") + writeFileSync(bad, "not json") + assert.equal(getApiKey({ env: {}, authPaths: [bad] }), undefined) + } finally { + rmSync(dir, { recursive: true, force: true }) + } + }) + + it("uses injected homeDir for default auth paths", () => { + const dir = mkdtempSync(join(tmpdir(), "cc-home-")) + try { + const authDir = join(dir, ".pi", "agent") + mkdirSync(authDir, { recursive: true }) + writeFileSync(join(authDir, "auth.json"), JSON.stringify({ commandcode: "pi-key" })) + assert.equal(getApiKey({ env: {}, homeDir: () => dir }), "pi-key") + } finally { + rmSync(dir, { recursive: true, force: true }) + } + }) +}) + +describe("error redaction", () => { + it("redacts bearer, credential, and query-string secrets", () => { + const redacted = redactCommandCodeErrorText( + "Bearer user_secret_value api_key=user_secret_value https://example.test/x?token=user_secret_value", + ) + assert.doesNotMatch(redacted, /user_secret_value/) + assert.match(redacted, /Bearer \[redacted\]/) + assert.doesNotMatch( + redactCommandCodeErrorText("provider returned sk-test-secret-value-1234567890"), + /sk-test-secret-value/, + ) + }) +}) + +describe("projectSlugFromPath()", () => { + it("matches the official CLI-style slug from an absolute working directory", () => { + assert.equal( + projectSlugFromPath("/Users/patwoz/dev/Personal/pi/pi-commandcode-provider"), + "users-patwoz-dev-personal-pi-pi-commandcode-provider", + ) + assert.equal(projectSlugFromPath("/repo"), "repo") + }) +}) + +describe("text-only image handling", () => { + it("rejects image content for models without image support", () => { + assert.throws( + () => + assertTextOnlyMessages([ + { + role: "user", + content: [{ type: "image", data: "base64-data", mimeType: "image/png" }], + }, + ]), + /does not support image content/i, + ) + assert.throws( + () => + assertTextOnlyMessages([ + { + role: "toolResult", + toolCallId: "c1", + content: [{ type: "image", data: "base64-data", mimeType: "image/png" }], + }, + ]), + /does not support image content/i, + ) + }) +}) + +describe("textContent()", () => { + it("extracts and joins text blocks", () => { + assert.equal( + textContent({ + content: [ + { type: "text", text: "hello" }, + { type: "text", text: "world" }, + ], + }), + "hello\nworld", + ) + }) + + it("extracts text while images are handled separately", () => { + assert.equal( + textContent({ + content: [ + { type: "text", text: "hello" }, + { type: "image", data: "x", mimeType: "image/png" }, + { type: "text", text: "world" }, + ], + }), + "hello\nworld", + ) + }) + + it("handles empty or missing content", () => { + assert.equal(textContent({ content: [] }), "") + assert.equal(textContent({}), "") + }) +}) + +describe("getEnvironmentInfo()", () => { + it("returns platform, arch, and Node version", () => { + const info = getEnvironmentInfo() + assert.match(info, /^(darwin|linux|win32)-/) + assert.ok(info.includes("Node.js")) + }) +}) + +describe("toJsonSchema()", () => { + it("converts scalar, enum, object, optional, array, and union schema shapes", () => { + assert.deepEqual(toJsonSchema({ kind: "string" }), { type: "string" }) + assert.deepEqual(toJsonSchema({ kind: "Number" }), { type: "number" }) + assert.deepEqual(toJsonSchema({ kind: "boolean" }), { type: "boolean" }) + assert.deepEqual(toJsonSchema({ kind: "string", enum: ["left", "right"] }), { + type: "string", + enum: ["left", "right"], + }) + assert.deepEqual( + toJsonSchema({ + kind: "object", + properties: { + name: { kind: "string" }, + tags: { kind: "array", items: { kind: "string" }, optional: true }, + }, + }), + { + type: "object", + properties: { + name: { type: "string" }, + tags: { type: "array", items: { type: "string" } }, + }, + required: ["name"], + }, + ) + assert.deepEqual(toJsonSchema({ kind: "optional", wrapped: { kind: "string" } }), { + type: "string", + }) + assert.deepEqual(toJsonSchema({ kind: "union", variants: [{}, { kind: "number" }] }), { + type: "number", + }) + }) + + it("preserves explicit required arrays and handles unknown values", () => { + assert.deepEqual( + toJsonSchema({ + type: "object", + properties: { name: { type: "string" }, nickname: { type: "string" } }, + required: ["name"], + }), + { + type: "object", + properties: { name: { type: "string" }, nickname: { type: "string" } }, + required: ["name"], + }, + ) + assert.deepEqual(toJsonSchema(undefined), {}) + assert.deepEqual(toJsonSchema({ kind: "wat" }), {}) + assert.deepEqual(toJsonSchema({ type: "wat", description: "not a schema" }), {}) + assert.deepEqual(toJsonSchema({}), {}) + assert.equal(toJsonSchema(true), true) + }) + + it("preserves complete JSON Schema metadata and nested schemas", () => { + assert.deepEqual( + toJsonSchema({ + type: "object", + description: "Search options", + properties: { + query: { + type: "string", + description: "Text to search for", + minLength: 2, + maxLength: 50, + pattern: "^[a-z]+$", + default: "pi", + }, + limit: { + type: "integer", + minimum: 1, + maximum: 100, + exclusiveMinimum: 0, + multipleOf: 1, + default: 10, + }, + tags: { + type: "array", + minItems: 1, + maxItems: 3, + uniqueItems: true, + items: { + type: "object", + properties: { name: { type: "string" } }, + required: ["name"], + additionalProperties: false, + }, + }, + }, + required: ["query", "limit"], + additionalProperties: false, + }), + { + type: "object", + description: "Search options", + properties: { + query: { + type: "string", + description: "Text to search for", + minLength: 2, + maxLength: 50, + pattern: "^[a-z]+$", + default: "pi", + }, + limit: { + type: "integer", + minimum: 1, + maximum: 100, + exclusiveMinimum: 0, + multipleOf: 1, + default: 10, + }, + tags: { + type: "array", + minItems: 1, + maxItems: 3, + uniqueItems: true, + items: { + type: "object", + properties: { name: { type: "string" } }, + required: ["name"], + additionalProperties: false, + }, + }, + }, + required: ["query", "limit"], + additionalProperties: false, + }, + ) + }) + + it("preserves JSON Schema composition and nullable forms", () => { + assert.deepEqual( + toJsonSchema({ + anyOf: [{ type: "string" }, { type: "number" }], + oneOf: [{ const: "a" }, { const: "b" }], + allOf: [{ minLength: 1 }, { maxLength: 10 }], + nullable: true, + }), + { + anyOf: [{ type: "string" }, { type: "number" }, { type: "null" }], + oneOf: [{ const: "a" }, { const: "b" }], + allOf: [{ minLength: 1 }, { maxLength: 10 }], + }, + ) + assert.deepEqual(toJsonSchema({ type: ["string", "null"] }), { + type: ["string", "null"], + }) + assert.deepEqual(toJsonSchema({ type: "string", nullable: true }), { + type: ["string", "null"], + }) + }) + + it("preserves dangerous schema property names", () => { + const inputProperties: Record = { + constructor: { type: "number" }, + } + Object.defineProperty(inputProperties, "__proto__", { + configurable: true, + enumerable: true, + value: { type: "string" }, + writable: true, + }) + const schema = toJsonSchema({ + type: "object", + properties: inputProperties, + required: ["__proto__", "constructor"], + }) + assert.ok(schema && typeof schema === "object" && !Array.isArray(schema)) + if (!schema || typeof schema !== "object" || Array.isArray(schema)) { + throw new Error("expected object schema") + } + const outputProperties: unknown = Object.getOwnPropertyDescriptor(schema, "properties")?.value + assert.ok(outputProperties && typeof outputProperties === "object") + if (!outputProperties || typeof outputProperties !== "object") { + throw new Error("expected object properties") + } + assert.ok(Object.prototype.hasOwnProperty.call(outputProperties, "__proto__")) + assert.deepEqual(Object.getOwnPropertyDescriptor(outputProperties, "__proto__")?.value, { + type: "string", + }) + assert.deepEqual(Object.getOwnPropertyDescriptor(outputProperties, "constructor")?.value, { + type: "number", + }) + }) + + it("converts legacy shapes without collapsing unions", () => { + assert.deepEqual( + toJsonSchema({ + kind: "Object", + description: "Legacy options", + properties: { + mode: { + kind: "union", + variants: [ + { kind: "string", enum: ["fast", "safe"] }, + { kind: "string", enum: ["debug"] }, + ], + }, + count: { kind: "Number", minimum: 1, optional: true }, + nested: { + kind: "Array", + element: { kind: "object", properties: { value: { kind: "boolean" } } }, + }, + }, + optional: ["count"], + additionalProperties: false, + }), + { + type: "object", + description: "Legacy options", + properties: { + mode: { + anyOf: [ + { type: "string", enum: ["fast", "safe"] }, + { type: "string", enum: ["debug"] }, + ], + }, + count: { type: "number", minimum: 1 }, + nested: { + type: "array", + items: { + type: "object", + properties: { value: { type: "boolean" } }, + required: ["value"], + }, + }, + }, + required: ["mode", "nested"], + additionalProperties: false, + }, + ) + assert.deepEqual( + toJsonSchema({ + kind: "intersect", + variants: [{ kind: "object", properties: { a: { kind: "string" } } }, { kind: "number" }], + }), + { + allOf: [ + { type: "object", properties: { a: { type: "string" } }, required: ["a"] }, + { type: "number" }, + ], + }, + ) + }) +}) + +describe("toolsToJson()", () => { + it("converts pi tools to Command Code tool JSON", () => { + assert.deepEqual( + toolsToJson([ + { + name: "get_weather", + description: "Get weather", + parameters: { + kind: "object", + properties: { city: { kind: "string" } }, + }, + }, + ]), + [ + { + type: "function", + name: "get_weather", + description: "Get weather", + input_schema: { + type: "object", + properties: { city: { type: "string" } }, + required: ["city"], + }, + }, + ], + ) + }) + + it("returns an empty array for missing tools", () => { + assert.deepEqual(toolsToJson(), []) + }) +}) + +describe("messagesToCC()", () => { + it("converts user, assistant, and tool result messages", () => { + const result = messagesToCC([ + { role: "user", content: "read /tmp/test" }, + { + role: "assistant", + content: [ + { type: "thinking", thinking: "I will read" }, + { type: "text", text: "Sure" }, + { + type: "toolCall", + id: "c1", + name: "read", + arguments: { path: "/tmp/test" }, + }, + ], + }, + { + role: "toolResult", + toolCallId: "c1", + toolName: "read", + isError: false, + content: [ + { type: "text", text: "hello" }, + { type: "text", text: "world" }, + ], + }, + ]) + + assert.equal(objectAt(result, ["0", "role"]), "user") + assert.equal(objectAt(result, ["1", "role"]), "assistant") + assert.equal(objectAt(result, ["1", "content", "0", "type"]), "text") + assert.equal(objectAt(result, ["1", "content", "1", "type"]), "tool-call") + assert.equal(objectAt(result, ["1", "content", "2"]), undefined) + assert.equal(objectAt(result, ["2", "role"]), "tool") + assert.equal(objectAt(result, ["2", "content", "0", "output", "value"]), "hello\nworld") + }) + + it("serializes image inputs in the current Command Code wire format", () => { + assert.deepEqual( + messagesToCC( + [ + { + role: "user", + content: [ + { type: "text", text: "inspect this" }, + { type: "image", data: "aGVsbG8=", mimeType: "image/png" }, + ], + }, + ], + { allowImages: true }, + ), + [ + { + role: "user", + content: [ + { type: "text", text: "inspect this" }, + { + type: "image", + image: "data:image/png;base64,aGVsbG8=", + mimeType: "image/png", + }, + ], + }, + ], + ) + }) + + it("preserves tool-result images as a following user image message", () => { + const result = messagesToCC( + [ + { role: "user", content: "read image" }, + { + role: "assistant", + content: [{ type: "toolCall", id: "c1", name: "read", arguments: {} }], + }, + { + role: "toolResult", + toolCallId: "c1", + toolName: "read", + content: [ + { type: "text", text: "image attached" }, + { type: "image", data: "aGVsbG8=", mimeType: "image/jpeg" }, + ], + }, + ], + { allowImages: true }, + ) + + assert.equal(objectAt(result, ["2", "role"]), "tool") + assert.equal(objectAt(result, ["2", "content", "0", "output", "value"]), "image attached") + assert.deepEqual(objectAt(result, ["3"]), { + role: "user", + content: [ + { + type: "image", + image: "data:image/jpeg;base64,aGVsbG8=", + mimeType: "image/jpeg", + }, + ], + }) + }) + + it("drops previous assistant reasoning while preserving text and tool calls", () => { + const result = messagesToCC([ + { role: "user", content: "first question" }, + { + role: "assistant", + content: [ + { type: "thinking", thinking: "private reasoning from turn one" }, + { type: "text", text: "first answer" }, + ], + }, + { role: "user", content: "follow-up question" }, + ]) + + assert.deepEqual(result, [ + { role: "user", content: "first question" }, + { role: "assistant", content: [{ type: "text", text: "first answer" }] }, + { role: "user", content: "follow-up question" }, + ]) + }) + + it("omits assistant turns that contain only previous reasoning", () => { + const result = messagesToCC([ + { role: "user", content: "first question" }, + { + role: "assistant", + content: [{ type: "thinking", thinking: "private reasoning" }], + }, + { role: "user", content: "follow-up question" }, + ]) + + assert.deepEqual(result, [ + { role: "user", content: "first question" }, + { role: "user", content: "follow-up question" }, + ]) + }) + + it("drops orphaned tool calls that have no matching tool result", () => { + const result = messagesToCC([ + { role: "user", content: "edit a file" }, + { + role: "assistant", + content: [ + { type: "text", text: "I will edit it" }, + { + type: "toolCall", + id: "missing-result", + name: "edit", + arguments: { path: "x" }, + }, + ], + }, + ]) + + assert.equal(objectAt(result, ["1", "role"]), "assistant") + assert.equal(objectAt(result, ["1", "content", "0", "type"]), "text") + assert.equal(objectAt(result, ["1", "content", "1"]), undefined) + }) + + it("handles empty conversations", () => { + assert.deepEqual(messagesToCC([]), []) + }) +}) + +describe("parseStreamEventLine()", () => { + it("parses plain JSON and SSE data lines", () => { + assert.deepEqual(parseStreamEventLine('{"type":"text-delta","text":"x"}'), { + type: "text-delta", + text: "x", + }) + assert.deepEqual(parseStreamEventLine('data: {"type":"finish","finishReason":"stop"}'), { + type: "finish", + finishReason: "stop", + }) + }) + + it("ignores comments, event labels, done markers, and malformed JSON", () => { + assert.equal(parseStreamEventLine(":"), undefined) + assert.equal(parseStreamEventLine("event: message"), undefined) + assert.equal(parseStreamEventLine("data: [DONE]"), undefined) + assert.equal(parseStreamEventLine("not-json"), undefined) + }) +}) + +describe("mapFinishReason()", () => { + it("maps provider finish reasons to pi stop reasons", () => { + assert.equal(mapFinishReason("stop"), "stop") + assert.equal(mapFinishReason("tool-calls"), "toolUse") + assert.equal(mapFinishReason("max_tokens"), "length") + assert.equal(mapFinishReason("max_output_tokens"), "length") + }) +}) diff --git a/tests/test-retry.ts b/tests/test-retry.ts new file mode 100644 index 0000000..e0366de --- /dev/null +++ b/tests/test-retry.ts @@ -0,0 +1,470 @@ +/** + * Tests for retry and timeout behaviour driven by pi settings.json + * retry config (timeoutMs, maxRetries, maxRetryDelayMs). + */ + +import assert from "node:assert/strict" +import { after, before, beforeEach, describe, it } from "node:test" + +import type { AssistantMessageEvent } from "../src/core.ts" +import { + collectEvents, + createTestDeps, + makeContext, + makeModel, + startMockCommandCodeServer, + type MockCommandCodeServer, +} from "./helpers.ts" + +const TEST_API_KEY = "option-key" + +let server: MockCommandCodeServer + +before(async () => { + server = await startMockCommandCodeServer() +}) + +after(async () => { + await server.close() +}) + +beforeEach(() => { + server.reset() +}) + +function eventTypes(events: readonly AssistantMessageEvent[]): string[] { + return events.map((event) => event.type) +} + +describe("streamCommandCode — retry on transient errors", () => { + it("retries on 429 and succeeds on the second attempt", async () => { + server.mockResponseQueue([ + { type: "error", status: 429, body: "rate limited" }, + { + type: "success", + events: [ + JSON.stringify({ type: "text-delta", text: "ok" }), + JSON.stringify({ type: "finish", finishReason: "stop" }), + ], + }, + ]) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { + apiKey: TEST_API_KEY, + maxRetries: 2, + }), + ) + + assert.equal(server.requestCount(), 2) + assert.deepEqual(eventTypes(events), ["start", "text_start", "text_delta", "text_end", "done"]) + const done = events.at(-1) + if (done?.type !== "done") throw new Error("expected done") + assert.equal(done.reason, "stop") + }) + + it("retries on 500 and succeeds on the second attempt", async () => { + server.mockResponseQueue([ + { type: "error", status: 500, body: "internal server error" }, + { + type: "success", + events: [JSON.stringify({ type: "finish", finishReason: "stop" })], + }, + ]) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { + apiKey: TEST_API_KEY, + maxRetries: 2, + }), + ) + + assert.equal(server.requestCount(), 2) + assert.equal(events.at(-1)?.type, "done") + }) + + it("does NOT retry on 400 (non-retryable client error)", async () => { + server.mockResponse({ type: "error", status: 400, body: "bad request" }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { apiKey: TEST_API_KEY }), + ) + + assert.equal(server.requestCount(), 1) + assert.deepEqual(eventTypes(events), ["start", "error"]) + const last = events.at(-1) + if (last?.type !== "error") throw new Error("expected error") + assert.match(last.error.errorMessage ?? "", /400/) + }) + + it("exhausts maxRetries and emits an error", async () => { + server.mockResponse({ type: "error", status: 503, body: "unavailable" }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { + apiKey: TEST_API_KEY, + maxRetries: 3, + }), + ) + + // initial attempt + 3 retries = 4 total + assert.equal(server.requestCount(), 4) + assert.deepEqual(eventTypes(events), ["start", "error"]) + const last503 = events.at(-1) + if (last503?.type !== "error") throw new Error("expected error") + assert.match(last503.error.errorMessage ?? "", /503/) + }) +}) + +describe("streamCommandCode — Retry-After header", () => { + it("respects Retry-After delay in seconds", async () => { + let delayCalled = false + server.mockResponseQueue([ + { + type: "error", + status: 429, + body: "rate limited", + headers: { "retry-after": "2" }, + }, + { + type: "success", + events: [JSON.stringify({ type: "finish", finishReason: "stop" })], + }, + ]) + const { streamCommandCode } = createTestDeps({ + apiBase: server.baseUrl(), + delay: async (ms: number) => { + delayCalled = true + assert.equal(ms, 2000) + }, + }) + + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { + apiKey: TEST_API_KEY, + maxRetries: 2, + }), + ) + + assert.equal(server.requestCount(), 2) + assert.equal(events.at(-1)?.type, "done") + assert.ok(delayCalled, "delay should have been called with Retry-After value") + }) + + it("fails immediately when Retry-After exceeds maxRetryDelayMs", async () => { + server.mockResponse({ + type: "error", + status: 429, + body: "rate limited", + headers: { "retry-after": "300" }, + }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { + apiKey: TEST_API_KEY, + maxRetryDelayMs: 10_000, + }), + ) + + assert.equal(server.requestCount(), 1) + assert.deepEqual(eventTypes(events), ["start", "error"]) + const lastMax = events.at(-1) + if (lastMax?.type !== "error") throw new Error("expected error") + assert.match(lastMax.error.errorMessage ?? "", /exceeds max/) + }) + + it("does not cap Retry-After when maxRetryDelayMs is 0", async () => { + let delayCalled = false + server.mockResponseQueue([ + { + type: "error", + status: 429, + body: "rate limited", + headers: { "retry-after": "120" }, + }, + { + type: "success", + events: [JSON.stringify({ type: "finish", finishReason: "stop" })], + }, + ]) + const { streamCommandCode } = createTestDeps({ + apiBase: server.baseUrl(), + delay: async (ms: number) => { + delayCalled = true + assert.equal(ms, 120_000) + }, + }) + + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { + apiKey: TEST_API_KEY, + maxRetries: 1, + maxRetryDelayMs: 0, + }), + ) + + assert.equal(server.requestCount(), 2) + assert.equal(events.at(-1)?.type, "done") + assert.ok(delayCalled) + }) +}) + +describe("streamCommandCode — timeout", () => { + it("retries on per-attempt timeout and succeeds", async () => { + server.mockResponseQueue([ + { + type: "success", + events: [JSON.stringify({ type: "finish", finishReason: "stop" })], + hangAfterLast: true, + responseDelay: 200, + }, + { + type: "success", + events: [ + JSON.stringify({ type: "text-delta", text: "fast" }), + JSON.stringify({ type: "finish", finishReason: "stop" }), + ], + }, + ]) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { + apiKey: TEST_API_KEY, + timeoutMs: 50, + maxRetries: 2, + }), + 5_000, + ) + + assert.equal(server.requestCount(), 2) + assert.deepEqual(eventTypes(events), ["start", "text_start", "text_delta", "text_end", "done"]) + }) + + it("retries when the response starts but the stream hangs before finish", async () => { + server.mockResponseQueue([ + { + type: "success", + events: [], + hangAfterLast: true, + }, + { + type: "success", + events: [ + JSON.stringify({ type: "text-delta", text: "ok" }), + JSON.stringify({ type: "finish", finishReason: "stop" }), + ], + }, + ]) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { + apiKey: TEST_API_KEY, + timeoutMs: 50, + maxRetries: 2, + }), + 5_000, + ) + + assert.equal(server.requestCount(), 2) + assert.deepEqual(eventTypes(events), ["start", "text_start", "text_delta", "text_end", "done"]) + }) + + it("does NOT retry on timeout after partial text-delta was emitted", async () => { + server.mockResponse({ + type: "success", + events: [JSON.stringify({ type: "text-delta", text: "partial" })], + hangAfterLast: true, + }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { + apiKey: TEST_API_KEY, + timeoutMs: 50, + maxRetries: 2, + }), + 5_000, + ) + + assert.equal(server.requestCount(), 1) + assert.deepEqual(eventTypes(events), ["start", "text_start", "text_delta", "error"]) + }) + + it("emits error when all retry attempts time out", async () => { + server.mockResponse({ + type: "success", + events: [JSON.stringify({ type: "finish", finishReason: "stop" })], + hangAfterLast: true, + responseDelay: 200, + }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { + apiKey: TEST_API_KEY, + timeoutMs: 50, + maxRetries: 1, + }), + 5_000, + ) + + // initial + 1 retry = 2 + assert.equal(server.requestCount(), 2) + assert.deepEqual(eventTypes(events), ["start", "error"]) + const error = events.at(-1) + if (error?.type !== "error") throw new Error("expected error") + assert.match(error.error.errorMessage ?? "", /timed out after 50ms/) + }) +}) + +describe("streamCommandCode — abort cancels retry loop", () => { + it("user abort stops retries immediately", async () => { + server.mockResponse({ type: "error", status: 500, body: "error" }) + const controller = new AbortController() + const { streamCommandCode } = createTestDeps({ + apiBase: server.baseUrl(), + delay: async (_ms: number, signal: AbortSignal) => { + // Abort during the retry delay + controller.abort() + // Simulate the real delay which rejects on abort + return new Promise((_, reject) => { + if (signal.aborted) reject(new DOMException("Aborted", "AbortError")) + signal.addEventListener("abort", () => reject(new DOMException("Aborted", "AbortError"))) + }) + }, + }) + + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { + apiKey: TEST_API_KEY, + signal: controller.signal, + maxRetries: 10, + }), + ) + + // Should only have made 1 request (the initial one), then aborted during delay + assert.equal(server.requestCount(), 1) + assert.deepEqual(eventTypes(events), ["start", "error"]) + const error = events.at(-1) + if (error?.type !== "error") throw new Error("expected error") + assert.equal(error.reason, "aborted") + }) +}) + +describe("streamCommandCode — retry defaults", () => { + it("uses default maxRetries of 0 when not specified", async () => { + server.mockResponse({ type: "error", status: 500, body: "error" }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + await collectEvents(streamCommandCode(makeModel(), makeContext(), { apiKey: TEST_API_KEY })) + + assert.equal(server.requestCount(), 1) + }) + + it("respects maxRetries: 0 (no retries)", async () => { + server.mockResponse({ type: "error", status: 500, body: "error" }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { + apiKey: TEST_API_KEY, + maxRetries: 0, + }), + ) + + assert.equal(server.requestCount(), 1) + assert.deepEqual(eventTypes(events), ["start", "error"]) + }) +}) + +describe("streamCommandCode — stream-level error retry", () => { + it("retries when API returns 200 OK but stream contains an error event", async () => { + server.mockResponseQueue([ + { + type: "success", + events: [ + JSON.stringify({ + type: "error", + error: "Service temporarily unavailable. Please try again shortly.", + }), + ], + }, + { + type: "success", + events: [ + JSON.stringify({ type: "text-delta", text: "ok" }), + JSON.stringify({ type: "finish", finishReason: "stop" }), + ], + }, + ]) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { + apiKey: TEST_API_KEY, + maxRetries: 2, + }), + ) + + assert.equal(server.requestCount(), 2) + assert.deepEqual(eventTypes(events), ["start", "text_start", "text_delta", "text_end", "done"]) + }) + + it("exhausts retries on persistent stream-level errors", async () => { + server.mockResponse({ + type: "success", + events: [ + JSON.stringify({ + type: "error", + error: "Service temporarily unavailable. Please try again shortly.", + }), + ], + }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { + apiKey: TEST_API_KEY, + maxRetries: 3, + }), + ) + + // initial + 3 retries = 4 + assert.equal(server.requestCount(), 4) + assert.deepEqual(eventTypes(events), ["start", "error"]) + const last = events.at(-1) + if (last?.type !== "error") throw new Error("expected error") + assert.match(last.error.errorMessage ?? "", /temporarily unavailable/) + }) + + it("does NOT retry stream error when content was already emitted", async () => { + server.mockResponseQueue([ + { + type: "success", + events: [ + JSON.stringify({ type: "text-delta", text: "partial" }), + JSON.stringify({ + type: "error", + error: "Service temporarily unavailable", + }), + ], + }, + ]) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { apiKey: TEST_API_KEY }), + ) + + // Only 1 request — no retry because content was already emitted. + assert.equal(server.requestCount(), 1) + assert.deepEqual(eventTypes(events), ["start", "text_start", "text_delta", "error"]) + }) +}) diff --git a/tests/test-stream.ts b/tests/test-stream.ts new file mode 100644 index 0000000..e30ac49 --- /dev/null +++ b/tests/test-stream.ts @@ -0,0 +1,771 @@ +/** + * Integration tests for the real streamCommandCode core using a local mock + * Command Code server. No real API key or pi runtime required. + */ + +import assert from "node:assert/strict" +import { after, before, beforeEach, describe, it } from "node:test" + +import type { AssistantMessageEvent } from "../src/core.ts" +import { MODEL_EFFORTS, thinkingLevelMapForEfforts } from "../src/models.ts" +import { + collectEvents, + createTestDeps, + makeContext, + makeModel, + objectAt, + startMockCommandCodeServer, + type MockCommandCodeServer, +} from "./helpers.ts" + +let server: MockCommandCodeServer + +before(async () => { + server = await startMockCommandCodeServer() +}) + +after(async () => { + await server.close() +}) + +beforeEach(() => { + server.reset() +}) + +function eventTypes(events: readonly AssistantMessageEvent[]): string[] { + return events.map((event) => event.type) +} + +describe("streamCommandCode — auth", () => { + it("emits a missing-key error without touching the network", async () => { + const { streamCommandCode } = createTestDeps({ + apiBase: server.baseUrl(), + env: {}, + authPaths: [], + }) + const stream = streamCommandCode(makeModel(), makeContext(), { + apiKey: "", + }) + const events = await collectEvents(stream) + + assert.deepEqual(eventTypes(events), ["error"]) + assert.equal(events[0].type, "error") + assert.equal(events[0].reason, "error") + assert.match(events[0].error.errorMessage ?? "", /No Command Code API key/) + assert.equal(server.requestCount(), 0) + }) + + it("ignores the literal env-var name and falls back to env", async () => { + server.mockResponse({ + type: "success", + events: [JSON.stringify({ type: "finish", finishReason: "stop" })], + }) + const { streamCommandCode } = createTestDeps({ + apiBase: server.baseUrl(), + env: { COMMANDCODE_API_KEY: "env-key" }, + }) + + await collectEvents( + streamCommandCode(makeModel(), makeContext(), { apiKey: "COMMANDCODE_API_KEY" }), + ) + + assert.equal( + server.lastRequestHeaders().authorization, + "Bearer env-key", + "should resolve from env, not send the literal var name as the token", + ) + }) + + it("uses options.apiKey in the Authorization header", async () => { + server.mockResponse({ + type: "success", + events: [JSON.stringify({ type: "finish", finishReason: "stop" })], + }) + const { streamCommandCode } = createTestDeps({ + apiBase: server.baseUrl(), + env: { COMMANDCODE_API_KEY: "env-key" }, + }) + + await collectEvents(streamCommandCode(makeModel(), makeContext(), { apiKey: "option-key" })) + + assert.equal(server.lastRequestHeaders().authorization, "Bearer option-key") + }) +}) + +describe("streamCommandCode — successful streams", () => { + it("emits start → text events → done and accumulates usage", async () => { + server.mockResponse({ + type: "success", + events: [ + JSON.stringify({ type: "text-delta", text: "Hel" }), + JSON.stringify({ type: "text-delta", text: "lo" }), + JSON.stringify({ + type: "finish", + finishReason: "stop", + totalUsage: { + inputTokens: 3124, + outputTokens: 15, + inputTokenDetails: { noCacheTokens: 52, cacheReadTokens: 3072 }, + }, + }), + ], + }) + const { streamCommandCode, calculatedUsages } = createTestDeps({ + apiBase: server.baseUrl(), + }) + + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { apiKey: "mock-key" }), + ) + + assert.deepEqual(eventTypes(events), [ + "start", + "text_start", + "text_delta", + "text_delta", + "text_end", + "done", + ]) + const done = events.at(-1) + assert.equal(done?.type, "done") + if (done?.type !== "done") throw new Error("expected done") + assert.equal(done.reason, "stop") + assert.equal(done.message.content[0]?.type, "text") + assert.equal( + done.message.content[0]?.type === "text" ? done.message.content[0].text : "", + "Hello", + ) + assert.equal(done.message.usage.input, 52) + assert.equal(done.message.usage.cacheRead, 3072) + assert.equal(done.message.usage.cacheWrite, 0) + assert.equal(done.message.usage.totalTokens, 3139) + assert.equal(calculatedUsages.length, 1) + }) + + it("sends images for vision-capable models", async () => { + server.mockResponse({ + type: "success", + events: [JSON.stringify({ type: "finish", finishReason: "stop" })], + }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + await collectEvents( + streamCommandCode( + makeModel({ id: "gpt-5.6-luna" }), + makeContext({ + messages: [ + { + role: "user", + content: [ + { type: "text", text: "inspect" }, + { type: "image", data: "aGVsbG8=", mimeType: "image/png" }, + ], + }, + ], + }), + { apiKey: "mock-key" }, + ), + ) + + assert.equal( + objectAt(server.lastRequestBody(), ["params", "messages", "0", "content", "1", "image"]), + "data:image/png;base64,aGVsbG8=", + ) + }) + + it("rejects images before network access for text-only models", async () => { + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + const events = await collectEvents( + streamCommandCode( + makeModel({ id: "deepseek/deepseek-v4-pro" }), + makeContext({ + messages: [ + { + role: "user", + content: [{ type: "image", data: "aGVsbG8=", mimeType: "image/png" }], + }, + ], + }), + { apiKey: "mock-key" }, + ), + ) + + assert.equal(events.at(-1)?.type, "error") + assert.equal(server.requestCount(), 0) + }) + + it("derives uncached input when noCacheTokens is missing", async () => { + server.mockResponse({ + type: "success", + events: [ + JSON.stringify({ + type: "finish", + finishReason: "stop", + totalUsage: { + inputTokens: 100, + outputTokens: 10, + inputTokenDetails: { cacheReadTokens: 75 }, + }, + }), + ], + }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { apiKey: "mock-key" }), + ) + + const done = events.at(-1) + assert.equal(done?.type, "done") + if (done?.type !== "done") throw new Error("expected done") + assert.equal(done.message.usage.input, 25) + assert.equal(done.message.usage.cacheRead, 75) + assert.equal(done.message.usage.cacheWrite, 0) + assert.equal(done.message.usage.totalTokens, 110) + }) + + it("accounts for cache writes separately from uncached input", async () => { + server.mockResponse({ + type: "success", + events: [ + JSON.stringify({ + type: "finish", + finishReason: "stop", + totalUsage: { + inputTokens: 100, + outputTokens: 10, + inputTokenDetails: { + noCacheTokens: 20, + cacheReadTokens: 70, + cacheWriteTokens: 10, + }, + }, + }), + ], + }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { apiKey: "mock-key" }), + ) + + const done = events.at(-1) + assert.equal(done?.type, "done") + if (done?.type !== "done") throw new Error("expected done") + assert.equal(done.message.usage.input, 20) + assert.equal(done.message.usage.cacheRead, 70) + assert.equal(done.message.usage.cacheWrite, 10) + assert.equal(done.message.usage.totalTokens, 110) + }) + + it("ends on finish without waiting for an open upstream connection", async () => { + server.mockResponse({ + type: "success", + events: [ + JSON.stringify({ type: "text-delta", text: "done" }), + JSON.stringify({ type: "finish", finishReason: "stop" }), + ], + hangAfterLast: true, + }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { apiKey: "mock-key" }), + 500, + ) + + assert.equal(events.at(-1)?.type, "done") + await new Promise((resolve) => setTimeout(resolve, 50)) + assert.ok(server.responseClosedBeforeEnd(), "client should cancel the still-open response body") + }) + + it("emits reasoning and tool-call blocks in order", async () => { + server.mockResponse({ + type: "success", + events: [ + JSON.stringify({ type: "reasoning-start" }), + JSON.stringify({ type: "reasoning-delta", text: "think" }), + JSON.stringify({ type: "reasoning-end" }), + JSON.stringify({ type: "text-delta", text: "Using tool" }), + JSON.stringify({ + type: "tool-call", + toolCallId: "call_1", + toolName: "read_file", + input: JSON.stringify({ path: "/tmp/x" }), + }), + JSON.stringify({ type: "finish", finishReason: "tool-calls" }), + ], + }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { apiKey: "mock-key" }), + ) + + assert.deepEqual(eventTypes(events), [ + "start", + "thinking_start", + "thinking_delta", + "thinking_end", + "text_start", + "text_delta", + "text_end", + "toolcall_start", + "toolcall_end", + "done", + ]) + const done = events.at(-1) + if (done?.type !== "done") throw new Error("expected done") + assert.equal(done.reason, "toolUse") + assert.deepEqual( + done.message.content.map((content) => content.type), + ["thinking", "text", "toolCall"], + ) + const toolCall = done.message.content[2] + assert.equal(toolCall?.type === "toolCall" ? toolCall.name : "", "read_file") + }) + + it("flushes reasoning if finish arrives without reasoning-end", async () => { + server.mockResponse({ + type: "success", + events: [ + JSON.stringify({ type: "reasoning-delta", text: "unfinished thought" }), + JSON.stringify({ type: "finish", finishReason: "stop" }), + ], + }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { apiKey: "mock-key" }), + ) + + const done = events.at(-1) + if (done?.type !== "done") throw new Error("expected done") + assert.equal(done.message.content[0]?.type, "thinking") + }) + + it("closes thinking block before text when reasoning-end is missing", async () => { + server.mockResponse({ + type: "success", + events: [ + JSON.stringify({ type: "reasoning-start" }), + JSON.stringify({ type: "reasoning-delta", text: "thinking" }), + JSON.stringify({ type: "text-delta", text: "answer" }), + JSON.stringify({ type: "finish", finishReason: "stop" }), + ], + }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { apiKey: "mock-key" }), + ) + + assert.deepEqual(eventTypes(events), [ + "start", + "thinking_start", + "thinking_delta", + "thinking_end", + "text_start", + "text_delta", + "text_end", + "done", + ]) + }) + + it("closes thinking block before tool-call when reasoning-end is missing", async () => { + server.mockResponse({ + type: "success", + events: [ + JSON.stringify({ type: "reasoning-start" }), + JSON.stringify({ type: "reasoning-delta", text: "thinking" }), + JSON.stringify({ + type: "tool-call", + toolCallId: "call_1", + toolName: "read_file", + input: JSON.stringify({ path: "/tmp/x" }), + }), + JSON.stringify({ type: "finish", finishReason: "tool-calls" }), + ], + }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { apiKey: "mock-key" }), + ) + + assert.deepEqual(eventTypes(events), [ + "start", + "thinking_start", + "thinking_delta", + "thinking_end", + "toolcall_start", + "toolcall_end", + "done", + ]) + }) +}) + +describe("streamCommandCode — request serialization", () => { + it("rejects image input before sending a lossy request", async () => { + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + const events = await collectEvents( + streamCommandCode( + makeModel(), + makeContext({ + messages: [ + { + role: "user", + content: [{ type: "image", data: "base64-data", mimeType: "image/png" }], + }, + ], + }), + { apiKey: "mock-key" }, + ), + ) + + assert.deepEqual(eventTypes(events), ["start", "error"]) + const lastEvent = events.at(-1) + assert.equal(lastEvent?.type, "error") + if (lastEvent?.type === "error") { + assert.match(lastEvent.error.errorMessage ?? "", /does not support image content/i) + } + assert.equal(server.requestCount(), 0) + }) + it("sends the expected request body and default headers", async () => { + server.mockResponse({ + type: "success", + events: [JSON.stringify({ type: "finish", finishReason: "stop" })], + }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + const context = makeContext({ + messages: [ + { role: "user", content: "first" }, + { + role: "assistant", + content: [{ type: "text", text: "first response" }], + }, + { role: "user", content: "second" }, + ], + tools: [ + { + name: "get_weather", + description: "Get weather", + parameters: { + kind: "object", + properties: { city: { kind: "string" } }, + }, + }, + ], + }) + + await collectEvents( + streamCommandCode(makeModel(), context, { + apiKey: "mock-key", + maxTokens: 500, + }), + ) + + const body = server.lastRequestBody() + assert.equal(objectAt(body, ["config", "workingDir"]), "/repo") + assert.equal(objectAt(body, ["config", "date"]), "2026-05-05") + assert.equal(objectAt(body, ["params", "model"]), "deepseek/deepseek-v4-flash") + assert.equal(objectAt(body, ["params", "stream"]), true) + assert.equal(objectAt(body, ["params", "max_tokens"]), 500) + assert.equal(objectAt(body, ["params", "reasoning_effort"]), undefined) + assert.equal(objectAt(body, ["params", "temperature"]), 0.3) + assert.equal(objectAt(body, ["params", "system"]), "You are a test assistant.") + assert.equal(objectAt(body, ["memory"]), null) + assert.equal(objectAt(body, ["taste"]), null) + assert.equal(objectAt(body, ["skills"]), null) + assert.equal(objectAt(body, ["permissionMode"]), undefined) + assert.equal(objectAt(body, ["threadId"]), "00000000-0000-4000-8000-000000000000") + assert.equal( + objectAt(body, ["params", "messages", "1", "content", "0", "text"]), + "first response", + ) + assert.equal(objectAt(body, ["params", "tools", "0", "name"]), "get_weather") + + const headers = server.lastRequestHeaders() + assert.equal(headers.authorization, "Bearer mock-key") + assert.equal(headers["x-command-code-version"], "1.15.1") + assert.equal(headers["x-project-slug"], "repo") + assert.equal(headers["x-taste-learning"], "true") + assert.equal(headers["x-co-flag"], "false") + assert.equal(headers["x-session-id"], undefined) + }) + + it("accepts the legacy OMP nested reasoning map", async () => { + server.mockResponse({ + type: "success", + events: [JSON.stringify({ type: "finish", finishReason: "stop" })], + }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + const model = makeModel({ + id: "omp-compat-reasoning-model", + reasoning: true, + thinking: { effortMap: { high: "legacy-high" } }, + }) + + await collectEvents( + streamCommandCode(model, makeContext(), { apiKey: "mock-key", reasoning: "high" }), + ) + + assert.equal(objectAt(server.lastRequestBody(), ["params", "reasoning_effort"]), "legacy-high") + }) + + it("forwards a supported Pi reasoning level as reasoning_effort", async () => { + server.mockResponse({ + type: "success", + events: [JSON.stringify({ type: "finish", finishReason: "stop" })], + }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + const model = makeModel({ + id: "deepseek/deepseek-v4-flash", + reasoning: true, + thinkingLevelMap: thinkingLevelMapForEfforts(MODEL_EFFORTS["deepseek/deepseek-v4-flash"]), + }) + + await collectEvents( + streamCommandCode(model, makeContext(), { apiKey: "mock-key", reasoning: "max" }), + ) + + assert.equal(objectAt(server.lastRequestBody(), ["params", "reasoning_effort"]), "max") + }) + + it("omits reasoning_effort for off, unsupported, and unknown reasoning levels", async () => { + const model = makeModel({ + id: "deepseek/deepseek-v4-flash", + reasoning: true, + thinkingLevelMap: thinkingLevelMapForEfforts(MODEL_EFFORTS["deepseek/deepseek-v4-flash"]), + }) + + for (const reasoning of ["off", "low"] as const) { + server.mockResponse({ + type: "success", + events: [JSON.stringify({ type: "finish", finishReason: "stop" })], + }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + await collectEvents( + streamCommandCode(model, makeContext(), { apiKey: "mock-key", reasoning }), + ) + assert.equal( + objectAt(server.lastRequestBody(), ["params", "reasoning_effort"]), + undefined, + `${reasoning} should not be sent when it has no supported Command Code field`, + ) + server.reset() + } + + server.mockResponse({ + type: "success", + events: [JSON.stringify({ type: "finish", finishReason: "stop" })], + }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + await collectEvents( + streamCommandCode( + makeModel({ id: "new-model-without-metadata", reasoning: false }), + makeContext(), + { apiKey: "mock-key", reasoning: "high" }, + ), + ) + assert.equal(objectAt(server.lastRequestBody(), ["params", "reasoning_effort"]), undefined) + }) + + it("caps maxTokens and passes custom headers", async () => { + server.mockResponse({ + type: "success", + events: [JSON.stringify({ type: "finish", finishReason: "stop" })], + }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + await collectEvents( + streamCommandCode(makeModel({ maxTokens: 500_000 }), makeContext(), { + apiKey: "mock-key", + maxTokens: 500_000, + headers: { "x-custom": "value" }, + }), + ) + + assert.equal(objectAt(server.lastRequestBody(), ["params", "max_tokens"]), 64_000) + assert.equal(server.lastRequestHeaders()["x-custom"], "value") + }) + + it("caps default maxTokens by the selected model", async () => { + server.mockResponse({ + type: "success", + events: [JSON.stringify({ type: "finish", finishReason: "stop" })], + }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + await collectEvents( + streamCommandCode(makeModel({ maxTokens: 8_192 }), makeContext(), { + apiKey: "mock-key", + }), + ) + + assert.equal(objectAt(server.lastRequestBody(), ["params", "max_tokens"]), 8_192) + }) + + it("serializes OMP system prompt arrays as a string", async () => { + server.mockResponse({ + type: "success", + events: [JSON.stringify({ type: "finish", finishReason: "stop" })], + }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + await collectEvents( + streamCommandCode( + makeModel(), + makeContext({ + systemPrompt: ["You are a test assistant.", "Use concise answers."] as unknown as string, + }), + { apiKey: "mock-key" }, + ), + ) + + assert.equal( + objectAt(server.lastRequestBody(), ["params", "system"]), + "You are a test assistant.\n\nUse concise answers.", + ) + }) + + it("times out a hung onResponse callback", async () => { + server.mockResponse({ + type: "success", + events: [JSON.stringify({ type: "finish", finishReason: "stop" })], + }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + const started = Date.now() + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { + apiKey: "mock-key", + timeoutMs: 25, + onResponse: async () => new Promise(() => {}), + }), + 1_000, + ) + + assert.ok(Date.now() - started < 500) + assert.deepEqual(eventTypes(events), ["start", "error"]) + const error = events.at(-1) + assert.equal(error?.type, "error") + if (error?.type !== "error") throw new Error("expected error") + assert.match(error.error.errorMessage ?? "", /timed out after 25ms/) + }) + + it("times out a hung onPayload callback", async () => { + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + const started = Date.now() + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { + apiKey: "mock-key", + timeoutMs: 25, + onPayload: async () => new Promise(() => {}), + }), + 1_000, + ) + + assert.ok(Date.now() - started < 500) + assert.deepEqual(eventTypes(events), ["start", "error"]) + const error = events.at(-1) + assert.equal(error?.type, "error") + if (error?.type !== "error") throw new Error("expected error") + assert.match(error.error.errorMessage ?? "", /timed out after 25ms/) + assert.equal(server.requestCount(), 0) + }) + + it("runs onPayload and onResponse hooks", async () => { + server.mockResponse({ + type: "success", + events: [JSON.stringify({ type: "finish", finishReason: "stop" })], + }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + let responseStatus = 0 + + await collectEvents( + streamCommandCode(makeModel(), makeContext(), { + apiKey: "mock-key", + onPayload: () => ({ replaced: true }), + onResponse: (response) => { + responseStatus = response.status + }, + }), + ) + + assert.equal(objectAt(server.lastRequestBody(), ["replaced"]), true) + assert.equal(responseStatus, 200) + }) +}) + +describe("streamCommandCode — upstream errors and malformed streams", () => { + it("emits error for HTTP failures", async () => { + server.mockResponse({ type: "error", status: 429, body: "rate limited" }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { apiKey: "mock-key" }), + ) + + assert.deepEqual(eventTypes(events), ["start", "error"]) + const error = events.at(-1) + assert.equal(error?.type, "error") + if (error?.type !== "error") throw new Error("expected error") + assert.match(error.error.errorMessage ?? "", /429/) + }) + + it("emits error for provider error events", async () => { + server.mockResponse({ + type: "success", + events: [ + JSON.stringify({ + type: "error", + error: { message: "provider failed" }, + }), + ], + }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { apiKey: "mock-key" }), + ) + + const error = events.at(-1) + assert.equal(error?.type, "error") + if (error?.type !== "error") throw new Error("expected error") + assert.equal(error.error.errorMessage, "provider failed") + }) + + it("handles SSE lines, malformed lines, split chunks, and final line without newline", async () => { + const textEvent = `data: ${JSON.stringify({ type: "text-delta", text: "split" })}\n` + const finishEvent = JSON.stringify({ + type: "finish", + finishReason: "max_tokens", + }) + server.mockResponse({ + type: "success", + chunks: [ + "not json\n", + textEvent.slice(0, 12), + textEvent.slice(12), + "event: ignored\n", + "data: [DONE]\n", + finishEvent, + ], + }) + const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() }) + + const events = await collectEvents( + streamCommandCode(makeModel(), makeContext(), { apiKey: "mock-key" }), + ) + + const done = events.at(-1) + if (done?.type !== "done") throw new Error("expected done") + assert.equal(done.reason, "length") + assert.equal( + done.message.content[0]?.type === "text" ? done.message.content[0].text : "", + "split", + ) + }) +}) diff --git a/tests/test-transport.ts b/tests/test-transport.ts new file mode 100644 index 0000000..59ea51b --- /dev/null +++ b/tests/test-transport.ts @@ -0,0 +1,183 @@ +import assert from "node:assert/strict" +import { describe, it } from "node:test" + +import { createCommandCodeTransportRouter } from "../src/transport.ts" +import type { + AssistantMessageEvent, + AssistantMessageEventStreamLike, + StreamOptions, +} from "../src/types.ts" +import { collectEvents, createTestEventStream, makeContext, makeModel } from "./helpers.ts" + +function completedStream(text: string): AssistantMessageEventStreamLike { + const stream = createTestEventStream() + const model = makeModel() + const message = { + role: "assistant" as const, + content: [{ type: "text" as const, text }], + api: model.api, + provider: model.provider, + model: model.id, + usage: { + input: 1, + output: 1, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 2, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop" as const, + timestamp: Date.now(), + } + const events: AssistantMessageEvent[] = [ + { type: "start", partial: message }, + { type: "text_start", contentIndex: 0, partial: message }, + { type: "text_delta", contentIndex: 0, delta: text, partial: message }, + { type: "text_end", contentIndex: 0, content: text, partial: message }, + { type: "done", reason: "stop", message }, + ] + for (const event of events) stream.push(event) + stream.end() + return stream +} + +function providerStream( + response: Response, + text: string, + options?: StreamOptions, +): AssistantMessageEventStreamLike { + const stream = createTestEventStream() + const run = async () => { + const received = await (options?.fetch ?? fetch)("https://provider.test", {}) + await options?.onResponse?.( + { status: received.status, headers: {} }, + makeModel({ api: "openai-completions" }), + ) + const source = completedStream(text) + for await (const event of source) stream.push(event) + stream.end() + } + run().catch(() => stream.end()) + return stream +} + +describe("Command Code transport router", () => { + it("keeps using the Provider API after a successful request", async () => { + let providerCalls = 0 + let generateCalls = 0 + const router = createCommandCodeTransportRouter({ + createStream: createTestEventStream, + streamProvider: (_model, _context, options) => { + providerCalls += 1 + return providerStream(new Response("ok", { status: 200 }), "provider", options) + }, + streamGenerate: () => { + generateCalls += 1 + return completedStream("generate") + }, + }) + + const options: StreamOptions = { + fetch: () => Promise.resolve(new Response("ok", { status: 200 })), + } + const first = await collectEvents(router.stream(makeModel(), makeContext(), options)) + const second = await collectEvents(router.stream(makeModel(), makeContext(), options)) + + assert.equal(first.at(-1)?.type, "done") + assert.equal(second.at(-1)?.type, "done") + assert.equal(router.getTransport(), "provider") + assert.equal(providerCalls, 2) + assert.equal(generateCalls, 0) + }) + + it("falls back only for 403 upgrade_required and remembers generate", async () => { + let providerCalls = 0 + let generateCalls = 0 + const responseBody = JSON.stringify({ + error: { code: "upgrade_required", type: "permission_error" }, + }) + const router = createCommandCodeTransportRouter({ + createStream: createTestEventStream, + streamProvider: (_model, _context, options) => { + providerCalls += 1 + return providerStream(new Response(responseBody, { status: 403 }), "blocked", options) + }, + streamGenerate: () => { + generateCalls += 1 + return completedStream("generate") + }, + }) + const options: StreamOptions = { + fetch: () => Promise.resolve(new Response(responseBody, { status: 403 })), + } + + const first = await collectEvents(router.stream(makeModel(), makeContext(), options)) + const second = await collectEvents(router.stream(makeModel(), makeContext(), options)) + + assert.equal(first.at(-1)?.type, "done") + assert.equal(second.at(-1)?.type, "done") + assert.equal(router.getTransport(), "generate") + assert.equal(providerCalls, 1) + assert.equal(generateCalls, 2) + }) + + it("re-detects the transport after the API key changes", async () => { + let providerCalls = 0 + let generateCalls = 0 + const upgradeBody = JSON.stringify({ error: { code: "upgrade_required" } }) + const router = createCommandCodeTransportRouter({ + createStream: createTestEventStream, + streamProvider: (_model, _context, options) => { + providerCalls += 1 + const response = + options?.apiKey === "go-key" + ? new Response(upgradeBody, { status: 403 }) + : new Response("ok", { status: 200 }) + return providerStream(response, "provider", options) + }, + streamGenerate: () => { + generateCalls += 1 + return completedStream("generate") + }, + }) + + await collectEvents( + router.stream(makeModel(), makeContext(), { + apiKey: "go-key", + fetch: () => Promise.resolve(new Response(upgradeBody, { status: 403 })), + }), + ) + await collectEvents( + router.stream(makeModel(), makeContext(), { + apiKey: "provider-key", + fetch: () => Promise.resolve(new Response("ok", { status: 200 })), + }), + ) + + assert.equal(router.getTransport(), "provider") + assert.equal(providerCalls, 2) + assert.equal(generateCalls, 1) + }) + + it("does not fall back for other 403 errors", async () => { + let generateCalls = 0 + const responseBody = JSON.stringify({ error: { code: "permission_denied" } }) + const router = createCommandCodeTransportRouter({ + createStream: createTestEventStream, + streamProvider: (_model, _context, options) => + providerStream(new Response(responseBody, { status: 403 }), "blocked", options), + streamGenerate: () => { + generateCalls += 1 + return completedStream("generate") + }, + }) + const options: StreamOptions = { + fetch: () => Promise.resolve(new Response(responseBody, { status: 403 })), + } + + await collectEvents(router.stream(makeModel(), makeContext(), options)) + + assert.equal(router.getTransport(), "provider") + assert.equal(generateCalls, 0) + }) +})