pi-commandcode-provider/src/core.ts

/**
 * 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. Run /login commandcode, 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 ?? 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);
              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;
  };
}