internal: (studio) enable studio to access the protocol database with a new connection (#31502)

* internal: (studio) improvements to how studio can access the protocol database

* fix

* fix build

* refactor due to downstream changes

* Update packages/app/src/studio/studio-app-types.ts

* Update packages/server/lib/cloud/studio.ts

Author: ryanthemanuel

.circleci/workflows.yml

@@ -38,7 +38,7 @@ mainBuildFilters: &mainBuildFilters
         - /^release\/\d+\.\d+\.\d+$/
         # use the following branch as well to ensure that v8 snapshot cache updates are fully tested
         - 'update-v8-snapshot-cache-on-develop'
-        - 'chore/update_mobx_decoratorless'
+        - 'ryanm/chore/full-snapshot-threaded'
 
 # usually we don't build Mac app - it takes a long time
 # but sometimes we want to really confirm we are doing the right thing
@@ -49,7 +49,7 @@ macWorkflowFilters: &darwin-workflow-filters
       - equal: [ develop, << pipeline.git.branch >> ]
       # use the following branch as well to ensure that v8 snapshot cache updates are fully tested
       - equal: [ 'update-v8-snapshot-cache-on-develop', << pipeline.git.branch >> ]
-      - equal: [ 'chore/update_mobx_decoratorless', << pipeline.git.branch >> ]
+      - equal: [ 'ryanm/chore/full-snapshot-threaded', << pipeline.git.branch >> ]
       - matches:
           pattern: /^release\/\d+\.\d+\.\d+$/
           value: << pipeline.git.branch >>
@@ -60,7 +60,7 @@ linuxArm64WorkflowFilters: &linux-arm64-workflow-filters
       - equal: [ develop, << pipeline.git.branch >> ]
       # use the following branch as well to ensure that v8 snapshot cache updates are fully tested
       - equal: [ 'update-v8-snapshot-cache-on-develop', << pipeline.git.branch >> ]
-      - equal: [ 'chore/update_mobx_decoratorless', << pipeline.git.branch >> ]
+      - equal: [ 'ryanm/chore/full-snapshot-threaded', << pipeline.git.branch >> ]
       - matches:
           pattern: /^release\/\d+\.\d+\.\d+$/
           value: << pipeline.git.branch >>
@@ -83,7 +83,7 @@ windowsWorkflowFilters: &windows-workflow-filters
       - equal: [ develop, << pipeline.git.branch >> ]
       # use the following branch as well to ensure that v8 snapshot cache updates are fully tested
       - equal: [ 'update-v8-snapshot-cache-on-develop', << pipeline.git.branch >> ]
-      - equal: [ 'chore/update_mobx_decoratorless', << pipeline.git.branch >> ]
+      - equal: [ 'ryanm/chore/full-snapshot-threaded', << pipeline.git.branch >> ]
       - matches:
           pattern: /^release\/\d+\.\d+\.\d+$/
           value: << pipeline.git.branch >>
@@ -157,7 +157,7 @@ commands:
           name: Set environment variable to determine whether or not to persist artifacts
           command: |
             echo "Setting SHOULD_PERSIST_ARTIFACTS variable"
-            echo 'if ! [[ "$CIRCLE_BRANCH" != "develop" && "$CIRCLE_BRANCH" != "release/"* && "$CIRCLE_BRANCH" != "chore/update_mobx_decoratorless" ]]; then
+            echo 'if ! [[ "$CIRCLE_BRANCH" != "develop" && "$CIRCLE_BRANCH" != "release/"* && "$CIRCLE_BRANCH" != "ryanm/chore/full-snapshot-threaded" ]]; then
                 export SHOULD_PERSIST_ARTIFACTS=true
             fi' >> "$BASH_ENV"
   # You must run `setup_should_persist_artifacts` command and be using bash before running this command

packages/server/lib/cloud/protocol.ts

@@ -62,8 +62,8 @@ export class ProtocolManager implements ProtocolManagerShape {
     } : undefined
   }
 
