Add Command Code browser login fallback

README.md

@@ -40,19 +40,33 @@ Then reload pi:
 
 Set your Command Code API key using one of these methods:
 
-### 1. Environment variable
+### 1. Browser login (recommended)
 
-```sh
+In pi, run:
-export COMMANDCODE_API_KEY="cc-..."
+
+```txt
+/login
 ```
 
-### 2. Auth file (recommended)
+Then select **Command Code** from the provider list.
+
+This opens Command Code in your browser and stores the returned API key in pi's auth file. If the browser shows "Copy your API key" because automatic transfer failed, copy that key and paste it into the pi terminal prompt.
+
+> Note: `/login commandcode` is not supported by pi currently; use interactive `/login` and select Command Code.
+
+### 2. Environment variable
+
+```sh
+export COMMANDCODE_API_KEY="user_..."
+```
+
+### 3. Auth file
 
 Create `~/.commandcode/auth.json`:
 
 ```json
 {
-  "apiKey": "cc-..."
+  "apiKey": "user_..."
 }
 ```
 
@@ -60,7 +74,7 @@ Or use pi's auth file at `~/.pi/agent/auth.json`:
 
 ```json
 {
-  "commandcode": "cc-..."
+  "commandcode": "user_..."
 }
 ```
 
@@ -69,7 +83,7 @@ Or use pi's auth file at `~/.pi/agent/auth.json`:
 After installing and setting your API key, select a Command Code model in pi:
 
 ```txt
