Fix Parallel OAuth callback and browser handling

packages/parallel-oauth/src/__tests__/index.test.ts

@@ -1,6 +1,7 @@
 import { afterEach, describe, expect, it, vi } from 'vitest';
 import { createHash } from 'node:crypto';
 import { get } from 'node:http';
+import { connect, type Socket } from 'node:net';
 import { loginWithParallel } from '../index.js';
 
 const PLATFORM_ORIGIN = 'https://example.test';
@@ -162,4 +163,92 @@ describe('loginWithParallel', () => {
 
     expect(fetchMock).not.toHaveBeenCalled();
   });
+
+  it('finishes login without waiting for the callback client to disconnect', async () => {
+    stubTokenExchange();
+    let socket: Socket | undefined;
+
+    const loginPromise = loginWithParallel({
+      platformOrigin: PLATFORM_ORIGIN,
+      openBrowser: false,
+      onAuthUrl: (rawUrl) => {
+        const authUrl = new URL(rawUrl);
+        const redirectUri = new URL(
+          authUrl.searchParams.get('redirect_uri') ?? ''
+        );
+        const state = authUrl.searchParams.get('state') ?? '';
+        socket = connect(Number(redirectUri.port), redirectUri.hostname, () => {
+          socket?.write(
+            `POST ${redirectUri.pathname}?code=abc&state=${encodeURIComponent(state)} HTTP/1.1\r\n` +
+              `Host: ${redirectUri.host}\r\n` +
+              'Transfer-Encoding: chunked\r\n' +
+              'Connection: keep-alive\r\n\r\n' +
+              '1\r\nx\r\n'
+          );
+        });
+        socket.on('error', () => {});
+        socket.resume();
+      },
+    });
+
+    let timer: ReturnType<typeof setTimeout> | undefined;
+    const outcome = await Promise.race([
+      loginPromise.then((result) => ({ result })),
+      new Promise<{ timedOut: true }>((resolve) => {
+        timer = setTimeout(() => resolve({ timedOut: true }), 1_000);
+      }),
+    ]);
+
+    clearTimeout(timer);
+    socket?.destroy();
+    await loginPromise;
+
+    expect(outcome).toEqual({ result: { apiKey: 'sk-test-key' } });
+  });
+
+  it('handles a follow-up request while the callback listener is closing', async () => {
+    let socket: Socket | undefined;
+    vi.stubGlobal(
+      'fetch',
+      vi.fn(async () => {
+        setImmediate(() => {
+          socket?.write(
+            '0\r\n\r\n' +
+              'GET /favicon.ico HTTP/1.1\r\n' +
+              'Host: 127.0.0.1\r\n' +
+              'Connection: close\r\n\r\n'
+          );
+        });
+        return new Response(JSON.stringify({ access_token: 'sk-test-key' }), {
+          status: 200,
+        });
+      })
+    );
+
+    const result = await loginWithParallel({
+      platformOrigin: PLATFORM_ORIGIN,
+      openBrowser: false,
+      onAuthUrl: (rawUrl) => {
+        const authUrl = new URL(rawUrl);
+        const redirectUri = new URL(
+          authUrl.searchParams.get('redirect_uri') ?? ''
+        );
+        const state = authUrl.searchParams.get('state') ?? '';
+        socket = connect(Number(redirectUri.port), redirectUri.hostname, () => {
+          socket?.write(
+            `POST ${redirectUri.pathname}?code=abc&state=${encodeURIComponent(state)} HTTP/1.1\r\n` +
+              `Host: ${redirectUri.host}\r\n` +
+              'Transfer-Encoding: chunked\r\n' +
+              'Connection: keep-alive\r\n\r\n' +
+              '1\r\nx\r\n'
+          );
+        });
+        socket.on('error', () => {});
+        socket.resume();
+        socket.setTimeout(1_000, () => socket?.destroy());
+      },
+    });
+
+    expect(result.apiKey).toBe('sk-test-key');
+  });
 });

packages/parallel-oauth/src/index.ts

@@ -95,8 +95,7 @@ async function startCallbackListener() {
 
   const server = createServer((req: IncomingMessage, res: ServerResponse) => {
     const requestUrl = req.url ?? '/';
-    const address = server.address() as AddressInfo;
-    const callbackUrl = `http://${LOOPBACK_HOST}:${address.port}${requestUrl}`;
+    const callbackUrl = `${callbackOrigin}${requestUrl}`;
     const url = new URL(callbackUrl);
 
     if (url.pathname !== '/callback') {
@@ -130,7 +129,8 @@ async function startCallbackListener() {
   });
 
   const address = server.address() as AddressInfo;
-  const redirectUri = `http://${LOOPBACK_HOST}:${address.port}/callback`;
+  const callbackOrigin = `http://${LOOPBACK_HOST}:${address.port}`;
+  const redirectUri = `${callbackOrigin}/callback`;
 
   return {
     redirectUri,
@@ -147,6 +147,7 @@ async function startCallbackListener() {
     async close() {
       await new Promise<void>((resolve) => {
         server.close(() => resolve());
+        server.closeAllConnections();
       });
     },
   };

packages/pi-extension/src/__tests__/parallel-auth.test.ts

@@ -115,9 +115,9 @@ describe('parallel-auth', () => {
     });
   });
 
-  it('runs the browser OAuth flow on login and returns an api_key credential', async () => {
+  it('delegates browser opening to Pi and returns an api_key credential', async () => {
     mocks.runParallelOAuth.mockImplementation(async (options) => {
-      options.onAuthUrl('https://platform.parallel.ai/oauth', true);
+      options.onAuthUrl('https://platform.parallel.ai/oauth', false);
       return { apiKey: 'fresh-key' };
     });
 
@@ -129,6 +129,9 @@ describe('parallel-auth', () => {
       key: 'fresh-key',
     });
 
+    expect(mocks.runParallelOAuth).toHaveBeenCalledWith(
+      expect.objectContaining({ openBrowser: false })
+    );
     expect(interaction.notify).toHaveBeenCalledWith(
       expect.objectContaining({
         type: 'auth_url',

packages/pi-extension/src/parallel-auth.ts

@@ -22,13 +22,12 @@ async function loginToParallel(
   interaction.signal.throwIfAborted();
 
   const { apiKey } = await runParallelOAuth({
-    onAuthUrl: (url, browserOpened) => {
+    openBrowser: false,
+    onAuthUrl: (url) => {
       interaction.notify({
         type: 'auth_url',
         url,
-        instructions: browserOpened
-          ? 'Opening Parallel login in your browser.'
-          : 'Open this URL to sign in to Parallel.',
+        instructions: 'Opening Parallel login in your browser.',
       });
     },
     promptForCallback: async (authUrl) => {