-  get db () {
-    return this._db
+  get dbPath () {
+    return this._dbPath
   }
 
   async prepareProtocol (script: string, options: ProtocolManagerOptions) {

packages/server/lib/cloud/studio.ts

@@ -1,4 +1,4 @@
-import type { StudioErrorReport, StudioManagerShape, StudioStatus, StudioServerDefaultShape, StudioServerShape, ProtocolManagerShape, StudioCloudApi } from '@packages/types'
+import type { StudioErrorReport, StudioManagerShape, StudioStatus, StudioServerDefaultShape, StudioServerShape, ProtocolManagerShape, StudioCloudApi, StudioAIInitializeOptions } from '@packages/types'
 import type { Router } from 'express'
 import type { Socket } from 'socket.io'
 import fetch from 'cross-fetch'
@@ -7,7 +7,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'
+import path from 'path'
 
 interface StudioServer { default: StudioServerDefaultShape }
 
@@ -39,17 +39,14 @@ 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
 
     this._studioServer = await createStudioServer({
       studioPath,
       projectSlug,
       cloudApi,
+      betterSqlite3Path: path.dirname(require.resolve('better-sqlite3/package.json')),
     })
 
     this._studioHash = studioHash
@@ -72,6 +69,14 @@ export class StudioManager implements StudioManagerShape {
     return (await this.invokeAsync('canAccessStudioAI', { isEssential: true }, browser)) ?? false
   }
 
+  async initializeStudioAI (options: StudioAIInitializeOptions): Promise<void> {
+    await this.invokeAsync('initializeStudioAI', { isEssential: true }, options)
+  }
+
+  async destroy (): Promise<void> {
+    await this.invokeAsync('destroy', { isEssential: true })
+  }
+
   private async reportError (error: Error): Promise<void> {
     try {
       const payload: StudioErrorReport = {

packages/server/lib/project-base.ts

@@ -432,25 +432,32 @@ export class ProjectBase extends EE {
 
       onStudioInit: async () => {
         if (this.spec && this.ctx.coreData.studio?.protocolManager) {
-          const canAccessStudioAI = await this.ctx.coreData.studio?.canAccessStudioAI(this.browser) ?? false
+          const canAccessStudioAI = await this.ctx.coreData.studio.canAccessStudioAI(this.browser) ?? false
 
           if (!canAccessStudioAI) {
             return { canAccessStudioAI }
           }
 
-          this.protocolManager = this.ctx.coreData.studio?.protocolManager
-          this.protocolManager?.setupProtocol()
-          this.protocolManager?.beforeSpec({
+          this.ctx.coreData.studio.protocolManager.setupProtocol()
+          this.ctx.coreData.studio.protocolManager.beforeSpec({
             ...this.spec,
             instanceId: v4(),
           })
 
-          await browsers.connectProtocolToBrowser({ browser: this.browser, foundBrowsers: this.options.browsers, protocolManager: this.protocolManager })
+          await browsers.connectProtocolToBrowser({ browser: this.browser, foundBrowsers: this.options.browsers, protocolManager: this.ctx.coreData.studio.protocolManager })
 
-          if (this.protocolManager.db) {
-            this.ctx.coreData.studio?.setProtocolDb(this.protocolManager.db)
+          if (!this.ctx.coreData.studio.protocolManager.dbPath) {
+            debug('Protocol database path is not set after initializing protocol manager')
+
+            return { canAccessStudioAI: false }
           }
 
+          this.protocolManager = this.ctx.coreData.studio.protocolManager
+
+          await this.ctx.coreData.studio.initializeStudioAI({
+            protocolDbPath: this.ctx.coreData.studio.protocolManager.dbPath,
+          })
+
           return { canAccessStudioAI: true }
         }
 
@@ -464,6 +471,7 @@ export class ProjectBase extends EE {
           await browsers.closeProtocolConnection({ browser: this.browser, foundBrowsers: this.options.browsers })
           this.protocolManager?.close()
           this.protocolManager = undefined
+          await this.ctx.coreData.studio.destroy()
         }
       },
 

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

@@ -1,7 +1,6 @@
 /// <reference types="cypress" />
 
 import type { StudioServerShape, StudioServerDefaultShape } from '@packages/types'
-import type Database from 'better-sqlite3'
 import type { Router } from 'express'
 import type { Socket } from '@packages/socket'
 
@@ -14,11 +13,15 @@ class StudioServer implements StudioServerShape {
     return Promise.resolve(true)
   }
 
-  addSocketListeners (socket: Socket): void {
-    // This is a test implementation that does nothing
+  initializeStudioAI (): Promise<void> {
+    return Promise.resolve()
   }
 
-  setProtocolDb (db: Database.Database): void {
+  destroy (): Promise<void> {
+    return Promise.resolve()
+  }
+  
+  addSocketListeners (socket: Socket): void {
     // This is a test implementation that does nothing
   }
 }

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

@@ -350,17 +350,17 @@ describe('lib/cloud/protocol', () => {
     })
   })
 
-  describe('.db', () => {
-    it('returns the database instance', () => {
-      const mockDb = { test: 'db' }
+  describe('.dbPath', () => {
+    it('returns the database path', () => {
+      const mockDbPath = '/path/to/db'
 
-      protocolManager['_db'] = mockDb
+      protocolManager['_dbPath'] = mockDbPath
 
-      expect(protocolManager.db).to.equal(mockDb)
+      expect(protocolManager.dbPath).to.equal(mockDbPath)
     })
 
-    it('returns undefined when no database is set', () => {
-      expect(protocolManager.db).to.be.undefined
+    it('returns undefined when no database path is set', () => {
+      expect(protocolManager.dbPath).to.be.undefined
     })
   })
 

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

@@ -197,15 +197,27 @@ describe('lib/cloud/studio', () => {
     })
   })
 
-  describe('setProtocolDb', () => {
-    it('sets the protocol database on the studio server', () => {
-      const mockDb = { test: 'db' }
+  describe('initializeStudioAI', () => {
+    it('initializes Studio AI on the studio server', async () => {
+      sinon.stub(studio, 'initializeStudioAI').resolves()
 
-      sinon.stub(studio, 'setProtocolDb')
+      await studioManager.initializeStudioAI({
+        protocolDbPath: 'test-db-path',
+      })
+
+      expect(studio.initializeStudioAI).to.be.calledWith({
+        protocolDbPath: 'test-db-path',
+      })
+    })
+  })
+
+  describe('destroy', () => {
+    it('destroys the studio server', async () => {
+      sinon.stub(studio, 'destroy').resolves()
 
-      studioManager.setProtocolDb(mockDb as any)
+      await studioManager.destroy()
 
-      expect(studio.setProtocolDb).to.be.calledWith(mockDb)
+      expect(studio.destroy).to.be.called
     })
   })
 })

packages/server/test/unit/project_spec.js

@@ -43,6 +43,7 @@ describe('lib/project-base', () => {
     this.testStudioManager = {
       initializeRoutes: () => {},
       status: 'INITIALIZED',
+      destroy: () => Promise.resolve(),
     }
 
     sinon.stub(studio, 'getAndInitializeStudioManager').resolves(this.testStudioManager)
@@ -736,18 +737,20 @@ This option will not have an effect in Some-other-name. Tests that rely on web s
     it('passes onStudioInit callback with AI enabled and a protocol manager', async function () {
       const mockSetupProtocol = sinon.stub()
       const mockBeforeSpec = sinon.stub()
-      const mockAccessStudioLLM = sinon.stub().resolves(true)
-      const mockSetProtocolDb = sinon.stub()
+      const mockAccessStudioAI = sinon.stub().resolves(true)
+      const mockSetProtocolDbPath = sinon.stub()
+      const mockInitializeStudioAI = sinon.stub().resolves()
 
       this.project.spec = {}
       this.project.ctx.coreData.studio = {
-        canAccessStudioAI: mockAccessStudioLLM,
+        canAccessStudioAI: mockAccessStudioAI,
         protocolManager: {
           setupProtocol: mockSetupProtocol,
           beforeSpec: mockBeforeSpec,
-          db: { test: 'db' },
+          dbPath: 'test-db-path',
         },
-        setProtocolDb: mockSetProtocolDb,
+        setProtocolDbPath: mockSetProtocolDbPath,
+        initializeStudioAI: mockInitializeStudioAI,
       }
 
       sinon.stub(browsers, 'connectProtocolToBrowser').resolves()
@@ -783,7 +786,7 @@ This option will not have an effect in Some-other-name. Tests that rely on web s
 
       expect(mockSetupProtocol).to.be.calledOnce
       expect(mockBeforeSpec).to.be.calledOnce
-      expect(mockAccessStudioLLM).to.be.calledWith({
+      expect(mockAccessStudioAI).to.be.calledWith({
         family: 'chromium',
         name: 'chrome',
         channel: 'stable',
@@ -796,17 +799,66 @@ 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' })
+      expect(mockInitializeStudioAI).to.be.calledWith({
+        protocolDbPath: 'test-db-path',
+      })
+    })
+
+    it('handles case where protocol manager db path is not set', async function () {
+      const mockSetupProtocol = sinon.stub()
+      const mockBeforeSpec = sinon.stub()
+      const mockAccessStudioAI = sinon.stub().resolves(true)
+      const mockSetProtocolDbPath = sinon.stub()
+      const mockInitializeStudioAI = sinon.stub().resolves()
+
+      this.project.spec = {}
+      this.project.ctx.coreData.studio = {
+        canAccessStudioAI: mockAccessStudioAI,
+        protocolManager: {
+          setupProtocol: mockSetupProtocol,
+          beforeSpec: mockBeforeSpec,
+          dbPath: null,
+        },
+        setProtocolDbPath: mockSetProtocolDbPath,
+        initializeStudioAI: mockInitializeStudioAI,
+      }
+
+      sinon.stub(browsers, 'connectProtocolToBrowser').resolves()
+
+      this.project.browser = {
+        name: 'chrome',
+        family: 'chromium',
+        channel: 'stable',
+      }
+
+      this.project.options.browsers = [{
+        name: 'chrome',
+        family: 'chromium',
+        channel: 'stable',
+      }]
+
+      let studioInitPromise
+
+      this.project.server.startWebsockets.callsFake(async (automation, config, callbacks) => {
+        studioInitPromise = callbacks.onStudioInit()
+      })
+
+      this.project.startWebsockets({}, {})
+
+      const { canAccessStudioAI } = await studioInitPromise
+
+      expect(canAccessStudioAI).to.be.false
+      expect(this.project['_protocolManager']).to.be.undefined
     })
 
     it('passes onStudioInit callback with AI enabled but no protocol manager', async function () {
       const mockSetupProtocol = sinon.stub()
       const mockBeforeSpec = sinon.stub()
-      const mockAccessStudioLLM = sinon.stub().resolves(true)
+      const mockAccessStudioAI = sinon.stub().resolves(true)
 
       this.project.spec = {}
       this.project.ctx.coreData.studio = {
-        canAccessStudioAI: mockAccessStudioLLM,
+        canAccessStudioAI: mockAccessStudioAI,
       }
 
       this.project.browser = {
@@ -836,20 +888,20 @@ This option will not have an effect in Some-other-name. Tests that rely on web s
 
       expect(mockSetupProtocol).not.to.be.called
       expect(mockBeforeSpec).not.to.be.called
-      expect(mockAccessStudioLLM).not.to.be.called
+      expect(mockAccessStudioAI).not.to.be.called
 
       expect(browsers.connectProtocolToBrowser).not.to.be.called
       expect(this.project['_protocolManager']).to.be.undefined
     })
 
-    it('passes onStudioInit callback with llm disabled', async function () {
+    it('passes onStudioInit callback with AI disabled', async function () {
       const mockSetupProtocol = sinon.stub()
       const mockBeforeSpec = sinon.stub()
-      const mockAccessStudioLLM = sinon.stub().resolves(false)
+      const mockAccessStudioAI = sinon.stub().resolves(false)
 
       this.project.spec = {}
       this.project.ctx.coreData.studio = {
-        canAccessStudioAI: mockAccessStudioLLM,
+        canAccessStudioAI: mockAccessStudioAI,
         protocolManager: {
           setupProtocol: mockSetupProtocol,
           beforeSpec: mockBeforeSpec,
@@ -887,9 +939,11 @@ This option will not have an effect in Some-other-name. Tests that rely on web s
 
     it('passes onStudioDestroy callback', async function () {
       const mockClose = sinon.stub()
+      const mockDestroy = sinon.stub().resolves()
 
       this.project.ctx.coreData.studio = {
         protocolManager: {},
+        destroy: mockDestroy,
       }
 
       sinon.stub(browsers, 'closeProtocolConnection').resolves()
@@ -928,7 +982,7 @@ This option will not have an effect in Some-other-name. Tests that rely on web s
       })
 
       expect(mockClose).to.be.calledOnce
-
+      expect(mockDestroy).to.be.calledOnce
       expect(this.project['_protocolManager']).to.be.undefined
     })
   })