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")
   })
 })