diff --git a/packages/parallel-oauth/src/__tests__/index.test.ts b/packages/parallel-oauth/src/__tests__/index.test.ts index 02c57ca..c2fd837 100644 --- a/packages/parallel-oauth/src/__tests__/index.test.ts +++ b/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 | 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'); + }); }); diff --git a/packages/parallel-oauth/src/index.ts b/packages/parallel-oauth/src/index.ts index 60bec35..028aa29 100644 --- a/packages/parallel-oauth/src/index.ts +++ b/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((resolve) => { server.close(() => resolve()); + server.closeAllConnections(); }); }, }; diff --git a/packages/pi-extension/src/__tests__/parallel-auth.test.ts b/packages/pi-extension/src/__tests__/parallel-auth.test.ts index 6de427b..b25e1b7 100644 --- a/packages/pi-extension/src/__tests__/parallel-auth.test.ts +++ b/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', diff --git a/packages/pi-extension/src/parallel-auth.ts b/packages/pi-extension/src/parallel-auth.ts index d0abbd5..616514e 100644 --- a/packages/pi-extension/src/parallel-auth.ts +++ b/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) => {