Extract Command Code stream core
This commit is contained in:
Patrick Wozniak
+398
@@ -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<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 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<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);
|
||||
},
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
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<Uint8Array> | 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;
|
||||
};
|
||||
}
|
||||
Reference in New Issue
Block a user