internal: (studio) make the protocol database accessible to studio (#31443)

* internal: (studio) make the protocol database accessible to studio

* PR comment

Author: ryanthemanuel

packages/server/lib/cloud/protocol.ts

@@ -62,6 +62,10 @@ export class ProtocolManager implements ProtocolManagerShape {
     } : undefined
   }
 
+  get db () {
+    return this._db
+  }
+
   async prepareProtocol (script: string, options: ProtocolManagerOptions) {
     this._captureHash = base64url.fromBase64(crypto.createHash('SHA256').update(script).digest('base64'))
 

packages/server/lib/cloud/studio.ts

@@ -6,6 +6,7 @@ import os from 'os'
 import { agent } from '@packages/network'
 import Debug from 'debug'
 import { requireScript } from './require_script'
+import type Database from 'better-sqlite3'
 
 interface StudioServer { default: StudioServerDefaultShape }
 
@@ -37,6 +38,10 @@ export class StudioManager implements StudioManagerShape {
     return manager
   }
 
+  setProtocolDb (db: Database.Database): void {
+    this.invokeSync('setProtocolDb', { isEssential: true }, db)
+  }
+
   async setup ({ script, studioPath, studioHash, projectSlug, cloudApi }: SetupOptions): Promise<void> {
     const { createStudioServer } = requireScript<StudioServer>(script).default
 
@@ -100,6 +105,7 @@ export class StudioManager implements StudioManagerShape {
     }
 
     try {
+      // @ts-expect-error - TS not associating the method & args properly, even though we know it's correct
       return this._studioServer[method].apply(this._studioServer, args)
     } catch (error: unknown) {
       let actualError: Error

packages/server/lib/project-base.ts

@@ -447,6 +447,10 @@ export class ProjectBase extends EE {
 
           await browsers.connectProtocolToBrowser({ browser: this.browser, foundBrowsers: this.options.browsers, protocolManager: this.protocolManager })
 
+          if (this.protocolManager.db) {
+            this.ctx.coreData.studio?.setProtocolDb(this.protocolManager.db)
+          }
+
           return { canAccessStudioAI: true }
         }
 

packages/server/test/support/fixtures/cloud/protocol/test-protocol.ts

@@ -1,4 +1,4 @@
-import type { Database } from 'better-sqlite3'
+import type Database from 'better-sqlite3'
 import type { AppCaptureProtocolInterface, CDPClient, ResponseStreamOptions } from '@packages/types'
 import { Readable } from 'stream'
 
@@ -12,7 +12,7 @@ export class AppCaptureProtocol implements AppCaptureProtocolInterface {
   responseStreamReceived (options: ResponseStreamOptions): Readable {
     return Readable.from([])
   }
-  beforeSpec ({ workingDirectory, archivePath, dbPath, db }: { workingDirectory: string, archivePath: string, dbPath: string, db: Database }): void {}
+  beforeSpec ({ workingDirectory, archivePath, dbPath, db }: { workingDirectory: string, archivePath: string, dbPath: string, db: Database.Database }): void {}
   addRunnables (runnables: any): void {}
   commandLogAdded (log: any): void {}
   commandLogChanged (log: any): void {}

packages/server/test/support/fixtures/cloud/studio/test-studio.ts

@@ -1,4 +1,5 @@
 import type { StudioServerShape, StudioServerDefaultShape, StudioBrowser } from '@packages/types'
+import type Database from 'better-sqlite3'
 import type { Router } from 'express'
 
 class StudioServer implements StudioServerShape {
@@ -9,6 +10,9 @@ class StudioServer implements StudioServerShape {
   canAccessStudioAI (browser: StudioBrowser): Promise<boolean> {
     return Promise.resolve(true)
   }
+
+  setProtocolDb (db: Database.Database): void {
+  }
 }
 
 const studioServerDefault: StudioServerDefaultShape = {

packages/server/test/unit/cloud/protocol_spec.ts

@@ -350,6 +350,20 @@ describe('lib/cloud/protocol', () => {
     })
   })
 
+  describe('.db', () => {
+    it('returns the database instance', () => {
+      const mockDb = { test: 'db' }
+
+      protocolManager['_db'] = mockDb
+
+      expect(protocolManager.db).to.equal(mockDb)
+    })
+
+    it('returns undefined when no database is set', () => {
+      expect(protocolManager.db).to.be.undefined
+    })
+  })
+
   describe('.uploadCaptureArtifact()', () => {
     let filePath: string
     let fileSize: number

packages/server/test/unit/cloud/studio_spec.ts

@@ -160,9 +160,27 @@ describe('lib/cloud/studio', () => {
         name: 'chrome',
         family: 'chromium',
         channel: 'stable',
+        displayName: 'Chrome',
+        version: '120.0.0',
+        majorVersion: '120',
+        path: '/path/to/chrome',
+        isHeaded: true,
+        isHeadless: false,
       })
 
       expect(result).to.be.true
     })
   })
+
+  describe('setProtocolDb', () => {
+    it('sets the protocol database on the studio server', () => {
+      const mockDb = { test: 'db' }
+
+      sinon.stub(studio, 'setProtocolDb')
+
+      studioManager.setProtocolDb(mockDb as any)
+
+      expect(studio.setProtocolDb).to.be.calledWith(mockDb)
+    })
+  })
 })

