internal: rework studio handshake to allow better caching

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

@@ -1,8 +1,10 @@
 export interface StudioPanelProps {
   canAccessStudioAI: boolean
-  onStudioPanelClose: () => void
-  useStudioEventManager?: StudioEventManagerShape
+  onStudioPanelClose?: () => void
+  useRunnerStatus?: RunnerStatusShape
+  useTestContentRetriever?: TestContentRetrieverShape
   useStudioAIStream?: StudioAIStreamShape
+  useCypress?: CypressShape
 }
 
 export type StudioPanelShape = (props: StudioPanelProps) => JSX.Element
@@ -18,21 +20,49 @@ CyEventEmitter & {
   state: (key: string) => any
 }
 
-export interface StudioEventManagerProps {
-  Cypress: CypressInternal
+export interface TestBlock {
+  content: string
+  testBodyPosition: {
+    contentStart: number
+    contentEnd: number
+    indentation: number
+  }
 }
 
 export type RunnerStatus = 'running' | 'finished'
 
-export type StudioEventManagerShape = (props: StudioEventManagerProps) => {
+export interface RunnerStatusProps {
+  Cypress: CypressInternal
+}
+
+export interface CypressProps {
+  Cypress: CypressInternal
+}
+
+export type CypressShape = (props: CypressProps) => {
+  currentCypress: CypressInternal
+}
+
+export type RunnerStatusShape = (props: RunnerStatusProps) => {
   runnerStatus: RunnerStatus
-  testBlock: string | null
 }
 
 export interface StudioAIStreamProps {
   canAccessStudioAI: boolean
   AIOutputRef: { current: HTMLTextAreaElement | null }
   runnerStatus: RunnerStatus
+  testCode?: string
+  isCreatingNewTest: boolean
 }
 
 export type StudioAIStreamShape = (props: StudioAIStreamProps) => void
+
+export interface TestContentRetrieverProps {
+  Cypress: CypressInternal
+}
+
+export type TestContentRetrieverShape = (props: TestContentRetrieverProps) => {
+  isLoading: boolean
+  testBlock: TestBlock | null
+  isCreatingNewTest: boolean
+}

packages/server/lib/StudioLifecycleManager.ts

@@ -11,6 +11,8 @@ import { reportStudioError } from './cloud/api/studio/report_studio_error'
 import { CloudRequest } from './cloud/api/cloud_request'
 import { isRetryableError } from './cloud/network/is_retryable_error'
 import { asyncRetry } from './util/async_retry'
+import { postStudioSession } from './cloud/api/studio/post_studio_session'
+
 const debug = Debug('cypress:server:studio-lifecycle-manager')
 const routes = require('./cloud/routes')
 
@@ -42,41 +44,11 @@ export class StudioLifecycleManager {
   }): void {
     debug('Initializing studio manager')
 
-    const studioManagerPromise = getAndInitializeStudioManager({
+    const studioManagerPromise = this.createStudioManager({
       projectId,
       cloudDataSource,
-    }).then(async (studioManager) => {
-      if (studioManager.status === 'ENABLED') {
-        debug('Cloud studio is enabled - setting up protocol')
-        const protocolManager = new ProtocolManager()
-        const protocolUrl = routes.apiRoutes.captureProtocolCurrent()
-        const script = await api.getCaptureProtocolScript(protocolUrl)
-
-        await protocolManager.prepareProtocol(script, {
-          runId: 'studio',
-          projectId: cfg.projectId,
-          testingType: cfg.testingType,
-          cloudApi: {
-            url: routes.apiUrl,
-            retryWithBackoff: api.retryWithBackoff,
-            requestPromise: api.rp,
-          },
-          projectConfig: _.pick(cfg, ['devServerPublicPathRoute', 'port', 'proxyUrl', 'namespace']),
-          mountVersion: api.runnerCapabilities.protocolMountVersion,
-          debugData,
-          mode: 'studio',
-        })
-
-        studioManager.protocolManager = protocolManager
-      } else {
-        debug('Cloud studio is not enabled - skipping protocol setup')
-      }
-
-      debug('Studio is ready')
-      this.studioManager = studioManager
-      this.callRegisteredListeners()
-
-      return studioManager
+      cfg,
+      debugData,
     }).catch(async (error) => {
       debug('Error during studio manager setup: %o', error)
 
@@ -125,6 +97,59 @@ export class StudioLifecycleManager {
     return await this.studioManagerPromise
   }
 
+  private async createStudioManager ({
+    projectId,
+    cloudDataSource,
+    cfg,
+    debugData,
+  }: {
+    projectId?: string
+    cloudDataSource: CloudDataSource
+    cfg: Cfg
+    debugData: any
+  }): Promise<StudioManager> {
+    const studioSession = await postStudioSession({
+      projectId,
+    })
+
+    const studioManager = await getAndInitializeStudioManager({
+      studioUrl: studioSession.studioUrl,
+      projectId,
+      cloudDataSource,
+    })
+
+    if (studioManager.status === 'ENABLED') {
+      debug('Cloud studio is enabled - setting up protocol')
+      const protocolManager = new ProtocolManager()
+      const script = await api.getCaptureProtocolScript(studioSession.protocolUrl)
+
+      await protocolManager.prepareProtocol(script, {
+        runId: 'studio',
+        projectId: cfg.projectId,
+        testingType: cfg.testingType,
+        cloudApi: {
+          url: routes.apiUrl,
+          retryWithBackoff: api.retryWithBackoff,
+          requestPromise: api.rp,
+        },
+        projectConfig: _.pick(cfg, ['devServerPublicPathRoute', 'port', 'proxyUrl', 'namespace']),
+        mountVersion: api.runnerCapabilities.protocolMountVersion,
+        debugData,
+        mode: 'studio',
+      })
+
+      studioManager.protocolManager = protocolManager
+    } else {
+      debug('Cloud studio is not enabled - skipping protocol setup')
+    }
+
+    debug('Studio is ready')
+    this.studioManager = studioManager
+    this.callRegisteredListeners()
+
+    return studioManager
+  }
+
   private callRegisteredListeners () {
     if (!this.studioManager) {
       throw new Error('Studio manager has not been initialized')

packages/server/lib/cloud/api/studio/get_and_initialize_studio_manager.ts

@@ -14,8 +14,12 @@ import { PUBLIC_KEY_VERSION } from '../../constants'
 import { CloudRequest } from '../cloud_request'
 import type { CloudDataSource } from '@packages/data-context/src/sources'
 
+interface Options {
+  studioUrl: string
+  projectId?: string
+}
+
 const pkg = require('@packages/root')
-const routes = require('../../routes')
 
 const _delay = linearDelay(500)
 
@@ -24,11 +28,11 @@ export const studioPath = path.join(os.tmpdir(), 'cypress', 'studio')
 const bundlePath = path.join(studioPath, 'bundle.tar')
 const serverFilePath = path.join(studioPath, 'server', 'index.js')
 
-const downloadStudioBundleToTempDirectory = async (projectId?: string): Promise<void> => {
+const downloadStudioBundleToTempDirectory = async ({ studioUrl, projectId }: Options): Promise<void> => {
   let responseSignature: string | null = null
 
   await (asyncRetry(async () => {
-    const response = await fetch(routes.apiRoutes.studio() as string, {
+    const response = await fetch(studioUrl, {
       // @ts-expect-error - this is supported
       agent,
       method: 'GET',
@@ -90,7 +94,7 @@ const getTarHash = (): Promise<string> => {
   })
 }
 
-export const retrieveAndExtractStudioBundle = async ({ projectId }: { projectId?: string } = {}): Promise<{ studioHash: string | undefined }> => {
+export const retrieveAndExtractStudioBundle = async ({ studioUrl, projectId }: Options): Promise<{ studioHash: string | undefined }> => {
   // First remove studioPath to ensure we have a clean slate
   await fs.promises.rm(studioPath, { recursive: true, force: true })
   await ensureDir(studioPath)
@@ -106,7 +110,7 @@ export const retrieveAndExtractStudioBundle = async ({ projectId }: { projectId?
     return { studioHash: undefined }
   }
 
-  await downloadStudioBundleToTempDirectory(projectId)
+  await downloadStudioBundleToTempDirectory({ studioUrl, projectId })
 
   const studioHash = await getTarHash()
 
@@ -118,7 +122,7 @@ export const retrieveAndExtractStudioBundle = async ({ projectId }: { projectId?
   return { studioHash }
 }
 
-export const getAndInitializeStudioManager = async ({ projectId, cloudDataSource }: { projectId?: string, cloudDataSource: CloudDataSource }): Promise<StudioManager> => {
+export const getAndInitializeStudioManager = async ({ studioUrl, projectId, cloudDataSource }: { studioUrl: string, projectId?: string, cloudDataSource: CloudDataSource }): Promise<StudioManager> => {
   let script: string
 
   const cloudEnv = (process.env.CYPRESS_INTERNAL_ENV || 'production') as 'development' | 'staging' | 'production'
@@ -128,7 +132,7 @@ export const getAndInitializeStudioManager = async ({ projectId, cloudDataSource
   let studioHash: string | undefined
 
   try {
-    ({ studioHash } = await retrieveAndExtractStudioBundle({ projectId }))
+    ({ studioHash } = await retrieveAndExtractStudioBundle({ studioUrl, projectId }))
 
     script = await readFile(serverFilePath, 'utf8')
 

packages/server/lib/cloud/api/studio/post_studio_session.ts

@@ -0,0 +1,42 @@
+import { asyncRetry, linearDelay } from '../../../util/async_retry'
+import { isRetryableError } from '../../network/is_retryable_error'
+import fetch from 'cross-fetch'
+import os from 'os'
+
+const pkg = require('@packages/root')
+const routes = require('../../routes') as typeof import('../../routes')
+
+interface GetStudioSessionOptions {
+  projectId?: string
+}
+
+const _delay = linearDelay(500)
+
+export const postStudioSession = async ({ projectId }: GetStudioSessionOptions) => {
+  return await (asyncRetry(async () => {
+    const response = await fetch(routes.apiRoutes.studioSession(), {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+        'x-os-name': os.platform(),
+        'x-cypress-version': pkg.version,
+      },
+      body: JSON.stringify({ projectSlug: projectId, studioMountVersion: 1, protocolMountVersion: 2 }),
+    })
+
+    if (!response.ok) {
+      throw new Error('Failed to create studio session')
+    }
+
+    const data = await response.json()
+
+    return {
+      studioUrl: data.studioUrl,
+      protocolUrl: data.protocolUrl,
+    }
+  }, {
+    maxAttempts: 3,
+    retryDelay: _delay,
+    shouldRetry: isRetryableError,
+  }))()
+}

packages/server/lib/cloud/routes.ts

@@ -16,8 +16,7 @@ const CLOUD_ENDPOINTS = {
   instanceStdout: 'instances/:id/stdout',
   instanceArtifacts: 'instances/:id/artifacts',
   captureProtocolErrors: 'capture-protocol/errors',
-  captureProtocolCurrent: 'capture-protocol/script/current.js',
-  studio: 'studio/bundle/current.tgz',
+  studioSession: 'studio/session',
   studioErrors: 'studio/errors',
   exceptions: 'exceptions',
   telemetry: 'telemetry',

packages/server/lib/project-base.ts

@@ -476,11 +476,12 @@ export class ProjectBase extends EE {
 
         const studio = await this.ctx.coreData.studioLifecycleManager?.getStudio()
 
-        if (studio?.protocolManager) {
+        await studio?.destroy()
+
+        if (this.protocolManager) {
           await browsers.closeProtocolConnection({ browser: this.browser, foundBrowsers: this.options.browsers })
           this.protocolManager?.close()
           this.protocolManager = undefined
-          await studio.destroy()
         }
       },
 

packages/server/test/unit/StudioLifecycleManager_spec.ts

@@ -9,6 +9,7 @@ import * as getAndInitializeStudioManagerModule from '../../lib/cloud/api/studio
 import * as reportStudioErrorPath from '../../lib/cloud/api/studio/report_studio_error'
 import ProtocolManager from '../../lib/cloud/protocol'
 const api = require('../../lib/cloud/api').default
+import * as postStudioSessionModule from '../../lib/cloud/api/studio/post_studio_session'
 
 // Helper to wait for next tick in event loop
 const nextTick = () => new Promise((resolve) => process.nextTick(resolve))
@@ -19,6 +20,7 @@ describe('StudioLifecycleManager', () => {
   let mockCtx: DataContext
   let mockCloudDataSource: CloudDataSource
   let mockCfg: Cfg
+  let postStudioSessionStub: sinon.SinonStub
   let getAndInitializeStudioManagerStub: sinon.SinonStub
   let getCaptureProtocolScriptStub: sinon.SinonStub
   let prepareProtocolStub: sinon.SinonStub
@@ -53,6 +55,12 @@ describe('StudioLifecycleManager', () => {
       namespace: '__cypress',
     } as unknown as Cfg
 
+    postStudioSessionStub = sinon.stub(postStudioSessionModule, 'postStudioSession')
+    postStudioSessionStub.resolves({
+      studioUrl: 'https://cloud.cypress.io/studio/bundle/abc.tgz',
+      protocolUrl: 'https://cloud.cypress.io/capture-protocol/script/def.js',
+    })
+
     getAndInitializeStudioManagerStub = sinon.stub(getAndInitializeStudioManagerModule, 'getAndInitializeStudioManager')
     getAndInitializeStudioManagerStub.resolves(mockStudioManager)
 
@@ -107,7 +115,11 @@ describe('StudioLifecycleManager', () => {
 
       await studioReadyPromise
 
-      expect(getCaptureProtocolScriptStub).to.be.calledWith('http://localhost:1234/capture-protocol/script/current.js')
+      expect(postStudioSessionStub).to.be.calledWith({
+        projectId: 'abc123',
+      })
+
+      expect(getCaptureProtocolScriptStub).to.be.calledWith('https://cloud.cypress.io/capture-protocol/script/def.js')
       expect(prepareProtocolStub).to.be.calledWith('console.log("hello")', {
         runId: 'studio',
         projectId: 'abc123',