feat: implement streamCommandCode with SSE parsing

index.ts

@@ -7,6 +7,15 @@ 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";
@@ -40,37 +49,25 @@ const MODELS = [
 ];
 
 // ---------------------------------------------------------------------------
-// Typebox → JSON Schema conversion
+// Helpers
 // ---------------------------------------------------------------------------
 
 function toJsonSchema(schema: any): any {
   if (!schema) return {};
   const s = schema as Record<string, any>;
   const kind = s.kind ?? s.type;
-
+  if (s.enum) return { type: typeof s.enum[0], enum: s.enum };
-  if (s.enum) {
-    return { type: typeof s.enum[0], enum: s.enum };
-  }
-
   switch (kind) {
-    case "string":
+    case "string": case "String": return { type: "string" };
-    case "String":
+    case "number": case "Number": return { type: "number" };
-      return { type: "string" };
+    case "boolean": case "Boolean": return { type: "boolean" };
-    case "number":
+    case "object": case "Object": {
-    case "Number":
-      return { type: "number" };
-    case "boolean":
-    case "Boolean":
-      return { type: "boolean" };
-    case "object":
-    case "Object": {
       const props: Record<string, any> = {};
       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))
+          if (!(v as any).optional && !s.optional?.includes?.(k)) inferredRequired.push(k);
-            inferredRequired.push(k);
         }
       }
       const required = Array.isArray(s.required) ? s.required : inferredRequired;
@@ -79,82 +76,43 @@ function toJsonSchema(schema: any): any {
       if (required.length) out.required = required;
       return out;
     }
-    case "array":
+    case "array": case "Array": return { type: "array", items: toJsonSchema(s.items ?? s.element) };
-    case "Array":
+    case "union": case "Union": {
-      return { type: "array", items: toJsonSchema(s.items ?? s.element) };
-    case "union":
-    case "Union": {
       const variants = s.variants ?? s.anyOf ?? [];
-      for (const v of variants) {
+      for (const v of variants) { const sch = toJsonSchema(v); if (sch && Object.keys(sch).length) return sch; }
-        const schema = toJsonSchema(v);
-        if (schema && Object.keys(schema).length) return schema;
-      }
       return {};
     }
-    case "optional":
+    case "optional": case "Optional": return toJsonSchema(s.wrapped ?? s.inner);
-    case "Optional":
+    default: return {};
-      return toJsonSchema(s.wrapped ?? s.inner);
-    default:
-      return {};
   }
 }
 
 function toolsToJson(tools: any[]): any[] {
   if (!tools) return [];
-  return tools.map((t) => {
+  return tools.map((t) => ({
-    const schema = t.parameters ? toJsonSchema(t.parameters) : {};
-    return {
     type: "function",
     name: t.name,
     description: t.description,
-      input_schema: schema,
+    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({
+      out.push({ role: "user", content: typeof m.content === "string" ? m.content : m.content });
-        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") {
+        if (c.type === "text") parts.push({ type: "text", text: c.text });
-          parts.push({ type: "text", text: c.text });
+        else if (c.type === "thinking") parts.push({ type: "reasoning", text: c.thinking });
-        } else if (c.type === "thinking") {
+        else if (c.type === "toolCall") parts.push({ type: "tool-call", toolCallId: c.id, toolName: c.name, input: c.arguments });
-          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 ?? [])
+      const text = (m.content ?? []).filter((c: any) => c.type === "text").map((c: any) => c.text ?? "").join("\n");
-        .filter((c: any) => c.type === "text")
+      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 } }] });
-        .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;
@@ -168,6 +126,205 @@ 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<any>,
+  context: Context,
+  options?: SimpleStreamOptions,
+): AssistantMessageEventStream {
+  const stream = createAssistantMessageEventStream();
+
+  (async () => {
+    const apiKey = options?.apiKey ?? process.env.COMMANDCODE_API_KEY;
+    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.",
+        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();
+
+    try {
+      stream.push({ type: "start", partial: output });
+
+      const response = await 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(),
+        },
+        body: JSON.stringify({
+          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,
+          },
+        }),
+        signal: controller.signal,
+      });
+
+      if (!response.ok) {
+        const errBody = await response.text().catch(() => "");
+        throw new Error(`Command Code API error ${response.status}: ${errBody.slice(0, 500)}`);
+      }
+
+      const 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 (;;) {
+        const { done, value } = await reader.read();
+        if (done) break;
+        buffer += decoder.decode(value, { stream: true });
+        const lines = buffer.split("\n");
+        buffer = lines.pop() ?? "";
+
+        for (const line of lines) {
+          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;
+        }
+      }
+
+      if (textBlock) {
+        stream.push({ type: "text_end", contentIndex: currentTextIdx, content: textBlock.text, partial: output });
+      }
+
+      stream.push({ type: "done", reason: output.stopReason as "stop" | "length" | "toolUse", message: output });
+      stream.end();
+    } catch (error: any) {
+      output.stopReason = "error";
+      output.errorMessage = error?.message ?? String(error);
+      stream.push({ type: "error", reason: "error", error: output });
+      stream.end();
+    }
+  })();
+
+  return stream;
+}
+
 // ---------------------------------------------------------------------------
 // Extension entry point
 // ---------------------------------------------------------------------------
@@ -176,7 +333,10 @@ export default function (pi: ExtensionAPI) {
   pi.registerProvider("commandcode", {
     name: "Command Code",
     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,
+    streamSimple: streamCommandCode,
     headers: {
       "x-command-code-version": "0.24.1",
       "x-cli-environment": "production",