diff --git a/index.ts b/index.ts index 69f46a3..31f2484 100644 --- a/index.ts +++ b/index.ts @@ -9,22 +9,15 @@ * Models: deepseek-v4-pro, deepseek-v4-flash, claude-sonnet-4-6, claude-opus-4-7, etc. */ -import { existsSync, readFileSync } from "node:fs"; -import { homedir } from "node:os"; -import { join } from "node:path"; - import { calculateCost, - type AssistantMessage, - type AssistantMessageEventStream, - type Context, createAssistantMessageEventStream, - type Model, - type SimpleStreamOptions, } from "@mariozechner/pi-ai"; import type { ExtensionAPI } from "@mariozechner/pi-coding-agent"; -const API_BASE = "https://api.commandcode.ai"; +import { createStreamCommandCode, DEFAULT_API_BASE } from "./src/core.ts"; + +const API_BASE = process.env.COMMANDCODE_API_BASE ?? DEFAULT_API_BASE; // --------------------------------------------------------------------------- // Model definitions @@ -54,366 +47,11 @@ const MODELS = [ { id: "Qwen/Qwen3.6-Plus", name: "Qwen 3.6 Plus (CC)", reasoning: true, contextWindow: 1_000_000, maxTokens: 131_072 }, ]; -// --------------------------------------------------------------------------- -// Helpers -// --------------------------------------------------------------------------- - -function getApiKey(): string | undefined { - const env = process.env.COMMANDCODE_API_KEY; - if (env) return env; - - const authPaths = [ - join(homedir(), ".commandcode", "auth.json"), - join(homedir(), ".pi", "agent", "auth.json"), - ]; - - for (const authPath of authPaths) { - try { - if (!existsSync(authPath)) continue; - const auth = JSON.parse(readFileSync(authPath, "utf-8")); - if (typeof auth.apiKey === "string" && auth.apiKey) return auth.apiKey; - if (typeof auth.commandcode === "string" && auth.commandcode) return auth.commandcode; - } catch { - // ignore malformed/missing auth files - } - } - - return undefined; -} - -function toJsonSchema(schema: any): any { - if (!schema) return {}; - const s = schema as Record; - const kind = s.kind ?? s.type; - if (s.enum) return { type: typeof s.enum[0], enum: s.enum }; - switch (kind) { - case "string": case "String": return { type: "string" }; - case "number": case "Number": return { type: "number" }; - case "boolean": case "Boolean": return { type: "boolean" }; - case "object": case "Object": { - const props: Record = {}; - const inferredRequired: string[] = []; - if (s.properties) { - for (const [k, v] of Object.entries(s.properties)) { - props[k] = toJsonSchema(v); - if (!(v as any).optional && !s.optional?.includes?.(k)) inferredRequired.push(k); - } - } - const required = Array.isArray(s.required) ? s.required : inferredRequired; - const out: any = { type: "object" }; - if (Object.keys(props).length) out.properties = props; - if (required.length) out.required = required; - return out; - } - case "array": case "Array": return { type: "array", items: toJsonSchema(s.items ?? s.element) }; - case "union": case "Union": { - const variants = s.variants ?? s.anyOf ?? []; - for (const v of variants) { const sch = toJsonSchema(v); if (sch && Object.keys(sch).length) return sch; } - return {}; - } - case "optional": case "Optional": return toJsonSchema(s.wrapped ?? s.inner); - default: return {}; - } -} - -function toolsToJson(tools: any[]): any[] { - if (!tools) return []; - return tools.map((t) => ({ - type: "function", - name: t.name, - description: t.description, - input_schema: t.parameters ? toJsonSchema(t.parameters) : {}, - })); -} - -function messagesToCC(msgs: any[]): any[] { - const out: any[] = []; - for (const m of msgs) { - if (m.role === "user") { - out.push({ role: "user", content: typeof m.content === "string" ? m.content : m.content }); - } else if (m.role === "assistant") { - const parts: any[] = []; - for (const c of m.content) { - if (c.type === "text") parts.push({ type: "text", text: c.text }); - else if (c.type === "thinking") parts.push({ type: "reasoning", text: c.thinking }); - else if (c.type === "toolCall") parts.push({ type: "tool-call", toolCallId: c.id, toolName: c.name, input: c.arguments }); - } - out.push({ role: "assistant", content: parts }); - } else if (m.role === "toolResult") { - const text = (m.content ?? []).filter((c: any) => c.type === "text").map((c: any) => c.text ?? "").join("\n"); - out.push({ role: "tool", content: [{ type: "tool-result", toolCallId: m.toolCallId, toolName: m.toolName, output: m.isError ? { type: "error-text", value: text } : { type: "text", value: text } }] }); - } - } - return out; -} - -function getEnvironmentInfo(): string { - return `${process.platform}-${process.arch}, Node.js ${process.version}`; -} - -function uuid(): string { - return crypto.randomUUID(); -} - -function parseStreamEventLine(line: string): any | 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 { return JSON.parse(trimmed); } catch { return undefined; } -} - -function mapFinishReason(reason: unknown): "stop" | "length" | "toolUse" { - if (reason === "tool-calls") return "toolUse"; - if (reason === "length" || reason === "max_tokens" || reason === "max-tokens" || reason === "max_output_tokens") return "length"; - return "stop"; -} - -// --------------------------------------------------------------------------- -// Stream implementation -// --------------------------------------------------------------------------- - -function streamCommandCode( - model: Model, - context: Context, - options?: SimpleStreamOptions, -): AssistantMessageEventStream { - const stream = createAssistantMessageEventStream(); - - (async () => { - const apiKey = options?.apiKey ?? getApiKey(); - if (!apiKey) { - const msg: AssistantMessage = { - 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: "No Command Code API key. Set COMMANDCODE_API_KEY env var or configure ~/.commandcode/auth.json or ~/.pi/agent/auth.json.", - timestamp: Date.now(), - }; - stream.push({ type: "error", reason: "error", error: msg }); - stream.end(); - return; - } - - const output: AssistantMessage = { - role: "assistant", content: [], api: model.api, provider: model.provider, model: model.id, - usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 } }, - stopReason: "stop", timestamp: Date.now(), - }; - - const controller = new AbortController(); - let reader: ReadableStreamDefaultReader | undefined; - - const abortUpstream = () => { - if (!controller.signal.aborted) controller.abort(); - try { reader?.cancel().catch(() => undefined); } catch { /* best-effort */ } - }; - - if (options?.signal?.aborted) { - abortUpstream(); - } else { - options?.signal?.addEventListener("abort", abortUpstream, { once: true }); - } - - // Helper: race a promise against the abort signal. - const raceAbort = (promise: Promise): Promise => { - if (controller.signal.aborted) { - return Promise.reject(new DOMException("The operation was aborted", "AbortError")); - } - return new Promise((resolve, reject) => { - const onAbort = () => reject(new DOMException("The operation was aborted", "AbortError")); - controller.signal.addEventListener("abort", onAbort, { once: true }); - promise.then( - (v) => { controller.signal.removeEventListener("abort", onAbort); resolve(v); }, - (e) => { controller.signal.removeEventListener("abort", onAbort); reject(e); }, - ); - }); - }; - - try { - stream.push({ type: "start", partial: output }); - - let body: unknown = { - config: { - workingDir: process.cwd(), - date: new Date().toISOString().split("T")[0], - environment: getEnvironmentInfo(), - structure: [], - isGitRepo: false, - currentBranch: "", - mainBranch: "", - gitStatus: "", - recentCommits: [], - }, - memory: "", taste: "", skills: null, - permissionMode: "standard" as const, - params: { - model: model.id, - messages: messagesToCC(context.messages), - tools: toolsToJson(context.tools), - system: context.systemPrompt ?? "", - max_tokens: Math.min(options?.maxTokens ?? model.maxTokens, 200_000), - stream: true, - }, - }; - - const nextBody = await raceAbort(Promise.resolve(options?.onPayload?.(body, model))); - if (nextBody !== undefined) body = nextBody; - - const response = await raceAbort(fetch(`${API_BASE}/alpha/generate`, { - method: "POST", - headers: { - "Content-Type": "application/json", - Authorization: `Bearer ${apiKey}`, - "x-command-code-version": "0.24.1", - "x-cli-environment": "production", - "x-project-slug": "pi-cc", - "x-taste-learning": "false", - "x-co-flag": "false", - "x-session-id": uuid(), - ...options?.headers, - }, - body: JSON.stringify(body), - signal: controller.signal, - })); - - await raceAbort(Promise.resolve(options?.onResponse?.({ - status: response.status, - headers: Object.fromEntries(response.headers.entries()), - }, model))); - - if (!response.ok) { - const errBody = await raceAbort(response.text().catch(() => "")); - throw new Error(`Command Code API error ${response.status}: ${errBody.slice(0, 500)}`); - } - - reader = response.body?.getReader(); - if (!reader) throw new Error("No response body"); - - const decoder = new TextDecoder(); - let buffer = ""; - let currentTextIdx = -1; - let textBlock: any = null; - let reasoningActive = false; - let thinkingBlock: string[] = []; - let finished = false; - - readLoop: for (;;) { - if (controller.signal.aborted) throw new DOMException("Aborted", "AbortError"); - const { done, value } = await raceAbort(reader.read()); - if (done) break; - if (controller.signal.aborted) throw new DOMException("Aborted", "AbortError"); - buffer += decoder.decode(value, { stream: true }); - const lines = buffer.split("\n"); - buffer = lines.pop() ?? ""; - - for (const line of lines) { - if (controller.signal.aborted) throw new DOMException("Aborted", "AbortError"); - const event = parseStreamEventLine(line); - if (!event) continue; - - switch (event.type) { - case "text-delta": { - if (!textBlock) { - textBlock = { type: "text", text: "" }; - output.content.push(textBlock); - currentTextIdx = output.content.length - 1; - stream.push({ type: "text_start", contentIndex: currentTextIdx, partial: output }); - } - textBlock.text += event.text ?? ""; - stream.push({ type: "text_delta", contentIndex: currentTextIdx, delta: event.text ?? "", partial: output }); - break; - } - case "reasoning-delta": { - if (!reasoningActive) reasoningActive = true; - thinkingBlock.push(event.text ?? ""); - break; - } - case "reasoning-end": { - if (thinkingBlock.length > 0) { - const thinkingText = thinkingBlock.join(""); - thinkingBlock = []; - output.content.push({ type: "thinking", thinking: thinkingText }); - const idx = output.content.length - 1; - stream.push({ type: "thinking_start", contentIndex: idx, partial: output }); - stream.push({ type: "thinking_delta", contentIndex: idx, delta: thinkingText, partial: output }); - stream.push({ type: "thinking_end", contentIndex: idx, content: thinkingText, partial: output }); - } - reasoningActive = false; - break; - } - case "tool-call": { - if (textBlock) { - stream.push({ type: "text_end", contentIndex: currentTextIdx, content: textBlock.text, partial: output }); - textBlock = null; - currentTextIdx = -1; - } - output.content.push({ type: "toolCall", id: event.toolCallId, name: event.toolName, arguments: event.input ?? event.args ?? {} }); - const idx = output.content.length - 1; - stream.push({ type: "toolcall_start", contentIndex: idx, partial: output }); - stream.push({ type: "toolcall_end", contentIndex: idx, toolCall: { type: "toolCall", id: event.toolCallId, name: event.toolName, arguments: event.input ?? event.args ?? {} }, partial: output }); - break; - } - case "finish": { - const usage = event.totalUsage; - if (usage) { - output.usage.input = usage.inputTokens ?? 0; - output.usage.output = usage.outputTokens ?? 0; - output.usage.cacheRead = usage.inputTokenDetails?.cacheReadTokens ?? 0; - output.usage.cacheWrite = usage.inputTokenDetails?.cacheWriteTokens ?? 0; - output.usage.totalTokens = output.usage.input + output.usage.output + output.usage.cacheRead + output.usage.cacheWrite; - calculateCost(model, output.usage); - } - output.stopReason = mapFinishReason(event.finishReason); - finished = true; - break; - } - case "error": { - const msg = event.error?.message ?? event.error ?? "Stream error"; - output.stopReason = "error"; - output.errorMessage = typeof msg === "string" ? msg : String(msg); - throw new Error(output.errorMessage); - } - } - if (finished) break readLoop; - } - } - - // End any lingering text block - if (textBlock) { - stream.push({ type: "text_end", contentIndex: currentTextIdx, content: textBlock.text, partial: output }); - } - - // Emit remaining thinking (may arrive after finish without reasoning-end) - if (thinkingBlock.length > 0) { - const thinkingText = thinkingBlock.join(""); - output.content.push({ type: "thinking", thinking: thinkingText }); - const idx = output.content.length - 1; - stream.push({ type: "thinking_start", contentIndex: idx, partial: output }); - stream.push({ type: "thinking_delta", contentIndex: idx, delta: thinkingText, partial: output }); - stream.push({ type: "thinking_end", contentIndex: idx, content: thinkingText, partial: output }); - } - - stream.push({ type: "done", reason: output.stopReason as "stop" | "length" | "toolUse", message: output }); - stream.end(); - } catch (error: any) { - if (controller.signal.aborted) { - output.stopReason = "aborted"; - output.errorMessage = "Request aborted"; - } else { - output.stopReason = "error"; - output.errorMessage = error?.message ?? String(error); - } - stream.push({ type: "error", reason: output.stopReason, error: output }); - stream.end(); - } finally { - options?.signal?.removeEventListener("abort", abortUpstream); - try { await reader?.cancel(); } catch { /* best-effort */ } - try { reader?.releaseLock(); } catch { /* may already be released */ } - } - })(); - - return stream; -} +const streamCommandCode = createStreamCommandCode({ + createStream: createAssistantMessageEventStream, + calculateCost, + apiBase: API_BASE, +}); // --------------------------------------------------------------------------- // Extension entry point @@ -425,20 +63,20 @@ export default function (pi: ExtensionAPI) { baseUrl: API_BASE, apiKey: "!python3 -c 'import json,pathlib; key=\"\"; paths=[pathlib.Path.home()/\".commandcode/auth.json\", pathlib.Path.home()/\".pi/agent/auth.json\"];\nfor p in paths:\n try:\n data=json.loads(p.read_text()); key=data.get(\"apiKey\") or data.get(\"commandcode\") or key\n if key: break\n except Exception: pass\nprint(key)'", authHeader: true, - api: "commandcode-custom" as any, + api: "commandcode-custom", streamSimple: streamCommandCode, headers: { "x-command-code-version": "0.24.1", "x-cli-environment": "production", }, - models: MODELS.map((m) => ({ - id: m.id, - name: m.name, - reasoning: m.reasoning, - input: ["text"] as ("text" | "image")[], + models: MODELS.map((model) => ({ + id: model.id, + name: model.name, + reasoning: model.reasoning, + input: ["text"], cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, - contextWindow: m.contextWindow, - maxTokens: m.maxTokens, + contextWindow: model.contextWindow, + maxTokens: model.maxTokens, })), }); } diff --git a/src/converters.ts b/src/converters.ts new file mode 100644 index 0000000..16d8e57 --- /dev/null +++ b/src/converters.ts @@ -0,0 +1,228 @@ +import { existsSync, readFileSync } from "node:fs"; +import { homedir } from "node:os"; +import { join } from "node:path"; + +import type { MessageLike, StopReason, ToolLike } from "./types.ts"; + +export function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +export function stringValue(value: unknown): string | undefined { + return typeof value === "string" ? value : undefined; +} + +function booleanValue(value: unknown): boolean | undefined { + return typeof value === "boolean" ? value : undefined; +} + +export function recordArray(value: unknown): readonly Record[] { + if (!Array.isArray(value)) return []; + return value.filter(isRecord); +} + +export function recordOrEmpty(value: unknown): Record { + return isRecord(value) ? value : {}; +} + +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, ".pi", "agent", "auth.json"), + ]; +} + +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; + const apiKey = stringValue(parsed.apiKey); + if (apiKey) return apiKey; + const commandcode = stringValue(parsed.commandcode); + if (commandcode) return commandcode; + } 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 toJsonSchema(schema: unknown): unknown { + if (!isRecord(schema)) return {}; + + const kind = stringValue(schema.kind) ?? stringValue(schema.type); + const enumValues = Array.isArray(schema.enum) ? schema.enum : undefined; + if (enumValues) { + return { type: typeof enumValues[0], enum: enumValues }; + } + + switch (kind) { + case "string": + case "String": + return { type: "string" }; + case "number": + case "Number": + return { type: "number" }; + case "boolean": + case "Boolean": + return { type: "boolean" }; + case "object": + case "Object": { + const properties: Record = {}; + const inferredRequired: string[] = []; + const sourceProperties = isRecord(schema.properties) ? schema.properties : undefined; + const optional = Array.isArray(schema.optional) + ? schema.optional.filter((item): item is string => typeof item === "string") + : []; + + if (sourceProperties) { + for (const [key, value] of Object.entries(sourceProperties)) { + properties[key] = toJsonSchema(value); + const valueRecord = isRecord(value) ? value : undefined; + if (booleanValue(valueRecord?.optional) !== true && !optional.includes(key)) { + inferredRequired.push(key); + } + } + } + + const explicitRequired = Array.isArray(schema.required) + ? schema.required.filter((item): item is string => typeof item === "string") + : undefined; + const required = explicitRequired ?? inferredRequired; + const out: Record = { type: "object" }; + if (Object.keys(properties).length > 0) out.properties = properties; + if (required.length > 0) out.required = required; + return out; + } + case "array": + case "Array": + return { type: "array", items: toJsonSchema(schema.items ?? schema.element) }; + case "union": + case "Union": { + const variants = Array.isArray(schema.variants) + ? schema.variants + : Array.isArray(schema.anyOf) + ? schema.anyOf + : []; + for (const variant of variants) { + const converted = toJsonSchema(variant); + if (isRecord(converted) && Object.keys(converted).length > 0) return converted; + } + return {}; + } + case "optional": + case "Optional": + return toJsonSchema(schema.wrapped ?? schema.inner); + default: + return {}; + } +} + +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) : {}, + })); +} + +export function messagesToCC(messages?: readonly MessageLike[]): unknown[] { + const out: unknown[] = []; + for (const message of messages ?? []) { + if (message.role === "user") { + out.push({ + role: "user", + content: typeof message.content === "string" ? message.content : message.content, + }); + } 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 === "thinking") { + parts.push({ type: "reasoning", text: stringValue(content.thinking) ?? "" }); + } else if (content.type === "toolCall") { + parts.push({ + type: "tool-call", + toolCallId: stringValue(content.id) ?? "", + toolName: stringValue(content.name) ?? "", + input: recordOrEmpty(content.arguments), + }); + } + } + out.push({ role: "assistant", content: parts }); + } else if (message.role === "toolResult") { + 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) }, + }, + ], + }); + } + } + 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"; +} diff --git a/src/core.ts b/src/core.ts new file mode 100644 index 0000000..f70a445 --- /dev/null +++ b/src/core.ts @@ -0,0 +1,398 @@ +/** + * 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 { + getApiKey, + getEnvironmentInfo, + isRecord, + mapFinishReason, + messagesToCC, + numberValue, + parseStreamEventLine, + recordOrEmpty, + stringValue, + toolsToJson, +} 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 "./types.ts"; + +export const DEFAULT_API_BASE = "https://api.commandcode.ai"; + +function defaultUsage(): Usage { + return { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }; +} + +function commandCodeUsage(event: Record): Record | undefined { + return isRecord(event.totalUsage) ? event.totalUsage : undefined; +} + +function commandCodeInputTokenDetails(usage: Record): Record | undefined { + return isRecord(usage.inputTokenDetails) ? usage.inputTokenDetails : undefined; +} + +function headersToRecord(headers: Headers): Record { + const out: Record = {}; + headers.forEach((value, key) => { + out[key] = value; + }); + return out; +} + +function abortError(message = "The operation was aborted"): DOMException { + return new DOMException(message, "AbortError"); +} + +function successStopReason(reason: TerminalReason): StopReason { + if (reason === "length" || reason === "toolUse") return reason; + return "stop"; +} + +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()); + + function raceAbort(promise: Promise, signal: AbortSignal): Promise { + if (signal.aborted) return Promise.reject(abortError()); + + return new Promise((resolve, reject) => { + const onAbort = () => reject(abortError()); + signal.addEventListener("abort", onAbort, { once: true }); + promise.then( + (value) => { + signal.removeEventListener("abort", onAbort); + resolve(value); + }, + (error: unknown) => { + signal.removeEventListener("abort", onAbort); + reject(error); + }, + ); + }); + } + + return function streamCommandCode( + model: ModelLike, + context: ContextLike, + options?: StreamOptions, + ): AssistantMessageEventStreamLike { + const stream = deps.createStream(); + + async function run() { + const apiKey = options?.apiKey ?? 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. Set COMMANDCODE_API_KEY env var or configure ~/.commandcode/auth.json or ~/.pi/agent/auth.json.", + timestamp: now(), + }; + stream.push({ type: "error", reason: "error", error: msg }); + stream.end(); + return; + } + + const output: AssistantMessageLike = { + role: "assistant", + content: [], + api: model.api, + provider: model.provider, + model: model.id, + usage: defaultUsage(), + stopReason: "stop", + timestamp: now(), + }; + + const controller = new AbortController(); + let reader: ReadableStreamDefaultReader | undefined; + let textBlock: TextContent | undefined; + let currentTextIdx = -1; + let thinkingBlock: string[] = []; + 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 flushThinkingBlock = () => { + if (thinkingBlock.length === 0) return; + const thinkingText = thinkingBlock.join(""); + thinkingBlock = []; + output.content.push({ type: "thinking", thinking: thinkingText }); + const idx = output.content.length - 1; + stream.push({ type: "thinking_start", contentIndex: idx, partial: output }); + stream.push({ type: "thinking_delta", contentIndex: idx, delta: thinkingText, partial: output }); + stream.push({ type: "thinking_end", contentIndex: idx, content: thinkingText, partial: output }); + }; + + const handleEvent = (event: unknown) => { + if (!isRecord(event)) return; + + switch (event.type) { + case "text-delta": { + 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-delta": { + thinkingBlock.push(stringValue(event.text) ?? ""); + break; + } + + case "reasoning-end": { + flushThinkingBlock(); + break; + } + + case "tool-call": { + endTextBlock(); + const toolCall: ToolCallContent = { + type: "toolCall", + id: stringValue(event.toolCallId) ?? "", + name: stringValue(event.toolName) ?? "", + arguments: recordOrEmpty(event.input ?? event.args), + }; + 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); + output.usage.input = numberValue(usage.inputTokens) ?? 0; + output.usage.output = numberValue(usage.outputTokens) ?? 0; + output.usage.cacheRead = numberValue(details?.cacheReadTokens) ?? 0; + output.usage.cacheWrite = numberValue(details?.cacheWriteTokens) ?? 0; + 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 errorRecord = isRecord(event.error) ? event.error : undefined; + const message = stringValue(errorRecord?.message) ?? stringValue(event.error) ?? "Stream error"; + output.stopReason = "error"; + output.errorMessage = message; + throw new Error(message); + } + } + }; + + try { + stream.push({ type: "start", partial: output }); + + let body: unknown = { + config: { + workingDir: cwd(), + date: new Date(now()).toISOString().split("T")[0], + environment: getEnvironmentInfo(), + structure: [], + isGitRepo: false, + currentBranch: "", + mainBranch: "", + gitStatus: "", + recentCommits: [], + }, + memory: "", + taste: "", + skills: null, + permissionMode: "standard", + params: { + model: model.id, + messages: messagesToCC(context.messages), + tools: toolsToJson(context.tools), + system: context.systemPrompt ?? "", + max_tokens: Math.min(options?.maxTokens ?? model.maxTokens, 200_000), + stream: true, + }, + }; + + const nextBody = await raceAbort(Promise.resolve(options?.onPayload?.(body, model)), controller.signal); + if (nextBody !== undefined) body = nextBody; + + const response = await raceAbort( + fetchImpl(`${apiBase}/alpha/generate`, { + method: "POST", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${apiKey}`, + "x-command-code-version": "0.24.1", + "x-cli-environment": "production", + "x-project-slug": "pi-cc", + "x-taste-learning": "false", + "x-co-flag": "false", + "x-session-id": uuid(), + ...options?.headers, + }, + body: JSON.stringify(body), + signal: controller.signal, + }), + controller.signal, + ); + + await raceAbort( + Promise.resolve(options?.onResponse?.({ status: response.status, headers: headersToRecord(response.headers) }, model)), + controller.signal, + ); + + if (!response.ok) { + const errBody = await raceAbort(response.text().catch(() => ""), controller.signal); + throw new Error(`Command Code API error ${response.status}: ${errBody.slice(0, 500)}`); + } + + reader = response.body?.getReader(); + if (!reader) throw new Error("No response body"); + + const decoder = new TextDecoder(); + let buffer = ""; + + readLoop: for (;;) { + if (controller.signal.aborted) throw abortError("Aborted"); + const { done, value } = await raceAbort(reader.read(), controller.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; + } + } + + endTextBlock(); + flushThinkingBlock(); + + stream.push({ type: "done", reason: successStopReason(output.stopReason), message: output }); + stream.end(); + } catch (error: unknown) { + const reason: ErrorReason = controller.signal.aborted ? "aborted" : "error"; + output.stopReason = reason; + output.errorMessage = reason === "aborted" + ? "Request aborted" + : 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: error instanceof Error ? error.message : String(error), + timestamp: now(), + }; + stream.push({ type: "error", reason: "error", error: msg }); + stream.end(); + }); + + return stream; + }; +} diff --git a/src/types.ts b/src/types.ts new file mode 100644 index 0000000..4af022d --- /dev/null +++ b/src/types.ts @@ -0,0 +1,123 @@ +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; + totalTokens: number; + cost: UsageCost; +} + +export interface TextContent { + type: "text"; + text: string; +} + +export interface ThinkingContent { + type: "thinking"; + thinking: string; +} + +export interface ToolCallContent { + type: "toolCall"; + id: string; + name: string; + arguments: Record; +} + +export type AssistantContent = TextContent | ThinkingContent | ToolCallContent; + +export interface AssistantMessageLike { + role: "assistant"; + content: AssistantContent[]; + api: unknown; + provider: string; + model: string; + usage: Usage; + stopReason: TerminalReason; + errorMessage?: string; + timestamp: number; +} + +export interface ModelLike { + id: string; + api: unknown; + provider: string; + maxTokens: number; +} + +export interface MessageLike { + role: string; + content?: unknown; + toolCallId?: string; + toolName?: string; + isError?: boolean; +} + +export interface ToolLike { + name: string; + description?: string; + parameters?: unknown; +} + +export interface ContextLike { + systemPrompt?: string; + messages?: readonly MessageLike[]; + tools?: readonly ToolLike[]; +} + +export interface ProviderResponseInfo { + status: number; + headers: Record; +} + +export interface StreamOptions { + apiKey?: string; + signal?: AbortSignal; + headers?: Record; + maxTokens?: number; + onPayload?: (payload: unknown, model: ModelLike) => unknown | Promise; + onResponse?: (response: ProviderResponseInfo, model: ModelLike) => void | Promise; +} + +export type AssistantMessageEvent = + | { type: "start"; partial: AssistantMessageLike } + | { type: "text_start"; contentIndex: number; partial: AssistantMessageLike } + | { type: "text_delta"; contentIndex: number; delta: string; partial: AssistantMessageLike } + | { type: "text_end"; contentIndex: number; content: string; partial: AssistantMessageLike } + | { type: "thinking_start"; contentIndex: number; partial: AssistantMessageLike } + | { type: "thinking_delta"; contentIndex: number; delta: string; partial: AssistantMessageLike } + | { type: "thinking_end"; contentIndex: number; content: string; partial: AssistantMessageLike } + | { type: "toolcall_start"; contentIndex: number; partial: AssistantMessageLike } + | { type: "toolcall_end"; contentIndex: number; toolCall: ToolCallContent; partial: AssistantMessageLike } + | { type: "done"; reason: StopReason; message: AssistantMessageLike } + | { type: "error"; reason: ErrorReason; error: AssistantMessageLike }; + +export interface AssistantMessageEventStreamLike extends AsyncIterable { + push(event: AssistantMessageEvent): void; + end(): void; +} + +export interface CoreDependencies { + createStream: () => AssistantMessageEventStreamLike; + calculateCost: (model: ModelLike, usage: Usage) => void; + apiBase?: string; + fetchImpl?: typeof fetch; + authPaths?: readonly string[]; + env?: NodeJS.ProcessEnv; + cwd?: () => string; + now?: () => number; + uuid?: () => string; + homeDir?: () => string; +}