pi-commandcode-provider 0.1.0

package/src/core.ts

@@ -0,0 +1,458 @@
+/**
+ * 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 and select Command Code, 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
+  }
+}

package/src/oauth.ts

@@ -0,0 +1,167 @@
+/**
+ * Command Code OAuth provider for pi's /login flow.
+ *
+ * Implements a browser-assisted API key retrieval flow:
+ * 1. Starts a local HTTP server on a Command Code CLI-compatible port
+ * 2. Opens the Command Code Studio auth page in the browser
+ * 3. The user authenticates on the Command Code website
+ * 4. The website POSTs the API key back to the local server
+ * 5. If browser transfer fails, the user can paste the API key manually
+ * 6. The API key is stored in pi's auth.json as OAuth credentials
+ *
+ * Since Command Code API keys don't expire, we store them as
+ * OAuth credentials with a far-future expiry.
+ */
+
+import { randomBytes } from "node:crypto"
+import { startAuthServer } from "./auth-server.ts"
+
+const STUDIO_BASE_URL = "https://commandcode.ai"
+const TEN_YEARS_MS = 10 * 365 * 24 * 60 * 60 * 1000 // API keys don't expire
+const DEFAULT_AUTH_TIMEOUT_MS = 15_000
+
+export interface OAuthLoginCallbacks {
+  onAuth(params: { url: string }): void
+  onPrompt(params: { message: string }): Promise<string>
+}
+
+export interface OAuthCredentials {
+  refresh: string
+  access: string
+  expires: number
+}
+
+class AuthTimeoutError extends Error {
+  constructor() {
+    super("Browser authentication timed out")
+    this.name = "AuthTimeoutError"
+  }
+}
+
+function generateStateToken(): string {
+  return randomBytes(32).toString("base64url")
+}
+
+function getAuthTimeoutMs(): number {
+  const raw = process.env.COMMANDCODE_AUTH_TIMEOUT_MS
+  if (!raw) return DEFAULT_AUTH_TIMEOUT_MS
+
+  const parsed = Number(raw)
+  return Number.isFinite(parsed) && parsed > 0 ? parsed : DEFAULT_AUTH_TIMEOUT_MS
+}
+
+function withTimeout<T>(promise: Promise<T>, timeoutMs: number): Promise<T> {
+  return new Promise((resolve, reject) => {
+    const timer = setTimeout(() => reject(new AuthTimeoutError()), timeoutMs)
+
+    promise.then(
+      (value) => {
+        clearTimeout(timer)
+        resolve(value)
+      },
+      (error) => {
+        clearTimeout(timer)
+        reject(error)
+      },
+    )
+  })
+}
+
+function credentialsFromApiKey(apiKey: string): OAuthCredentials {
+  return {
+    refresh: apiKey,
+    access: apiKey,
+    expires: Date.now() + TEN_YEARS_MS,
+  }
+}
+
+/**
+ * Remove common terminal paste wrappers/control chars and surrounding whitespace.
+ */
+export function sanitizeApiKey(input: string): string {
+  const esc = String.fromCharCode(27)
+  return Array.from(
+    input
+      .replaceAll(`${esc}[200~`, "")
+      .replaceAll(`${esc}[201~`, "")
+      .replaceAll("[200~", "")
+      .replaceAll("[201~", ""),
+  )
+    .filter((char) => {
+      const code = char.charCodeAt(0)
+      return code > 31 && code !== 127
+    })
+    .join("")
+    .trim()
+}
+
+async function promptForApiKey(callbacks: OAuthLoginCallbacks, message: string) {
+  const apiKey = sanitizeApiKey(await callbacks.onPrompt({ message }))
+  if (!apiKey) throw new Error("No Command Code API key provided")
+  return credentialsFromApiKey(apiKey)
+}
+
+/**
+ * Starts the browser-based login flow for Command Code.
+ *
+ * Returns OAuth credentials where access == refresh == the user's API key.
+ * The keys don't expire, so we set a far-future expiry.
+ */
+export async function login(callbacks: OAuthLoginCallbacks): Promise<OAuthCredentials> {
+  let authServer
+  try {
+    authServer = await startAuthServer()
+  } catch {
+    return promptForApiKey(
+      callbacks,
+      "Could not start browser auth. Paste your Command Code API key:",
+    )
+  }
+
+  const stateToken = generateStateToken()
+  const callbackUrl = `http://localhost:${authServer.port}/callback`
+  const authUrl = `${STUDIO_BASE_URL}/studio/auth/cli?callback=${encodeURIComponent(callbackUrl)}&state=${encodeURIComponent(stateToken)}`
+
+  // Tell pi to open the browser.
+  callbacks.onAuth({ url: authUrl })
+
+  // Wait for the Command Code Studio to POST the API key back. If the browser
+  // cannot reach localhost (Command Code shows "Copy your API key"), fall back
+  // to pi's prompt so the user can paste the key from the browser.
+  let callback: { apiKey: string; state: string }
+  try {
+    callback = await withTimeout(authServer.waitForCallback, getAuthTimeoutMs())
+  } catch (error) {
+    authServer.server.close()
+    if (error instanceof AuthTimeoutError) {
+      return promptForApiKey(
+        callbacks,
+        "Automatic transfer failed or timed out. Paste your Command Code API key:",
+      )
+    }
+    throw error
+  }
+
+  // Validate state token to prevent CSRF.
+  if (callback.state !== stateToken) {
+    authServer.server.close()
+    throw new Error("State token mismatch. Authentication may have been tampered with.")
+  }
+
+  return credentialsFromApiKey(callback.apiKey)
+}
+
+/**
+ * Command Code API keys don't expire, so "refresh" is a no-op.
+ * Returns the same credentials with an updated far-future expiry.
+ */
+export async function refreshToken(credentials: OAuthCredentials): Promise<OAuthCredentials> {
+  return credentialsFromApiKey(credentials.refresh)
+}
+
+/**
+ * Returns the access token (API key) from OAuth credentials.
+ */
+export function getApiKey(credentials: OAuthCredentials): string {
+  return credentials.access
+}