packages/server/test/unit/project_spec.js

@@ -737,14 +737,17 @@ This option will not have an effect in Some-other-name. Tests that rely on web s
       const mockSetupProtocol = sinon.stub()
       const mockBeforeSpec = sinon.stub()
       const mockAccessStudioLLM = sinon.stub().resolves(true)
+      const mockSetProtocolDb = sinon.stub()
 
       this.project.spec = {}
       this.project.ctx.coreData.studio = {
         canAccessStudioAI: mockAccessStudioLLM,
         protocolManager: {
           setupProtocol: mockSetupProtocol,
           beforeSpec: mockBeforeSpec,
+          db: { test: 'db' },
         },
+        setProtocolDb: mockSetProtocolDb,
       }
 
       sinon.stub(browsers, 'connectProtocolToBrowser').resolves()
@@ -793,6 +796,7 @@ This option will not have an effect in Some-other-name. Tests that rely on web s
       })
 
       expect(this.project['_protocolManager']).to.eq(this.project.ctx.coreData.studio.protocolManager)
+      expect(mockSetProtocolDb).to.be.calledWith({ test: 'db' })
     })
 
     it('passes onStudioInit callback with AI enabled but no protocol manager', async function () {

packages/types/src/protocol.ts

@@ -1,4 +1,4 @@
-import type { Database } from 'better-sqlite3'
+import type Database from 'better-sqlite3'
 import type ProtocolMapping from 'devtools-protocol/types/protocol-mapping'
 import type { IncomingHttpHeaders } from 'http'
 import type { Readable } from 'stream'
@@ -38,7 +38,7 @@ export interface AppCaptureProtocolCommon {
 
 export interface AppCaptureProtocolInterface extends AppCaptureProtocolCommon {
   getDbMetadata (): { offset: number, size: number } | undefined
-  beforeSpec ({ spec, workingDirectory, archivePath, dbPath, db }: { spec: FoundSpec & { instanceId: string }, workingDirectory: string, archivePath: string, dbPath: string, db: Database }): void
+  beforeSpec ({ spec, workingDirectory, archivePath, dbPath, db }: { spec: FoundSpec & { instanceId: string }, workingDirectory: string, archivePath: string, dbPath: string, db: Database.Database }): void
   uploadStallSamplingInterval: () => number
   connectToBrowser (cdpClient: CDPClient): Promise<void>
 }
@@ -138,6 +138,7 @@ export interface ProtocolManagerShape extends AppCaptureProtocolCommon {
   uploadCaptureArtifact(artifact: CaptureArtifact): Promise<UploadCaptureArtifactResult | undefined>
   connectToBrowser (cdpClient: CDPClient): Promise<void>
   close (): void
+  db?: Database.Database
 }
 
 type Response = {

packages/types/src/studio/studio-server-types.ts

@@ -2,6 +2,7 @@
 
 import type { Router } from 'express'
 import type { AxiosInstance } from 'axios'
+import type Database from 'better-sqlite3'
 
 interface RetryOptions {
   maxAttempts: number
@@ -32,6 +33,7 @@ export interface StudioServerOptions {
 export interface StudioServerShape {
   initializeRoutes(router: Router): void
   canAccessStudioAI(browser: Cypress.Browser): Promise<boolean>
+  setProtocolDb(database: Database.Database): void
 }
 
 export interface StudioServerDefaultShape {