add sse test to make sure responses are not buffered

playwriter/src/mcp.test.ts

@@ -12,13 +12,15 @@ import { getCDPSessionForPage } from './cdp-session.js'
 import { Debugger } from './debugger.js'
 import { Editor } from './editor.js'
 import type { CDPCommand } from './cdp-types.js'
-import { setupTestContext, cleanupTestContext, getExtensionServiceWorker, type TestContext } from './test-utils.js'
+import { setupTestContext, cleanupTestContext, getExtensionServiceWorker, createSseServer, type TestContext, withTimeout } from './test-utils.js'
+import { screenshotWithAccessibilityLabels } from './aria-snapshot.js'
 
 declare const window: any
 declare const document: any
 
 const TEST_PORT = 19987
 
+
 function js(strings: TemplateStringsArray, ...values: any[]): string {
     return strings.reduce(
         (result, str, i) => result + str + (values[i] || ''),
@@ -26,6 +28,7 @@ function js(strings: TemplateStringsArray, ...values: any[]): string {
     )
 }
 
+
 declare global {
     var toggleExtensionForActiveTab: () => Promise<{ isConnected: boolean; state: ExtensionState }>;
     var getExtensionState: () => ExtensionState;
@@ -2347,7 +2350,7 @@ describe('MCP Server Tests', () => {
         // Verify the image is valid JPEG by checking base64
         const buffer = Buffer.from(imageContent.data, 'base64')
         const dimensions = imageSize(buffer)
-        
+
         // Get actual viewport size from page
         const viewport = await page.evaluate(() => ({
             innerWidth: window.innerWidth,
@@ -2357,7 +2360,7 @@ describe('MCP Server Tests', () => {
         }))
         console.log('Screenshot dimensions:', dimensions.width, 'x', dimensions.height)
         console.log('Window viewport:', viewport)
-        
+
         expect(dimensions.type).toBe('jpg')
         expect(dimensions.width).toBeGreaterThan(0)
         expect(dimensions.height).toBeGreaterThan(0)
@@ -2365,6 +2368,7 @@ describe('MCP Server Tests', () => {
         await page.close()
     }, 60000)
 
+
     it('should get clean HTML with getCleanHTML', async () => {
         const browserContext = getBrowserContext()
         const serviceWorker = await getExtensionServiceWorker(browserContext)
@@ -2573,7 +2577,7 @@ describe('MCP Server Tests', () => {
         const listJson = await listRes.json() as Array<{ url?: string }>
         expect(Array.isArray(listJson)).toBe(true)
         expect(listJson.length).toBeGreaterThan(0)
-        
+
         // Find the example.com page (there may be other pages like about:blank)
         const examplePage = listJson.find((t) => t.url?.includes('example.com'))
         expect(examplePage).toBeDefined()
@@ -3651,6 +3655,114 @@ describe('Service Worker Target Tests', () => {
         await browser.close()
         await page.close()
     }, 60000)
+
+    it('should stream SSE without waiting for response end', async () => {
+        const browserContext = getBrowserContext()
+        const serviceWorker = await withTimeout({
+            promise: getExtensionServiceWorker(browserContext),
+            timeoutMs: 5000,
+            errorMessage: 'getExtensionServiceWorker timed out',
+        })
+        const sseServer = await withTimeout({
+            promise: createSseServer(),
+            timeoutMs: 5000,
+            errorMessage: 'createSseServer timed out',
+        })
+        let page: Awaited<ReturnType<typeof browserContext.newPage>> | null = null
+        let browser: Awaited<ReturnType<typeof chromium.connectOverCDP>> | null = null
+
+        try {
+            page = await withTimeout({
+                promise: browserContext.newPage(),
+                timeoutMs: 5000,
+                errorMessage: 'newPage timed out',
+            })
+            await withTimeout({
+                promise: page.goto(`${sseServer.baseUrl}/`),
+                timeoutMs: 5000,
+                errorMessage: 'page.goto timed out',
+            })
+            await page.bringToFront()
+
+            await withTimeout({
+                promise: serviceWorker.evaluate(async () => {
+                    await globalThis.toggleExtensionForActiveTab()
+                }),
+                timeoutMs: 5000,
+                errorMessage: 'toggleExtensionForActiveTab timed out',
+            })
+            await new Promise((resolve) => {
+                setTimeout(resolve, 100)
+            })
+
+            browser = await withTimeout({
+                promise: chromium.connectOverCDP(getCdpUrl({ port: TEST_PORT })),
+                timeoutMs: 5000,
+                errorMessage: 'connectOverCDP timed out',
+            })
+            const cdpPage = browser.contexts()[0].pages().find((p) => {
+                return p.url().startsWith(sseServer.baseUrl)
+            })
+            expect(cdpPage).toBeDefined()
+
+            await cdpPage!.evaluate(() => {
+                return window.startSse()
+            })
+            await withTimeout({
+                promise: cdpPage!.waitForFunction(() => {
+                    return window.__sseMessages.length > 0
+                }, { timeout: 5000 }),
+                timeoutMs: 7000,
+                errorMessage: 'SSE message not received in time',
+            })
+
+            const firstMessage = await cdpPage!.evaluate(() => {
+                return window.__sseMessages[0]
+            })
+            expect(firstMessage).toBe('hello')
+
+            const sseState = sseServer.getState()
+            expect(sseState.connected).toBe(true)
+            expect(sseState.finished).toBe(false)
+            expect(sseState.closed).toBe(false)
+            expect(sseState.writeCount).toBeGreaterThan(0)
+
+            const readyState = await cdpPage!.evaluate(() => {
+                if (!window.__sseSource) {
+                    return -1
+                }
+                return window.__sseSource.readyState
+            })
+            expect(readyState).toBe(1)
+
+            await cdpPage!.evaluate(() => {
+                window.stopSse()
+            })
+            await new Promise((resolve) => {
+                setTimeout(resolve, 100)
+            })
+        } finally {
+            if (browser) {
+                await withTimeout({
+                    promise: browser.close(),
+                    timeoutMs: 5000,
+                    errorMessage: 'browser.close timed out',
+                })
+            }
+            if (page) {
+                await withTimeout({
+                    promise: page.close(),
+                    timeoutMs: 5000,
+                    errorMessage: 'page.close timed out',
+                })
+            }
+            await withTimeout({
+                promise: sseServer.close(),
+                timeoutMs: 5000,
+                errorMessage: 'sseServer.close timed out',
+            })
+        }
+    }, 60000)
 })
 
 describe('Auto-enable Tests', () => {

playwriter/src/test-utils.ts

@@ -1,5 +1,7 @@
 import { exec } from 'node:child_process'
 import { promisify } from 'node:util'
+import http from 'node:http'
+import net from 'node:net'
 import { chromium, BrowserContext } from 'playwright-core'
 import path from 'node:path'
 import fs from 'node:fs'
@@ -101,3 +103,167 @@ export async function cleanupTestContext(ctx: TestContext | null, cleanup?: (()
     await cleanup()
   }
 }
+
+export type SseServerState = {
+  connected: boolean
+  finished: boolean
+  writeCount: number
+  closed: boolean
+}
+
+export type SseServer = {
+  baseUrl: string
+  getState: () => SseServerState
+  close: () => Promise<void>
+}
+
+export async function createSseServer(): Promise<SseServer> {
+  let sseResponse: http.ServerResponse | null = null
+  let sseFinished = false
+  let sseClosed = false
+  let sseWriteCount = 0
+  let sseInterval: NodeJS.Timeout | null = null
+  const openResponses: Set<http.ServerResponse> = new Set()
+  const openSockets: Set<net.Socket> = new Set()
+
+  const server = http.createServer((req, res) => {
+    if (req.url === '/') {
+      res.writeHead(200, { 'Content-Type': 'text/html' })
+      res.end(`<!doctype html>
+<html>
+  <head>
+    <meta charset="utf-8" />
+    <title>SSE Test</title>
+  </head>
+  <body>
+    <script>
+      window.__sseMessages = [];
+      window.__sseOpen = false;
+      window.__sseError = null;
+      window.startSse = function () {
+        const source = new EventSource('/sse');
+        window.__sseSource = source;
+        source.onopen = function () {
+          window.__sseOpen = true;
+        };
+        source.onmessage = function (event) {
+          window.__sseMessages.push(event.data);
+        };
+        source.onerror = function () {
+          window.__sseError = 'SSE error';
+        };
+        return true;
+      };
+      window.stopSse = function () {
+        if (window.__sseSource) {
+          window.__sseSource.close();
+        }
+      };
+    </script>
+  </body>
+</html>`)
+      return
+    }
+
+    if (req.url === '/sse') {
+      res.writeHead(200, {
+        'Content-Type': 'text/event-stream',
+        'Cache-Control': 'no-cache, no-transform',
+        'Connection': 'keep-alive',
+      })
+      res.write('retry: 1000\n\n')
+      res.write('data: hello\n\n')
+      sseResponse = res
+      sseWriteCount += 1
+      openResponses.add(res)
+
+      res.on('finish', () => {
+        sseFinished = true
+      })
+      res.on('close', () => {
+        sseClosed = true
+        openResponses.delete(res)
+        if (sseInterval) {
+          clearInterval(sseInterval)
+          sseInterval = null
+        }
+      })
+
+      sseInterval = setInterval(() => {
+        res.write('data: ping\n\n')
+        sseWriteCount += 1
+      }, 200)
+      return
+    }
+
+    res.writeHead(404)
+    res.end('Not found')
+  })
+
+  server.on('connection', (socket) => {
+    openSockets.add(socket)
+    socket.on('close', () => {
+      openSockets.delete(socket)
+    })
+  })
+
+  await new Promise<void>((resolve) => {
+    server.listen(0, '127.0.0.1', () => {
+      resolve()
+    })
+  })
+
+  const address = server.address()
+  if (!address || typeof address === 'string') {
+    throw new Error('Failed to bind SSE server')
+  }
+
+  return {
+    baseUrl: `http://127.0.0.1:${address.port}`,
+    getState: () => ({
+      connected: sseResponse !== null,
+      finished: sseFinished,
+      closed: sseClosed,
+      writeCount: sseWriteCount,
+    }),
+    close: async () => {
+      for (const response of openResponses) {
+        response.destroy()
+      }
+      for (const socket of openSockets) {
+        socket.destroy()
+      }
+      if (sseInterval) {
+        clearInterval(sseInterval)
+        sseInterval = null
+      }
+      await new Promise<void>((resolve, reject) => {
+        server.close((error) => {
+          if (error) {
+            reject(error)
+            return
+          }
+          resolve()
+        })
+      })
+    }
+  }
+}
+
+export async function withTimeout<T>({ promise, timeoutMs, errorMessage }: { promise: Promise<T>; timeoutMs: number; errorMessage: string }): Promise<T> {
+    return await new Promise<T>((resolve, reject) => {
+        const timeoutId = setTimeout(() => {
+            reject(new Error(errorMessage))
+        }, timeoutMs)
+
+        promise
+            .then((value) => {
+                clearTimeout(timeoutId)
+                resolve(value)
+            })
+            .catch((error) => {
+                clearTimeout(timeoutId)
+                reject(error)
+            })
+    })
+}