feat: Websocket communication between k6 Studio and browser (#382)

extension/src/background.ts

@@ -1,5 +1,80 @@
+import { MessageEnvelope, MessagePayload } from '@/services/browser/schemas'
 import { browser } from 'webextension-polyfill-ts'
 
+let socket: WebSocket | null = null
+let buffer: MessageEnvelope[] = []
+
+function send(message: MessagePayload) {
+  const envelope: MessageEnvelope = {
+    messageId: crypto.randomUUID(),
+    payload: message,
+  }
+
+  if (socket) {
+    socket.send(JSON.stringify(envelope))
+  } else {
+    buffer.push(envelope)
+  }
+}
+
+function flush(socket: WebSocket) {
+  for (const message of buffer) {
+    socket.send(JSON.stringify(message))
+  }
+
+  buffer = []
+}
+
+function reconnect() {
+  socket?.close()
+  socket = null
+
+  setTimeout(() => {
+    connect()
+  }, 1000)
+}
+
+function connect() {
+  const ws = new WebSocket('ws://localhost:7554')
+
+  ws.onopen = () => {
+    console.log('Connected to server...')
+
+    socket = ws
+
+    flush(ws)
+  }
+
+  ws.onerror = (err) => {
+    console.log('Connection error...', err)
+
+    reconnect()
+  }
+
+  ws.onclose = () => {
+    console.log('Connection closed...')
+
+    reconnect()
+  }
+}
+
+setInterval(() => {
+  send({
+    type: 'events-captured',
+    events: [
+      {
+        eventId: crypto.randomUUID(),
+        timestamp: Date.now(),
+        type: 'dummy',
+        selector: 'button',
+        message: 'Clicked button',
+      },
+    ],
+  })
+}, 5000)
+
+connect()
+
 browser.runtime.onInstalled.addListener(() => {
   console.log('Extension installed...')
 })

package.json

@@ -38,6 +38,7 @@
     "@types/node-forge": "^1.3.11",
     "@types/react": "^18.3.3",
     "@types/react-dom": "^18.3.0",
+    "@types/ws": "^8.5.13",
     "@typescript-eslint/eslint-plugin": "^7.16.1",
     "@typescript-eslint/parser": "^7.16.1",
     "electron": "30.0.8",
@@ -106,6 +107,7 @@
     "tree-kill": "^1.2.2",
     "update-electron-app": "^3.0.0",
     "webextension-polyfill-ts": "^0.26.0",
+    "ws": "^8.18.0",
     "zod": "^3.23.8",
     "zustand": "^4.5.3"
   },

src/browser.ts

@@ -10,6 +10,7 @@ import { mkdtemp } from 'fs/promises'
 import path from 'path'
 import os from 'os'
 import { appSettings } from './main'
+import { launchBrowserServer } from './services/browser/server'
 
 const createUserDataDir = async () => {
   return mkdtemp(path.join(os.tmpdir(), 'k6-studio-'))
@@ -59,7 +60,13 @@ export const launchBrowser = async (
   const extensionPath = getExtensionPath()
   console.info(`extension path: ${extensionPath}`)
 
-  const sendBrowserClosedEvent = (): Promise<void> => {
+  const disposeWebSockerServer = appSettings.recorder.enableBrowserRecorder
+    ? launchBrowserServer(browserWindow)
+    : () => {}
+
+  const handleBrowserClose = (): Promise<void> => {
+    disposeWebSockerServer()
+
     // we send the browser:stopped event when the browser is closed
     // NOTE: on macos pressing the X button does not close the application so it won't be fired
     browserWindow.webContents.send('browser:closed')
@@ -87,6 +94,6 @@ export const launchBrowser = async (
       disableChromeOptimizations,
       url ?? '',
     ],
-    onExit: sendBrowserClosedEvent,
+    onExit: handleBrowserClose,
   })
 }

src/hooks/useListenBrowserEvent.ts

@@ -0,0 +1,16 @@
+import { BrowserEvent } from '@/schemas/recording'
+import { useEffect, useState } from 'react'
+
+export function useListenBrowserEvent() {
+  const [events, setEvents] = useState<BrowserEvent[]>([])
+
+  useEffect(() => {
+    return window.studio.browser.onBrowserEvent((events: BrowserEvent[]) => {
+      console.log('Received browser events', events)
+
+      setEvents((prevEvents) => [...prevEvents, ...events])
+    })
+  }, [])
+
+  return events
+}

src/preload.ts

@@ -4,6 +4,7 @@ import { HarFile } from './types/har'
 import { GeneratorFile } from './types/generator'
 import { AddToastPayload } from './types/toast'
 import { AppSettings } from './types/settings'
+import { BrowserEvent } from './schemas/recording'
 
 interface GetFilesResponse {
   recordings: string[]
@@ -55,6 +56,9 @@ const browser = {
   openExternalLink: (url: string) => {
     return ipcRenderer.invoke('browser:open:external:link', url)
   },
+  onBrowserEvent: (callback: (event: BrowserEvent[]) => void) => {
+    return createListener('browser:event', callback)
+  },
 } as const
 
 const script = {

src/schemas/recording/index.ts

@@ -0,0 +1 @@
+export * from './v1/browser'

src/schemas/recording/v1/browser.ts

@@ -0,0 +1,14 @@
+import { z } from 'zod'
+
+const DummyEvent = z.object({
+  eventId: z.string(),
+  timestamp: z.number(),
+  type: z.literal('dummy'),
+  selector: z.string(),
+  message: z.string(),
+})
+
+export const BrowserEventSchema = DummyEvent
+
+export type DummyEvent = z.infer<typeof DummyEvent>
+export type BrowserEvent = z.infer<typeof BrowserEventSchema>

src/services/browser/schemas.ts

@@ -0,0 +1,18 @@
+import { BrowserEventSchema } from '@/schemas/recording'
+import { z } from 'zod'
+
+export const EventsCapturedMessage = z.object({
+  type: z.literal('events-captured'),
+  events: z.array(BrowserEventSchema),
+})
+
+export const MessagePayload = EventsCapturedMessage
+
+export const MessageEnvelope = z.object({
+  messageId: z.string(),
+  payload: MessagePayload,
+})
+
+export type EventsCapturedMessage = z.infer<typeof EventsCapturedMessage>
+export type MessagePayload = z.infer<typeof MessagePayload>
+export type MessageEnvelope = z.infer<typeof MessageEnvelope>