fix(stream): match Command Code CLI transport behavior

tests/test-oauth.ts

@@ -9,7 +9,7 @@ import assert from "node:assert/strict"
 import { describe, it } from "node:test"
 
 import { startAuthServer, type AuthCallback } from "../src/auth-server.ts"
-import { getApiKey, login, refreshToken, sanitizeApiKey } from "../src/oauth.ts"
+import { getApiKey, login, refreshToken, sanitizeApiKey, validateApiKey } from "../src/oauth.ts"
 
 /**
  * Helper: wait for an HTTP server to close, or resolve immediately if already closed.
@@ -24,9 +24,27 @@ function waitForClose(server: {
   })
 }
 
+async function withValidApiKeyFetch<T>(run: () => Promise<T>): Promise<T> {
+  const originalFetch = globalThis.fetch
+  globalThis.fetch = (input, init) => {
+    if (String(input).endsWith("/alpha/whoami")) {
+      return Promise.resolve(new Response(JSON.stringify({ user: {} }), { status: 200 }))
+    }
+    return originalFetch(input, init)
+  }
+  try {
+    return await run()
+  } finally {
+    globalThis.fetch = originalFetch
+  }
+}
+
 describe("startAuthServer()", () => {
   it("starts on a localhost port and accepts a valid callback POST", async () => {
-    const { server, port, waitForCallback } = await startAuthServer({ startPort: 0 })
+    const { server, port, waitForCallback } = await startAuthServer({
+      startPort: 0,
+      expectedState: "test-state-token",
+    })
 
     const callbackData: AuthCallback = {
       apiKey: "user_testKey123",
@@ -57,6 +75,42 @@ describe("startAuthServer()", () => {
     await waitForClose(server)
   })
 
+  it("rejects a mismatched state without closing the callback server", async () => {
+    const { server, port, waitForCallback } = await startAuthServer({
+      startPort: 0,
+      expectedState: "correct-state",
+    })
+
+    const invalidResponse = await fetch(`http://127.0.0.1:${port}/callback`, {
+      method: "POST",
+      headers: { "Content-Type": "application/json", Origin: "https://commandcode.ai" },
+      body: JSON.stringify({
+        apiKey: "user_badState",
+        state: "wrong-state",
+        userId: "user_789",
+        userName: "Attacker",
+        keyName: "evil-key",
+      }),
+    })
+    assert.equal(invalidResponse.status, 403)
+    assert.equal(server.listening, true)
+
+    const validResponse = await fetch(`http://127.0.0.1:${port}/callback`, {
+      method: "POST",
+      headers: { "Content-Type": "application/json", Origin: "https://commandcode.ai" },
+      body: JSON.stringify({
+        apiKey: "user_valid",
+        state: "correct-state",
+        userId: "user_123",
+        userName: "Valid User",
+        keyName: "valid-key",
+      }),
+    })
+    assert.equal(validResponse.status, 200)
+    assert.equal((await waitForCallback).apiKey, "user_valid")
+    await waitForClose(server)
+  })
+
   it("rejects when the callback indicates access_denied", async () => {
     const { server, port, waitForCallback } = await startAuthServer({ startPort: 0 })
 
@@ -176,6 +230,18 @@ describe("OAuth functions", () => {
   it("sanitizeApiKey removes paste markers, control chars, and whitespace", () => {
     assert.equal(sanitizeApiKey("\u001b[200~  user_manualKey\n\u001b[201~"), "user_manualKey")
   })
+
+  it("validates manual API keys through whoami", async () => {
+    await validateApiKey("valid-key", {
+      fetchImpl: () => Promise.resolve(new Response(JSON.stringify({ user: {} }), { status: 200 })),
+    })
+    await assert.rejects(
+      validateApiKey("invalid-key", {
+        fetchImpl: () => Promise.resolve(new Response("unauthorized", { status: 401 })),
+      }),
+      /Invalid Command Code API key/,
+    )
+  })
 })
 
 describe("login()", () => {
@@ -242,15 +308,17 @@ describe("login()", () => {
     const promptMessages: string[] = []
 
     try {
-      const result = await login({
-        onAuth(params: { url: string }) {
-          authUrl = params.url
-        },
-        async onPrompt(params: { message: string }): Promise<string> {
-          promptMessages.push(params.message)
-          return promptMessages.length === 1 ? "" : "\u001b[200~  user_manualApiKey\n\u001b[201~"
-        },
-      })
+      const result = await withValidApiKeyFetch(() =>
+        login({
+          onAuth(params: { url: string }) {
+            authUrl = params.url
+          },
+          async onPrompt(params: { message: string }): Promise<string> {
+            promptMessages.push(params.message)
+            return promptMessages.length === 1 ? "" : "\u001b[200~  user_manualApiKey\n\u001b[201~"
+          },
+        }),
+      )
 
       assert.match(authUrl, /^https:\/\/commandcode\.ai\/studio\/auth\/cli\?/)
       assert.match(promptMessages[1] ?? "", /Paste your Command Code API key/)
@@ -265,14 +333,16 @@ describe("login()", () => {
 
   it("accepts a directly pasted API key", async () => {
     let authOpened = false
-    const result = await login({
-      onAuth() {
-        authOpened = true
-      },
-      onPrompt(): Promise<string> {
-        return Promise.resolve("user_directApiKey")
-      },
-    })
+    const result = await withValidApiKeyFetch(() =>
+      login({
+        onAuth() {
+          authOpened = true
+        },
+        onPrompt(): Promise<string> {
+          return Promise.resolve("user_directApiKey")
+        },
+      }),
+    )
 
     assert.equal(authOpened, false)
     assert.equal(result.access, "user_directApiKey")
@@ -280,21 +350,23 @@ describe("login()", () => {
 
   it("offers an explicit API key prompt", async () => {
     let promptCount = 0
-    const result = await login({
-      onAuth() {
-        throw new Error("browser should not open")
-      },
-      onPrompt(): Promise<string> {
-        promptCount += 1
-        return Promise.resolve(promptCount === 1 ? "key" : "user_promptedApiKey")
-      },
-    })
+    const result = await withValidApiKeyFetch(() =>
+      login({
+        onAuth() {
+          throw new Error("browser should not open")
+        },
+        onPrompt(): Promise<string> {
+          promptCount += 1
+          return Promise.resolve(promptCount === 1 ? "key" : "user_promptedApiKey")
+        },
+      }),
+    )
 
     assert.equal(result.access, "user_promptedApiKey")
     assert.equal(promptCount, 2)
   })
 
-  it("rejects on state token mismatch", async () => {
+  it("keeps waiting after a state mismatch and accepts the legitimate callback", async () => {
     let authUrl = ""
     const callbacks = {
       onAuth(params: { url: string }) {
@@ -305,12 +377,7 @@ describe("login()", () => {
       },
     }
 
-    const loginPromise: Promise<string> = login(callbacks).then(
-      () => {
-        throw new Error("Expected login to reject")
-      },
-      (e: Error) => e.message,
-    )
+    const loginPromise = login(callbacks)
 
     // Wait for onAuth to be called asynchronously
     while (!authUrl) await new Promise((resolve) => setTimeout(resolve, 10))
@@ -318,8 +385,8 @@ describe("login()", () => {
     const url = new URL(authUrl)
     const port = parseInt(url.searchParams.get("callback")?.match(/localhost:(\d+)/)?.[1] ?? "0")
 
-    // Post back with a wrong state token
-    await fetch(`http://127.0.0.1:${port}/callback`, {
+    // Post back with a wrong state token.
+    const invalidResponse = await fetch(`http://127.0.0.1:${port}/callback`, {
       method: "POST",
       headers: { "Content-Type": "application/json", Origin: "https://commandcode.ai" },
       body: JSON.stringify({
@@ -331,7 +398,20 @@ describe("login()", () => {
       }),
     })
 
-    const errorMsg = await loginPromise
-    assert.match(errorMsg, /State token mismatch/)
+    assert.equal(invalidResponse.status, 403)
+
+    const validResponse = await fetch(`http://127.0.0.1:${port}/callback`, {
+      method: "POST",
+      headers: { "Content-Type": "application/json", Origin: "https://commandcode.ai" },
+      body: JSON.stringify({
+        apiKey: "user_goodState",
+        state: url.searchParams.get("state"),
+        userId: "user_123",
+        userName: "Real User",
+        keyName: "real-key",
+      }),
+    })
+    assert.equal(validResponse.status, 200)
+    assert.equal((await loginPromise).access, "user_goodState")
   })
 })

tests/test-pure-functions.ts

@@ -27,8 +27,18 @@ import { redactCommandCodeErrorText } from "../src/overflow.ts"
 import { objectAt } from "./helpers.ts"
 
 describe("getApiKey()", () => {
-  it("uses COMMANDCODE_API_KEY from provided env", () => {
-    assert.equal(getApiKey({ env: { COMMANDCODE_API_KEY: "env-key" }, authPaths: [] }), "env-key")
+  it("uses the official API key env var before the legacy alias", () => {
+    assert.equal(
+      getApiKey({
+        env: { COMMAND_CODE_API_KEY: "official-key", COMMANDCODE_API_KEY: "legacy-key" },
+        authPaths: [],
+      }),
+      "official-key",
+    )
+    assert.equal(
+      getApiKey({ env: { COMMANDCODE_API_KEY: "legacy-key" }, authPaths: [] }),
+      "legacy-key",
+    )
   })
 
   it("reads apiKey, commandcode, pi OAuth, and official CLI credential fields", () => {
@@ -109,11 +119,14 @@ describe("error redaction", () => {
 
 describe("pickCommandCodeApiKey()", () => {
   it("falls back to the host key for a placeholder registry value", () => {
+    assert.equal(pickCommandCodeApiKey("$COMMAND_CODE_API_KEY", "file-key"), "file-key")
+    assert.equal(pickCommandCodeApiKey("COMMAND_CODE_API_KEY", "file-key"), "file-key")
     assert.equal(pickCommandCodeApiKey("$COMMANDCODE_API_KEY", "file-key"), "file-key")
     assert.equal(pickCommandCodeApiKey("COMMANDCODE_API_KEY", "file-key"), "file-key")
   })
 
   it("returns undefined when only a placeholder is provided (no fallback)", () => {
+    assert.equal(pickCommandCodeApiKey("$COMMAND_CODE_API_KEY", undefined), undefined)
     assert.equal(pickCommandCodeApiKey("$COMMANDCODE_API_KEY", undefined), undefined)
   })
 
@@ -199,6 +212,12 @@ describe("textContent()", () => {
     )
   })
 
+  it("normalizes malformed string and object content", () => {
+    assert.equal(textContent({ content: "raw result" }), "raw result")
+    assert.equal(textContent({ content: { ok: true } }), '{"ok":true}')
+    assert.equal(textContent({ content: null }), "")
+  })
+
   it("handles empty or missing content", () => {
     assert.equal(textContent({ content: [] }), "")
     assert.equal(textContent({}), "")
@@ -531,6 +550,23 @@ describe("messagesToCC()", () => {
     assert.equal(objectAt(result, ["2", "content", "0", "output", "value"]), "hello\nworld")
   })
 
+  it("preserves malformed string tool results instead of sending empty output", () => {
+    const result = messagesToCC([
+      {
+        role: "assistant",
+        content: [{ type: "toolCall", id: "c1", name: "read", arguments: {} }],
+      },
+      {
+        role: "toolResult",
+        toolCallId: "c1",
+        toolName: "read",
+        content: "raw result",
+      },
+    ])
+
+    assert.equal(objectAt(result, ["1", "content", "0", "output", "value"]), "raw result")
+  })
+
   it("serializes image inputs in the current Command Code wire format", () => {
     assert.deepEqual(
       messagesToCC(
@@ -632,7 +668,7 @@ describe("messagesToCC()", () => {
     ])
   })
 
-  it("drops orphaned tool calls that have no matching tool result", () => {
+  it("synthesizes missing results for orphaned tool calls", () => {
     const result = messagesToCC([
       { role: "user", content: "edit a file" },
       {
@@ -651,7 +687,12 @@ describe("messagesToCC()", () => {
 
     assert.equal(objectAt(result, ["1", "role"]), "assistant")
     assert.equal(objectAt(result, ["1", "content", "0", "type"]), "text")
-    assert.equal(objectAt(result, ["1", "content", "1"]), undefined)
+    assert.equal(objectAt(result, ["1", "content", "1", "type"]), "tool-call")
+    assert.equal(objectAt(result, ["2", "role"]), "tool")
+    assert.match(
+      String(objectAt(result, ["2", "content", "0", "output", "value"])),
+      /did not complete/,
+    )
   })
 
   it("handles empty conversations", () => {

tests/test-stream.ts

@@ -77,6 +77,23 @@ describe("streamCommandCode — auth", () => {
     )
   })
 
+  it("accepts the official CLI API key environment variable", async () => {
+    server.mockResponse({
+      type: "success",
+      events: [JSON.stringify({ type: "finish", finishReason: "stop" })],
+    })
+    const { streamCommandCode } = createTestDeps({
+      apiBase: server.baseUrl(),
+      env: { COMMAND_CODE_API_KEY: "official-env-key" },
+    })
+
+    await collectEvents(
+      streamCommandCode(makeModel(), makeContext(), { apiKey: "$COMMAND_CODE_API_KEY" }),
+    )
+
+    assert.equal(server.lastRequestHeaders().authorization, "Bearer official-env-key")
+  })
+
   it("uses options.apiKey in the Authorization header", async () => {
     server.mockResponse({
       type: "success",
@@ -555,7 +572,7 @@ describe("streamCommandCode — request serialization", () => {
     assert.equal(objectAt(body, ["params", "stream"]), true)
     assert.equal(objectAt(body, ["params", "max_tokens"]), 500)
     assert.equal(objectAt(body, ["params", "reasoning_effort"]), undefined)
-    assert.equal(objectAt(body, ["params", "temperature"]), 0.3)
+    assert.equal(objectAt(body, ["params", "temperature"]), undefined)
     assert.equal(objectAt(body, ["params", "system"]), "You are a test assistant.")
     assert.equal(objectAt(body, ["memory"]), null)
     assert.equal(objectAt(body, ["taste"]), null)
@@ -573,10 +590,53 @@ describe("streamCommandCode — request serialization", () => {
     assert.equal(headers["x-command-code-version"], COMMAND_CODE_CLI_VERSION)
     assert.equal(headers["x-project-slug"], "repo")
     assert.equal(headers["x-taste-learning"], "true")
-    assert.equal(headers["x-co-flag"], "false")
+    assert.equal(headers["user-agent"], "cli")
+    assert.equal(headers["x-co-flag"], undefined)
     assert.equal(headers["x-session-id"], undefined)
   })
 
+  it("forwards explicit temperature and stable session metadata", async () => {
+    server.mockResponse({
+      type: "success",
+      events: [JSON.stringify({ type: "finish", finishReason: "stop" })],
+    })
+    const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() })
+
+    await collectEvents(
+      streamCommandCode(makeModel(), makeContext(), {
+        apiKey: "mock-key",
+        temperature: 0.7,
+        sessionId: "11111111-1111-4111-8111-111111111111",
+      }),
+    )
+
+    const body = server.lastRequestBody()
+    assert.equal(objectAt(body, ["params", "temperature"]), 0.7)
+    assert.equal(objectAt(body, ["threadId"]), "11111111-1111-4111-8111-111111111111")
+    assert.equal(
+      server.lastRequestHeaders()["x-session-id"],
+      "11111111-1111-4111-8111-111111111111",
+    )
+  })
+
+  it("omits non-UUID session ids from the generate thread id", async () => {
+    server.mockResponse({
+      type: "success",
+      events: [JSON.stringify({ type: "finish", finishReason: "stop" })],
+    })
+    const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() })
+
+    await collectEvents(
+      streamCommandCode(makeModel(), makeContext(), {
+        apiKey: "mock-key",
+        sessionId: "human-readable-session",
+      }),
+    )
+
+    assert.equal(objectAt(server.lastRequestBody(), ["threadId"]), undefined)
+    assert.equal(server.lastRequestHeaders()["x-session-id"], "human-readable-session")
+  })
+
   it("accepts the legacy OMP nested reasoning map", async () => {
     server.mockResponse({
       type: "success",
@@ -819,6 +879,63 @@ describe("streamCommandCode — upstream errors and malformed streams", () => {
     assert.equal(error.error.errorMessage, "provider failed")
   })
 
+  it("rejects a truncated stream without a finish event", async () => {
+    server.mockResponse({
+      type: "success",
+      events: [JSON.stringify({ type: "text-delta", text: "truncated" })],
+    })
+    const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() })
+
+    const events = await collectEvents(
+      streamCommandCode(makeModel(), makeContext(), { apiKey: "mock-key" }),
+    )
+
+    const error = events.at(-1)
+    assert.equal(error?.type, "error")
+    if (error?.type !== "error") throw new Error("expected error")
+    assert.match(error.error.errorMessage ?? "", /no finish event/i)
+  })
+
+  it("maps an upstream abort event to an aborted request", async () => {
+    server.mockResponse({
+      type: "success",
+      events: [JSON.stringify({ type: "abort" })],
+    })
+    const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() })
+
+    const events = await collectEvents(
+      streamCommandCode(makeModel(), makeContext(), { apiKey: "mock-key" }),
+    )
+
+    const error = events.at(-1)
+    assert.equal(error?.type, "error")
+    if (error?.type !== "error") throw new Error("expected error")
+    assert.equal(error.reason, "aborted")
+  })
+
+  it("rejects terminal upstream network failure reasons", async () => {
+    server.mockResponse({
+      type: "success",
+      events: [
+        JSON.stringify({
+          type: "finish",
+          finishReason: "stop",
+          rawFinishReason: "upstream_error",
+        }),
+      ],
+    })
+    const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() })
+
+    const events = await collectEvents(
+      streamCommandCode(makeModel(), makeContext(), { apiKey: "mock-key" }),
+    )
+
+    const error = events.at(-1)
+    assert.equal(error?.type, "error")
+    if (error?.type !== "error") throw new Error("expected error")
+    assert.match(error.error.errorMessage ?? "", /upstream connection failed/i)
+  })
+
   it("handles SSE lines, malformed lines, split chunks, and final line without newline", async () => {
     const textEvent = `data: ${JSON.stringify({ type: "text-delta", text: "split" })}\n`
     const finishEvent = JSON.stringify({