wait for tunnels to be ready (take 2)

packages/cli-common/src/lib/common-flags/index.ts

@@ -94,4 +94,9 @@ export const urlFlags = {
     summary: 'Timeout for fetching URLs request in milliseconds',
     default: 2500,
   }),
+  wait: Flags.boolean({
+    description: 'Wait for all tunnels to be ready',
+    default: true,
+    allowNo: true,
+  }),
 } as const

packages/cli/package.json

@@ -123,4 +123,4 @@
     "preview"
   ],
   "types": "dist/index.d.ts"
-}
+}

packages/cli/src/commands/up.ts

@@ -16,7 +16,7 @@ import { inspect } from 'util'
 import { editUrl, tunnelNameResolver } from '@preevy/common'
 import MachineCreationDriverCommand from '../machine-creation-driver-command.js'
 import { envIdFlags, urlFlags } from '../common-flags.js'
-import { filterUrls, printUrls, writeUrlsToFile } from './urls.js'
+import { filterUrls, printUrls, urlsRetryOpts, writeUrlsToFile } from './urls.js'
 import { connectToTunnelServerSsh } from '../tunnel-server-client.js'
 
 const fetchTunnelServerDetails = async ({
@@ -207,13 +207,9 @@ export default class Up extends MachineCreationDriverCommand<typeof Up> {
       tunnelingKey,
       includeAccessCredentials: flags['include-access-credentials'] && (flags['access-credentials-type'] as 'api' | 'browser'),
       showPreevyService: flags['show-preevy-service-urls'],
-      retryOpts: {
-        minTimeout: 1000,
-        maxTimeout: 2000,
-        retries: 10,
-        onFailedAttempt: e => { this.logger.debug(`Failed to query tunnels: ${inspect(e)}`) },
-      },
+      retryOpts: urlsRetryOpts(this.logger),
       fetchTimeout: flags['fetch-urls-timeout'],
+      waitForAllTunnels: flags.wait,
     }), { text: 'Getting tunnel URLs...' })
 
     const urls = await filterUrls({

packages/cli/src/commands/urls.ts

@@ -5,6 +5,7 @@ import { FlatTunnel, Logger, TunnelOpts, addBaseComposeTunnelAgentService, comma
 import { HooksListeners, PluginContext, parseTunnelServerFlags, tableFlags, text, tunnelServerFlags } from '@preevy/cli-common'
 import { asyncReduce } from 'iter-tools-es'
 import { tunnelNameResolver } from '@preevy/common'
+import { inspect } from 'util'
 import { connectToTunnelServerSsh } from '../tunnel-server-client.js'
 import ProfileCommand from '../profile-command.js'
 import { envIdFlags, urlFlags } from '../common-flags.js'
@@ -49,6 +50,15 @@ export const filterUrls = ({ flatTunnels, context, filters }: {
   filters,
 )
 
+export const urlsRetryOpts = (log: Logger) => ({
+  minTimeout: 1000,
+  maxTimeout: 2000,
+  retries: 10,
+  onFailedAttempt: (e: unknown) => { log.debug(`Failed to query tunnels: ${inspect(e)}`) },
+} as const)
+
+export const noWaitUrlsRetryOpts = { retries: 2 } as const
+
 // eslint-disable-next-line no-use-before-define
 export default class Urls extends ProfileCommand<typeof Urls> {
   static description = 'Show urls for an existing environment'
@@ -128,8 +138,9 @@ export default class Urls extends ProfileCommand<typeof Urls> {
       tunnelingKey,
       includeAccessCredentials: flags['include-access-credentials'] && (flags['access-credentials-type'] as 'api' | 'browser'),
       showPreevyService: flags['show-preevy-service-urls'],
-      retryOpts: { retries: 2 },
+      retryOpts: flags.wait ? urlsRetryOpts(this.logger) : noWaitUrlsRetryOpts,
       fetchTimeout: flags['fetch-urls-timeout'],
+      waitForAllTunnels: flags.wait,
     })
 
     const urls = await filterUrls({

packages/common/index.ts

@@ -30,6 +30,7 @@ export {
   COMPOSE_TUNNEL_AGENT_PORT,
   COMPOSE_TUNNEL_AGENT_SERVICE_LABELS,
   COMPOSE_TUNNEL_AGENT_SERVICE_NAME,
+  ComposeTunnelAgentState,
 } from './src/compose-tunnel-agent/index.js'
 export { MachineStatusCommand, DockerMachineStatusCommandRecipe } from './src/machine-status-command.js'
 export { ProcessOutputBuffers, orderedOutput, OrderedOutput } from './src/process-output-buffers.js'

packages/common/src/compose-tunnel-agent/index.ts

@@ -3,3 +3,13 @@ export { ScriptInjection, parseScriptInjectionLabels, scriptInjectionsToLabels }
 
 export const COMPOSE_TUNNEL_AGENT_SERVICE_NAME = 'preevy_proxy'
 export const COMPOSE_TUNNEL_AGENT_PORT = 3000
+
+export type ComposeTunnelAgentState = {
+  state: 'unknown'
+  reason: string
+} | {
+  state: 'pending'
+  pendingServices: string[]
+} | {
+  state: 'stable'
+}

packages/compose-tunnel-agent/index.ts

@@ -4,6 +4,7 @@ import Docker from 'dockerode'
 import { rimraf } from 'rimraf'
 import { pino } from 'pino'
 import pinoPrettyModule from 'pino-pretty'
+import yaml from 'yaml'
 import {
   requiredEnv,
   formatPublicKey,
@@ -20,11 +21,13 @@ import { runMachineStatusCommand } from './src/machine-status.js'
 import { envMetadata } from './src/metadata.js'
 import { readAllFiles } from './src/files.js'
 import { eventsClient as dockerEventsClient, filteredClient as dockerFilteredClient } from './src/docker/index.js'
+import { tunnelsStateCalculator } from './src/tunnels-state.js'
 
 const PinoPretty = pinoPrettyModule.default
 
 const homeDir = process.env.HOME || '/root'
 const dockerSocket = '/var/run/docker.sock'
+const COMPOSE_FILE_PATH = '/preevy/docker-compose.yaml'
 
 const targetComposeProject = process.env.COMPOSE_PROJECT
 const defaultAccess = process.env.DEFAULT_ACCESS_LEVEL === 'private' ? 'private' : 'public'
@@ -106,22 +109,39 @@ const main = async () => {
   })
 
   sshLog.info('ssh client connected to %j', sshUrl)
-  let currentTunnels = dockerClient.getRunningServices().then(services => sshClient.updateTunnels(services))
+  let currentState = dockerClient.getRunningServices().then(async runningServices => ({
+    runningServices,
+    sshTunnels: await sshClient.updateTunnels(runningServices),
+  }))
 
   void dockerClient.startListening({
-    onChange: async services => {
-      currentTunnels = sshClient.updateTunnels(services)
+    onChange: runningServices => {
+      currentState = (async () => ({
+        runningServices,
+        sshTunnels: await sshClient.updateTunnels(runningServices),
+      }))()
     },
   })
 
+  const calcTunnelsState = tunnelsStateCalculator({
+    composeProject: targetComposeProject,
+    composeModelReader: async () => yaml.parse(await fs.promises.readFile(COMPOSE_FILE_PATH, { encoding: 'utf8' })),
+  })
+
   const app = await createApp({
     log: log.child({ name: 'api' }),
-    currentSshState: async () => (await currentTunnels),
+    tunnels: async () => {
+      const { sshTunnels, runningServices } = await currentState
+      return {
+        ...sshTunnels,
+        state: await calcTunnelsState(runningServices),
+      }
+    },
     machineStatus: machineStatusCommand
       ? async () => await runMachineStatusCommand({ log, docker })(machineStatusCommand)
       : undefined,
     envMetadata: await envMetadata({ env: process.env, log }),
-    composeModelPath: '/preevy/docker-compose.yaml',
+    composeModelPath: COMPOSE_FILE_PATH,
     docker,
     dockerFilter: dockerFilteredClient({ docker, composeProject: targetComposeProject }),
   })

packages/compose-tunnel-agent/package.json

@@ -31,6 +31,7 @@
     "rimraf": "^5.0.5",
     "ssh2": "^1.12.0",
     "ws": "^8.13.0",
+    "yaml": "^2.3.2",
     "zod": "^3.21.4"
   },
   "devDependencies": {
@@ -65,4 +66,4 @@
     "prepare": "cd ../.. && husky install",
     "test": "node --no-warnings=ExperimentalWarning --experimental-vm-modules ../../node_modules/.bin/jest"
   }
-}
+}

packages/compose-tunnel-agent/src/api-server/env.ts

@@ -1,14 +1,19 @@
 import fs from 'node:fs'
 import { FastifyPluginAsync } from 'fastify'
+import { ComposeTunnelAgentState } from '@preevy/common'
 import { SshState } from '../ssh/index.js'
 
 export const env: FastifyPluginAsync<{
-  currentSshState: () => Promise<SshState>
+  tunnels: () => Promise<SshState & { state: ComposeTunnelAgentState }>
   machineStatus?: () => Promise<{ data: Buffer; contentType: string }>
   envMetadata?: Record<string, unknown>
   composeModelPath: string
-}> = async (app, { currentSshState, machineStatus, envMetadata, composeModelPath }) => {
-  app.get('/tunnels', async () => await currentSshState())
+}> = async (app, { tunnels, machineStatus, envMetadata, composeModelPath }) => {
+  app.get('/tunnels', async ({ log }) => {
+    const response = await tunnels()
+    log.debug('tunnels response: %j', response)
+    return response
+  })
 
   if (machineStatus) {
     app.get('/machine-status', async (_req, res) => {

packages/compose-tunnel-agent/src/api-server/index.test.ts

@@ -8,6 +8,7 @@ import { inspect, promisify } from 'node:util'
 import waitForExpect from 'wait-for-expect'
 import WebSocket from 'ws'
 import stripAnsi from 'strip-ansi'
+import { ComposeTunnelAgentState } from '@preevy/common'
 import { createApp } from './index.js'
 import { filteredClient } from '../docker/index.js'
 import { SshState } from '../ssh/index.js'
@@ -71,7 +72,7 @@ const setupApiServer = () => {
       docker,
       dockerFilter: filteredClient({ docker, composeProject: TEST_COMPOSE_PROJECT }),
       composeModelPath: '',
-      currentSshState: () => Promise.resolve({} as unknown as SshState),
+      tunnels: () => Promise.resolve({} as unknown as SshState & { state: ComposeTunnelAgentState }),
     })
     await app.listen({ port: 0 })
     const { port } = app.server.address() as AddressInfo

packages/compose-tunnel-agent/src/api-server/index.ts

@@ -3,22 +3,23 @@ import Dockerode from 'dockerode'
 import fastify from 'fastify'
 import cors from '@fastify/cors'
 import { validatorCompiler, serializerCompiler, ZodTypeProvider } from 'fastify-type-provider-zod'
+import { ComposeTunnelAgentState } from '@preevy/common'
 import { SshState } from '../ssh/index.js'
 import { DockerFilterClient } from '../docker/index.js'
 import { containers } from './containers/index.js'
 import { env } from './env.js'
 
 export const createApp = async ({
   log,
-  currentSshState,
+  tunnels,
   machineStatus,
   envMetadata,
   composeModelPath,
   dockerFilter,
   docker,
 }: {
   log: Logger
-  currentSshState: () => Promise<SshState>
+  tunnels: () => Promise<SshState & { state: ComposeTunnelAgentState }>
   machineStatus?: () => Promise<{ data: Buffer; contentType: string }>
   envMetadata?: Record<string, unknown>
   composeModelPath: string
@@ -39,7 +40,7 @@ export const createApp = async ({
 
   app.get('/healthz', { logLevel: 'warn' }, async () => 'OK')
 
-  await app.register(env, { composeModelPath, currentSshState, envMetadata, machineStatus })
+  await app.register(env, { composeModelPath, tunnels, envMetadata, machineStatus })
   await app.register(containers, { docker, dockerFilter, prefix: '/containers' })
 
   return app

packages/compose-tunnel-agent/src/tunnels-state.ts

@@ -0,0 +1,51 @@
+import { inspect } from 'util'
+import { ComposeTunnelAgentState } from '@preevy/common'
+import { RunningService } from './docker/index.js'
+
+const findPendingComposeServiceTunnels = ({
+  composeProject,
+  composeModel,
+  runningServices,
+}: {
+  composeProject: string
+  composeModel: { services: Record<string, unknown> }
+  runningServices: Pick<RunningService, 'name' | 'project'>[]
+}) => {
+  const composeServiceNames = Object.keys(composeModel.services)
+
+  const runningServiceNames = new Set(
+    runningServices
+      .filter(({ project }) => project === composeProject)
+      .map(({ name }) => name)
+  )
+
+  return composeServiceNames.filter(service => !runningServiceNames.has(service))
+}
+
+export const tunnelsStateCalculator = ({
+  composeProject,
+  composeModelReader,
+}: {
+  composeProject?: string
+  composeModelReader: () => Promise<{ services: Record<string, unknown> }>
+}) => async (
+  runningServices: Pick<RunningService, 'name' | 'project'>[]
+): Promise<ComposeTunnelAgentState> => {
+  if (!composeProject) {
+    return { state: 'unknown', reason: 'COMPOSE_PROJECT not set' }
+  }
+
+  let composeModel: { services: Record<string, unknown> }
+  try {
+    composeModel = await composeModelReader()
+  } catch (e) {
+    return { state: 'unknown', reason: `Could not read compose file: ${inspect(e)}` }
+  }
+
+  const pendingServices = findPendingComposeServiceTunnels({ composeProject, composeModel, runningServices })
+  if (pendingServices.length) {
+    return { state: 'pending', pendingServices }
+  }
+
+  return { state: 'stable' }
+}

packages/core/src/commands/urls.ts

@@ -26,6 +26,7 @@ export const urls = async ({
   showPreevyService,
   composeTunnelServiceUrl,
   fetchTimeout,
+  waitForAllTunnels,
 }: {
   serviceAndPort?: { service: string; port?: number }
   tunnelingKey: string | Buffer
@@ -34,6 +35,7 @@ export const urls = async ({
   showPreevyService: boolean
   composeTunnelServiceUrl: string
   fetchTimeout: number
+  waitForAllTunnels?: boolean
 }) => {
   const credentials = await generateBasicAuthCredentials(jwtGenerator(tunnelingKey))
 
@@ -43,6 +45,7 @@ export const urls = async ({
     credentials,
     includeAccessCredentials,
     fetchTimeout,
+    waitForAllTunnels,
   })
 
   return flattenTunnels(tunnels).filter(tunnelFilter({ serviceAndPort, showPreevyService }))

packages/core/src/compose-tunnel-agent-client.ts

@@ -1,9 +1,9 @@
 import path from 'path'
 import retry from 'p-retry'
-import util from 'util'
+import { inspect } from 'util'
 import { createRequire } from 'module'
 import { mapValues, merge } from 'lodash-es'
-import { COMPOSE_TUNNEL_AGENT_PORT, COMPOSE_TUNNEL_AGENT_SERVICE_LABELS, COMPOSE_TUNNEL_AGENT_SERVICE_NAME, MachineStatusCommand, ScriptInjection, dateReplacer } from '@preevy/common'
+import { COMPOSE_TUNNEL_AGENT_PORT, COMPOSE_TUNNEL_AGENT_SERVICE_LABELS, COMPOSE_TUNNEL_AGENT_SERVICE_NAME, ComposeTunnelAgentState, MachineStatusCommand, ScriptInjection, dateReplacer } from '@preevy/common'
 import { ComposeModel, ComposeService, composeModelFilename } from './compose/model.js'
 import { TunnelOpts } from './ssh/url.js'
 import { Tunnel } from './tunneling/index.js'
@@ -150,7 +150,7 @@ export const findComposeTunnelAgentUrl = (
   )?.url
 
   if (!serviceUrl) {
-    throw new Error(`Cannot find compose tunnel agent API service URL ${COMPOSE_TUNNEL_AGENT_SERVICE_NAME}:${COMPOSE_TUNNEL_AGENT_PORT} in: ${util.inspect(serviceUrls)}`)
+    throw new Error(`Cannot find compose tunnel agent API service URL ${COMPOSE_TUNNEL_AGENT_SERVICE_NAME}:${COMPOSE_TUNNEL_AGENT_PORT} in: ${inspect(serviceUrls)}`)
   }
 
   return serviceUrl
@@ -182,16 +182,27 @@ const fetchFromComposeTunnelAgent = async ({
   return r
 }, retryOpts)
 
+type TunnelsResponse = {
+  tunnels: Tunnel[]
+  state: ComposeTunnelAgentState
+}
+
 export const queryTunnels = async ({
   includeAccessCredentials,
+  waitForAllTunnels,
   ...fetchOpts
 }: ComposeTunnelAgentFetchOpts & {
   includeAccessCredentials: false | 'browser' | 'api'
+  waitForAllTunnels?: boolean
 }) => {
   const r = await fetchFromComposeTunnelAgent({ ...fetchOpts, pathAndQuery: 'tunnels' })
-  const { tunnels } = await (r.json() as Promise<{ tunnels: Tunnel[] }>)
+  const response = await (r.json() as Promise<TunnelsResponse>)
+
+  if (waitForAllTunnels && response.state.state !== 'stable') {
+    throw new AgentFetchError(`Not all configured tunnels are ready yet: ${inspect(response, { depth: null })}`)
+  }
 
-  return tunnels
+  return response.tunnels
     .map(tunnel => ({
       ...tunnel,
       ports: mapValues(