feat(api): fall back for go plan accounts
This commit is contained in:
Patrick Wozniak
@@ -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<string, string> | undefined {
|
||||
if (process.env.COMMANDCODE_ZDR === "1") {
|
||||
@@ -38,6 +44,7 @@ function commandCodeHeaders(): Record<string, string> | 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<ProviderConfig, ExtensionCommandContext>(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()
|
||||
|
||||
+7
-2
@@ -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",
|
||||
|
||||
@@ -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<string, unknown> {
|
||||
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<string, unknown>[] {
|
||||
if (!Array.isArray(value)) return []
|
||||
return value.filter(isRecord)
|
||||
}
|
||||
|
||||
export function recordOrEmpty(value: unknown): Record<string, unknown> {
|
||||
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<string, unknown>[] {
|
||||
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<string, unknown>): Record<string, string> {
|
||||
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<string> {
|
||||
const callIds = new Set<string>()
|
||||
const resultIds = new Set<string>()
|
||||
|
||||
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)
|
||||
}
|
||||
+741
@@ -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<string, unknown>): Record<string, unknown> | undefined {
|
||||
return isRecord(event.totalUsage) ? event.totalUsage : undefined
|
||||
}
|
||||
|
||||
function commandCodeInputTokenDetails(
|
||||
usage: Record<string, unknown>,
|
||||
): Record<string, unknown> | undefined {
|
||||
return isRecord(usage.inputTokenDetails) ? usage.inputTokenDetails : undefined
|
||||
}
|
||||
|
||||
function headersToRecord(headers: Headers): Record<string, string> {
|
||||
const out: Record<string, string> = {}
|
||||
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<void>((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<T>(promise: Promise<T>, signal: AbortSignal): Promise<T> {
|
||||
if (signal.aborted) return Promise.reject(abortError())
|
||||
|
||||
return new Promise<T>((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<T>(
|
||||
promise: Promise<T>,
|
||||
controller: AbortController,
|
||||
timeoutMs: number | undefined,
|
||||
): Promise<T> {
|
||||
if (timeoutMs === undefined) return raceAbort(promise, controller.signal)
|
||||
|
||||
return new Promise<T>((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<Uint8Array> | 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<typeof setTimeout> | 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 = <T>(promise: Promise<T>): Promise<T> =>
|
||||
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
|
||||
}
|
||||
}
|
||||
+31
@@ -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
|
||||
}
|
||||
@@ -0,0 +1,382 @@
|
||||
function isRecord(value: unknown): value is Record<string, unknown> {
|
||||
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<string, unknown>
|
||||
|
||||
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, unknown>): 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<string, unknown>): 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<object>): JsonSchemaValue {
|
||||
if (typeof value === "boolean") return value
|
||||
if (!isRecord(value)) return {}
|
||||
return convertSchema(value, seen)
|
||||
}
|
||||
|
||||
function setSchemaProperty(target: Record<string, unknown>, key: string, value: unknown): void {
|
||||
Object.defineProperty(target, key, {
|
||||
configurable: true,
|
||||
enumerable: true,
|
||||
value,
|
||||
writable: true,
|
||||
})
|
||||
}
|
||||
|
||||
function schemaMap(value: unknown, seen: WeakSet<object>): Record<string, unknown> {
|
||||
if (!isRecord(value)) return {}
|
||||
const out: Record<string, unknown> = {}
|
||||
for (const [key, item] of Object.entries(value)) {
|
||||
setSchemaProperty(out, key, schemaValue(item, seen))
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
function schemaArray(value: unknown, seen: WeakSet<object>): 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<string, unknown>,
|
||||
seen: WeakSet<object>,
|
||||
legacy: boolean,
|
||||
forcedType?: string,
|
||||
): JsonSchemaValue {
|
||||
const out: Record<string, unknown> = {}
|
||||
|
||||
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<string, unknown>): Record<string, unknown> {
|
||||
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<string, unknown>): unknown[] {
|
||||
if (Array.isArray(schema.variants)) return schema.variants
|
||||
if (Array.isArray(schema.anyOf)) return schema.anyOf
|
||||
return []
|
||||
}
|
||||
|
||||
function convertLegacySchema(
|
||||
source: Record<string, unknown>,
|
||||
kind: string,
|
||||
seen: WeakSet<object>,
|
||||
): 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<string, unknown> = {}
|
||||
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<string, unknown>, seen: WeakSet<object>): 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<object>())
|
||||
}
|
||||
+120
@@ -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<string, unknown> {
|
||||
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<T extends CommandCodeMessageLike>(
|
||||
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 } }
|
||||
}
|
||||
@@ -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<string, unknown> {
|
||||
return typeof value === "object" && value !== null && !Array.isArray(value)
|
||||
}
|
||||
|
||||
async function isUpgradeRequired(response: Response): Promise<boolean> {
|
||||
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<void> {
|
||||
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
|
||||
},
|
||||
}
|
||||
}
|
||||
+202
@@ -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<string, unknown>
|
||||
}
|
||||
|
||||
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<Record<string, string | null>>
|
||||
thinking?: {
|
||||
mode?: "effort"
|
||||
effortMap?: Partial<Record<string, string>>
|
||||
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<string, string>
|
||||
}
|
||||
|
||||
export interface StreamOptions {
|
||||
apiKey?: string
|
||||
signal?: AbortSignal
|
||||
headers?: Record<string, string>
|
||||
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<unknown>
|
||||
onResponse?: (response: ProviderResponseInfo, model: ModelLike) => void | Promise<void>
|
||||
/**
|
||||
* 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<AssistantMessageEvent> {
|
||||
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<void>
|
||||
}
|
||||
@@ -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<IteratorResult<AssistantMessageEvent>> {
|
||||
while (index >= events.length && !ended) {
|
||||
await new Promise<void>((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<AssistantMessageEvent[]> {
|
||||
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<AssistantMessageEvent[]>((_, reject) => {
|
||||
setTimeout(
|
||||
() => reject(new Error(`Timed out collecting stream events after ${timeoutMs}ms`)),
|
||||
timeoutMs,
|
||||
)
|
||||
}),
|
||||
])
|
||||
}
|
||||
|
||||
export function makeModel(overrides: Partial<ModelLike> = {}): 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> = {}): ContextLike {
|
||||
return {
|
||||
systemPrompt: "You are a test assistant.",
|
||||
messages: [{ role: "user", content: "hello" }],
|
||||
tools: [],
|
||||
...overrides,
|
||||
}
|
||||
}
|
||||
|
||||
export interface TestDepsResult {
|
||||
streamCommandCode: ReturnType<typeof createStreamCommandCode>
|
||||
calculatedUsages: Usage[]
|
||||
}
|
||||
|
||||
export function createTestDeps(overrides: Partial<CoreDependencies> = {}): 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<string, string>
|
||||
}
|
||||
|
||||
export type ResponsePlan = SuccessPlan | ErrorPlan
|
||||
|
||||
function headersToRecord(headers: IncomingHttpHeaders): Record<string, string> {
|
||||
const out: Record<string, string> = {}
|
||||
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<void>
|
||||
lastRequestBody(): unknown
|
||||
lastRequestHeaders(): Record<string, string>
|
||||
requestCount(): number
|
||||
responseClosedBeforeEnd(): boolean
|
||||
}
|
||||
|
||||
export async function startMockCommandCodeServer(): Promise<MockCommandCodeServer> {
|
||||
let planQueue: ResponsePlan[] = [{ type: "success", events: [] }]
|
||||
let lastBody: unknown
|
||||
let lastHeaders: Record<string, string> = {}
|
||||
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<string, string> = { "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<void>((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<void>((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
|
||||
}
|
||||
@@ -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")
|
||||
})
|
||||
})
|
||||
@@ -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<CostRates & { inputTokensAbove: number }>
|
||||
}
|
||||
|
||||
const COST_FIXTURES: Record<string, CostTable> = {
|
||||
"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)
|
||||
})
|
||||
})
|
||||
@@ -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)
|
||||
})
|
||||
})
|
||||
@@ -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<string, unknown> = {
|
||||
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")
|
||||
})
|
||||
})
|
||||
@@ -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<void>((_, 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"])
|
||||
})
|
||||
})
|
||||
@@ -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<void>(() => {}),
|
||||
}),
|
||||
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<unknown>(() => {}),
|
||||
}),
|
||||
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",
|
||||
)
|
||||
})
|
||||
})
|
||||
@@ -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)
|
||||
})
|
||||
})
|
||||
Reference in New Issue
Block a user