-/model claude-sonnet-4-6
+/model deepseek/deepseek-v4-flash
 ```
 
 Any query will then use the Command Code API. You can list available models:

index.ts

@@ -4,7 +4,7 @@
  * Connects pi to Command Code's API (https://api.commandcode.ai/alpha/generate).
  *
  * Authentication (pick one):
- *   1. Run `/login commandcode` — opens browser to commandcode.ai, auto-stores API key
+ *   1. Run `/login`, then select Command Code — opens browser to commandcode.ai, auto-stores API key
  *   2. Set COMMANDCODE_API_KEY environment variable
  *   3. Place API key in `~/.commandcode/auth.json` or `~/.pi/agent/auth.json`
  *      as {"apiKey": "user_..."} or {"commandcode": "user_..."}

package.json

@@ -25,11 +25,12 @@
     "LICENSE"
   ],
   "scripts": {
-    "test": "npm run typecheck && tsx tests/test-pure-functions.ts && tsx tests/test-abort.ts && tsx tests/test-stream.ts && node tests/test-pi-local.mjs",
+    "test": "npm run typecheck && tsx tests/test-pure-functions.ts && tsx tests/test-oauth.ts && tsx tests/test-abort.ts && tsx tests/test-stream.ts && node tests/test-pi-local.mjs",
     "typecheck": "tsc --noEmit",
     "format:check": "prettier --check '**/*.{ts,mjs,json,md}'",
     "format": "prettier --write '**/*.{ts,mjs,json,md}'",
     "test:unit": "tsx tests/test-pure-functions.ts",
+    "test:oauth": "tsx tests/test-oauth.ts",
     "test:abort": "tsx tests/test-abort.ts",
     "test:stream": "tsx tests/test-stream.ts",
     "test:pi-local": "node tests/test-pi-local.mjs",

src/auth-server.ts

@@ -1,13 +1,16 @@
 /**
  * Local HTTP callback server for the Command Code browser auth flow.
  *
- * Starts a one-shot server on a random port. The Command Code Studio
+ * Starts a one-shot server on a CLI-compatible localhost port. The Command Code
- * website POSTs the user's API key to /callback after they authenticate.
+ * Studio website POSTs the user's API key to /callback after they authenticate.
  */
 
 import { createServer, type Server } from "node:http"
 import type { AddressInfo } from "node:net"
 
+const DEFAULT_PORT = 5959
+const DEFAULT_PORT_RANGE = 10
+
 export interface AuthCallback {
   apiKey: string
   state: string
@@ -22,15 +25,65 @@ export interface AuthServer {
   waitForCallback: Promise<AuthCallback>
 }
 
+export interface AuthServerOptions {
+  startPort?: number
+  portRange?: number
+}
+
+function listenOnAvailablePort(
+  server: Server,
+  startPort = DEFAULT_PORT,
+  range = DEFAULT_PORT_RANGE,
+): Promise<number> {
+  return new Promise((resolve, reject) => {
+    let offset = 0
+
+    const tryListen = () => {
+      const useFallbackPort = startPort === 0 || offset >= range
+      const port = useFallbackPort ? 0 : startPort + offset
+
+      const onError = (err: NodeJS.ErrnoException) => {
+        server.off("listening", onListening)
+        if (err.code === "EADDRINUSE" && !useFallbackPort) {
+          offset += 1
+          tryListen()
+          return
+        }
+        reject(err)
+      }
+
+      const onListening = () => {
+        server.off("error", onError)
+        const address = server.address() as AddressInfo
+        resolve(address.port)
+      }
+
+      server.once("error", onError)
+      server.once("listening", onListening)
+      server.listen(port, "127.0.0.1")
+    }
+
+    tryListen()
+  })
+}
+
+function closeServer(server: Server) {
+  server.close((err: NodeJS.ErrnoException | undefined) => {
+    if (err && err.code !== "ERR_SERVER_NOT_RUNNING") {
+      // There is nowhere useful to report this during auth cleanup.
+    }
+  })
+}
+
 /**
  * Start a local HTTP server that listens for the Command Code Studio
  * to POST the API key after the user authenticates in their browser.
  *
  * The server accepts exactly one valid POST to /callback and then closes.
  */
-export function startAuthServer(): Promise<AuthServer> {
+export async function startAuthServer(options: AuthServerOptions = {}): Promise<AuthServer> {
-  let resolveCallback: (value: AuthCallback) => void
+  let resolveCallback!: (value: AuthCallback) => void
-  let rejectCallback: (error: Error) => void
+  let rejectCallback!: (error: Error) => void
 
   const waitForCallback = new Promise<AuthCallback>((resolve, reject) => {
     resolveCallback = resolve
@@ -38,7 +91,7 @@ export function startAuthServer(): Promise<AuthServer> {
   })
 
   const server = createServer((req, res) => {
-    // CORS: allow requests from Command Code domains and localhost for dev
+    // CORS: allow requests from Command Code domains and localhost for dev.
     const origin = req.headers.origin || ""
     const allowedOrigins = [
       "http://localhost:3000",
@@ -46,13 +99,22 @@ export function startAuthServer(): Promise<AuthServer> {
       "https://commandcode.ai",
     ]
     const responseOrigin = allowedOrigins.includes(origin) ? origin : allowedOrigins[0]
+    const requestedHeaders = req.headers["access-control-request-headers"]
 
     res.setHeader("Access-Control-Allow-Origin", responseOrigin)
     res.setHeader("Access-Control-Allow-Methods", "POST, OPTIONS")
-    res.setHeader("Access-Control-Allow-Headers", "Content-Type")
+    res.setHeader(
+      "Access-Control-Allow-Headers",
+      typeof requestedHeaders === "string" && requestedHeaders.length > 0
+        ? requestedHeaders
+        : "Content-Type",
+    )
+    // Chrome's Private Network Access preflight may require this for an HTTPS
+    // page posting to a localhost HTTP callback.
+    res.setHeader("Access-Control-Allow-Private-Network", "true")
     res.setHeader("Content-Type", "application/json")
 
-    // Handle CORS preflight
+    // Handle CORS preflight.
     if (req.method === "OPTIONS") {
       res.writeHead(204)
       res.end()
@@ -98,7 +160,7 @@ export function startAuthServer(): Promise<AuthServer> {
           } else {
             rejectCallback(new Error(description || String(parsed.error)))
           }
-          server.close()
+          closeServer(server)
           return
         }
 
@@ -123,7 +185,7 @@ export function startAuthServer(): Promise<AuthServer> {
         res.end(JSON.stringify({ success: true }))
 
         resolveCallback({ apiKey, state, userId, userName, keyName })
-        server.close()
+        closeServer(server)
       } catch {
         res.writeHead(400)
         res.end(JSON.stringify({ success: false, error: "Invalid JSON" }))
@@ -136,14 +198,17 @@ export function startAuthServer(): Promise<AuthServer> {
     })
   })
 
-  return new Promise((resolve) => {
+  try {
-    server.on("error", (err: NodeJS.ErrnoException) => {
+    const port = await listenOnAvailablePort(
-      rejectCallback(new Error(`Failed to start auth server: ${err.message}`))
+      server,
-    })
+      options.startPort ?? DEFAULT_PORT,
-
+      options.portRange ?? DEFAULT_PORT_RANGE,
-    server.listen(0, "127.0.0.1", () => {
+    )
-      const address = server.address() as AddressInfo
+    return { server, port, waitForCallback }
-      resolve({ server, port: address.port, waitForCallback })
+  } catch (err) {
-    })
+    const message = err instanceof Error ? err.message : String(err)
-  })
+    const error = new Error(`Failed to start auth server: ${message}`)
+    rejectCallback(error)
+    throw error
+  }
 }

src/core.ts

@@ -129,7 +129,7 @@ export function createStreamCommandCode(deps: CoreDependencies) {
           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.",
+            "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 })

src/oauth.ts

@@ -2,11 +2,12 @@
  * 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 random port
+ * 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. The API key is stored in pi's auth.json as OAuth credentials
+ * 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.
@@ -17,6 +18,7 @@ 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
@@ -29,10 +31,76 @@ export interface OAuthCredentials {
   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.
  *
@@ -40,37 +108,47 @@ function generateStateToken(): string {
  * The keys don't expire, so we set a far-future expiry.
  */
 export async function login(callbacks: OAuthLoginCallbacks): Promise<OAuthCredentials> {
-  const authServer = await startAuthServer()
+  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)}`
 
-  const authUrl = `${STUDIO_BASE_URL}/studio/auth/cli?callback=${encodeURIComponent(`http://localhost:${authServer.port}/callback`)}&state=${encodeURIComponent(stateToken)}`
+  // Tell pi to open the browser.
-
-  // Tell pi to open the browser
   callbacks.onAuth({ url: authUrl })
 
-  // Wait for the Command Code Studio to POST the API key back
+  // 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 authServer.waitForCallback
+    callback = await withTimeout(authServer.waitForCallback, getAuthTimeoutMs())
   } catch (error) {
-    // Clean up server on 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
+  // 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 as OAuth credentials. Since CC API keys don't expire,
+  return credentialsFromApiKey(callback.apiKey)
-  // we set a far-future expiry and use the API key as both access and refresh.
-  return {
-    refresh: callback.apiKey,
-    access: callback.apiKey,
-    expires: Date.now() + TEN_YEARS_MS,
-  }
 }
 
 /**
@@ -78,11 +156,7 @@ export async function login(callbacks: OAuthLoginCallbacks): Promise<OAuthCreden
  * Returns the same credentials with an updated far-future expiry.
  */
 export async function refreshToken(credentials: OAuthCredentials): Promise<OAuthCredentials> {
-  return {
+  return credentialsFromApiKey(credentials.refresh)
-    refresh: credentials.refresh,
-    access: credentials.access,
-    expires: Date.now() + TEN_YEARS_MS,
-  }
 }
 
 /**

tests/test-oauth.ts

@@ -9,11 +9,24 @@ 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 } from "../src/oauth.ts"
+import { getApiKey, login, refreshToken, sanitizeApiKey } from "../src/oauth.ts"
+
+/**
+ * Helper: wait for an HTTP server to close, or resolve immediately if already closed.
+ */
+function waitForClose(server: {
+  listening: boolean
+  on(event: "close", cb: () => void): void
+}): Promise<void> {
+  return new Promise((resolve) => {
+    if (!server.listening) return resolve(undefined)
+    server.on("close", resolve)
+  })
+}
 
 describe("startAuthServer()", () => {
-  it("starts on a random port and accepts a valid callback POST", async () => {
+  it("starts on a localhost port and accepts a valid callback POST", async () => {
-    const { server, port, waitForCallback } = await startAuthServer()
+    const { server, port, waitForCallback } = await startAuthServer({ startPort: 0 })
 
     const callbackData: AuthCallback = {
       apiKey: "user_testKey123",
@@ -26,10 +39,7 @@ describe("startAuthServer()", () => {
     // Simulate the Command Code Studio posting the API key back
     const response = await fetch(`http://127.0.0.1:${port}/callback`, {
       method: "POST",
-      headers: {
+      headers: { "Content-Type": "application/json", Origin: "https://commandcode.ai" },
-        "Content-Type": "application/json",
-        Origin: "https://commandcode.ai",
-      },
       body: JSON.stringify(callbackData),
     })
 
@@ -44,47 +54,40 @@ describe("startAuthServer()", () => {
     assert.equal(result.userName, "Test User")
     assert.equal(result.keyName, "test-key")
 
-    // Server closes itself after callback; ensure it's done
+    await waitForClose(server)
-    await new Promise((resolve) => {
-      if (!server.listening) return resolve(undefined)
-      server.on("close", resolve)
-    })
   })
 
   it("rejects when the callback indicates access_denied", async () => {
-    const { server, port, waitForCallback } = await startAuthServer()
+    const { server, port, waitForCallback } = await startAuthServer({ startPort: 0 })
+
+    // Attach rejection handler before posting to avoid unhandled rejection
+    const errorPromise: Promise<string> = waitForCallback.then(
+      () => {
+        throw new Error("Expected callback to reject")
+      },
+      (e: Error) => e.message,
+    )
 
     const response = await fetch(`http://127.0.0.1:${port}/callback`, {
       method: "POST",
-      headers: {
+      headers: { "Content-Type": "application/json", Origin: "https://commandcode.ai" },
-        "Content-Type": "application/json",
+      body: JSON.stringify({ error: "access_denied", error_description: "User cancelled" }),
-        Origin: "https://commandcode.ai",
-      },
-      body: JSON.stringify({
-        error: "access_denied",
-        error_description: "User cancelled",
-      }),
     })
 
     assert.equal(response.status, 200)
 
-    await assert.rejects(() => waitForCallback, /User cancelled/)
+    const errorMsg = await errorPromise
+    assert.match(errorMsg, /User cancelled/)
 
-    await new Promise((resolve) => {
+    await waitForClose(server)
-      if (!server.listening) return resolve(undefined)
-      server.on("close", resolve)
-    })
   })
 
   it("returns 400 for missing required fields", async () => {
-    const { server, port } = await startAuthServer()
+    const { server, port } = await startAuthServer({ startPort: 0 })
 
     const response = await fetch(`http://127.0.0.1:${port}/callback`, {
       method: "POST",
-      headers: {
+      headers: { "Content-Type": "application/json", Origin: "https://commandcode.ai" },
-        "Content-Type": "application/json",
-        Origin: "https://commandcode.ai",
-      },
       body: JSON.stringify({ apiKey: "key", state: "s" }),
     })
 
@@ -94,22 +97,32 @@ describe("startAuthServer()", () => {
     server.close()
   })
 
-  it("handles CORS preflight OPTIONS request", async () => {
+  it("handles CORS and private-network preflight OPTIONS request", async () => {
-    const { server, port } = await startAuthServer()
+    const { server, port } = await startAuthServer({ startPort: 0 })
 
     const response = await fetch(`http://127.0.0.1:${port}/callback`, {
       method: "OPTIONS",
-      headers: { Origin: "https://commandcode.ai" },
+      headers: {
+        Origin: "https://commandcode.ai",
+        "Access-Control-Request-Headers": "content-type,x-requested-with",
+        "Access-Control-Request-Private-Network": "true",
+      },
     })
 
     assert.equal(response.status, 204)
+    assert.equal(response.headers.get("access-control-allow-origin"), "https://commandcode.ai")
+    assert.equal(
+      response.headers.get("access-control-allow-headers"),
+      "content-type,x-requested-with",
+    )
+    assert.equal(response.headers.get("access-control-allow-private-network"), "true")
 
     await new Promise((resolve) => setTimeout(resolve, 100))
     server.close()
   })
 
   it("returns 404 for non-callback paths", async () => {
-    const { server, port } = await startAuthServer()
+    const { server, port } = await startAuthServer({ startPort: 0 })
 
     const response = await fetch(`http://127.0.0.1:${port}/other`, {
       method: "POST",
@@ -124,7 +137,7 @@ describe("startAuthServer()", () => {
   })
 
   it("returns 405 for GET on /callback", async () => {
-    const { server, port } = await startAuthServer()
+    const { server, port } = await startAuthServer({ startPort: 0 })
 
     const response = await fetch(`http://127.0.0.1:${port}/callback`, {
       method: "GET",
@@ -159,6 +172,10 @@ describe("OAuth functions", () => {
     assert.equal(result.refresh, "my-api-key")
     assert.ok(result.expires > Date.now(), "expiry should be in the future")
   })
+
+  it("sanitizeApiKey removes paste markers, control chars, and whitespace", () => {
+    assert.equal(sanitizeApiKey("\u001b[200~  user_manualKey\n\u001b[201~"), "user_manualKey")
+  })
 })
 
 describe("login()", () => {
@@ -168,7 +185,7 @@ describe("login()", () => {
       onAuth(params: { url: string }) {
         authUrl = params.url
       },
-      onPrompt(params: { message: string }): Promise<string> {
+      onPrompt(_params: { message: string }): Promise<string> {
         throw new Error("onPrompt should not be called in browser flow")
       },
     }
@@ -176,10 +193,13 @@ describe("login()", () => {
     // Start login in the background
     const loginPromise = login(callbacks)
 
-    // Verify the auth URL was passed to callbacks
+    // Wait for onAuth to be called (it fires asynchronously after the auth server starts)
+    while (!authUrl) await new Promise((resolve) => setTimeout(resolve, 10))
+
+    // Verify the auth URL was passed to callbacks (callback URL is encoded)
     assert.match(
       authUrl,
-      /^https:\/\/commandcode\.ai\/studio\/auth\/cli\?callback=http:\/\/localhost:\d+\/callback&state=/,
+      /^https:\/\/commandcode\.ai\/studio\/auth\/cli\?callback=http%3A%2F%2Flocalhost%3A\d+%2Fcallback&state=/,
     )
 
     // Extract port and state from the URL
@@ -214,18 +234,55 @@ describe("login()", () => {
     assert.ok(result.expires > Date.now(), "expiry should be far in the future")
   })
 
+  it("prompts for a manual API key if browser transfer times out", async () => {
+    const originalTimeout = process.env.COMMANDCODE_AUTH_TIMEOUT_MS
+    process.env.COMMANDCODE_AUTH_TIMEOUT_MS = "1"
+
+    let authUrl = ""
+    let promptMessage = ""
+
+    try {
+      const result = await login({
+        onAuth(params: { url: string }) {
+          authUrl = params.url
+        },
+        async onPrompt(params: { message: string }): Promise<string> {
+          promptMessage = params.message
+          return "\u001b[200~  user_manualApiKey\n\u001b[201~"
+        },
+      })
+
+      assert.match(authUrl, /^https:\/\/commandcode\.ai\/studio\/auth\/cli\?/)
+      assert.match(promptMessage, /Paste your Command Code API key/)
+      assert.equal(result.access, "user_manualApiKey")
+      assert.equal(result.refresh, "user_manualApiKey")
+      assert.ok(result.expires > Date.now(), "expiry should be far in the future")
+    } finally {
+      if (originalTimeout === undefined) delete process.env.COMMANDCODE_AUTH_TIMEOUT_MS
+      else process.env.COMMANDCODE_AUTH_TIMEOUT_MS = originalTimeout
+    }
+  })
+
   it("rejects on state token mismatch", async () => {
     let authUrl = ""
     const callbacks = {
       onAuth(params: { url: string }) {
         authUrl = params.url
       },
-      onPrompt(params: { message: string }): Promise<string> {
+      onPrompt(_params: { message: string }): Promise<string> {
         throw new Error("should not prompt")
       },
     }
 
-    const loginPromise = login(callbacks)
+    const loginPromise: Promise<string> = login(callbacks).then(
+      () => {
+        throw new Error("Expected login to reject")
+      },
+      (e: Error) => e.message,
+    )
+
+    // Wait for onAuth to be called asynchronously
+    while (!authUrl) await new Promise((resolve) => setTimeout(resolve, 10))
 
     const url = new URL(authUrl)
     const port = parseInt(url.searchParams.get("callback")?.match(/localhost:(\d+)/)?.[1] ?? "0")
@@ -233,10 +290,7 @@ describe("login()", () => {
     // Post back with a wrong state token
     await fetch(`http://127.0.0.1:${port}/callback`, {
       method: "POST",
-      headers: {
+      headers: { "Content-Type": "application/json", Origin: "https://commandcode.ai" },
-        "Content-Type": "application/json",
-        Origin: "https://commandcode.ai",
-      },
       body: JSON.stringify({
         apiKey: "user_badState",
         state: "wrong-state-token",
@@ -246,6 +300,7 @@ describe("login()", () => {
       }),
     })
 
-    await assert.rejects(() => loginPromise, /State token mismatch/)
+    const errorMsg = await loginPromise
+    assert.match(errorMsg, /State token mismatch/)
   })
 })