From 7860068ebc53cd8600a4b0b2cbafe3e642d8f428 Mon Sep 17 00:00:00 2001 From: vmelikyan Date: Sun, 19 Apr 2026 15:12:03 -0700 Subject: [PATCH] harden agent session runtime and prewarm matching --- .../__tests__/devModeManager.test.ts | 62 +++ .../agentSession/__tests__/podFactory.test.ts | 20 +- .../__tests__/runtimeConfig.test.ts | 8 +- src/server/lib/agentSession/devModeManager.ts | 48 ++ .../lib/agentSession/editorServiceFactory.ts | 2 +- src/server/lib/agentSession/podFactory.ts | 41 +- src/server/lib/agentSession/runtimeConfig.ts | 2 +- src/server/lib/kubernetes/JobMonitor.ts | 2 +- .../services/__tests__/agentPrewarm.test.ts | 415 +++++++++++++++++- .../services/__tests__/agentSession.test.ts | 141 +++++- src/server/services/agentPrewarm.ts | 197 +++++++-- src/server/services/agentSession.ts | 30 +- 12 files changed, 869 insertions(+), 99 deletions(-) diff --git a/src/server/lib/agentSession/__tests__/devModeManager.test.ts b/src/server/lib/agentSession/__tests__/devModeManager.test.ts index b741a17a..d149e9ec 100644 --- a/src/server/lib/agentSession/__tests__/devModeManager.test.ts +++ b/src/server/lib/agentSession/__tests__/devModeManager.test.ts @@ -131,6 +131,7 @@ describe('DevModeManager', () => { template: { spec: { nodeSelector: { 'kubernetes.io/hostname': 'agent-node-a' }, + securityContext: { fsGroup: 1000, fsGroupChangePolicy: 'OnRootMismatch' }, volumes: [{ name: 'workspace', persistentVolumeClaim: { claimName: 'agent-pvc-abc' } }], containers: [ expect.objectContaining({ @@ -186,6 +187,43 @@ describe('DevModeManager', () => { }); }); + it('preserves existing pod securityContext fields while normalizing shared workspace ownership', async () => { + mockReadDeployment.mockResolvedValue({ + body: { + spec: { + template: { + spec: { + securityContext: { + fsGroup: 2000, + runAsNonRoot: false, + supplementalGroups: [2000], + }, + containers: [{ name: 'web-app' }], + }, + }, + }, + }, + }); + + const opts: DevModeOptions = { + namespace: 'test-ns', + deploymentName: 'my-app', + serviceName: 'my-app', + pvcName: 'agent-pvc-abc', + devConfig: { image: 'node:20-slim', command: 'pnpm dev', workDir: '/workspace' }, + }; + + await manager.enableDevMode(opts); + + const patchBody = mockPatchDeployment.mock.calls[0][2]; + expect(patchBody.spec.template.spec.securityContext).toEqual({ + fsGroup: 1000, + fsGroupChangePolicy: 'OnRootMismatch', + runAsNonRoot: false, + supplementalGroups: [2000], + }); + }); + it('mounts the shared workspace root when workDir points at a service subdirectory', async () => { const opts: DevModeOptions = { namespace: 'test-ns', @@ -325,6 +363,7 @@ describe('DevModeManager', () => { replicas: 3, template: { spec: { + securityContext: { fsGroup: 2000, supplementalGroups: [2000] }, containers: [{ name: 'web-app', image: 'registry.example/my-app:built' }], }, }, @@ -343,6 +382,10 @@ describe('DevModeManager', () => { const snapshot = await manager.enableDevMode(opts); expect(snapshot.deployment.replicas).toBe(3); + expect(snapshot.deployment.securityContext).toEqual({ + fsGroup: 2000, + supplementalGroups: [2000], + }); }); it('pins an attached HorizontalPodAutoscaler to a single replica during dev mode', async () => { @@ -413,6 +456,7 @@ describe('DevModeManager', () => { volumeMounts: [{ name: 'config-volume', mountPath: '/config' }], }, ], + securityContext: { fsGroup: 2000 }, volumes: [{ name: 'config-volume' }], }, }, @@ -424,6 +468,7 @@ describe('DevModeManager', () => { replicas: 1, template: { spec: { + securityContext: { fsGroup: 1000, fsGroupChangePolicy: 'OnRootMismatch' }, containers: [ { name: 'web-app', @@ -458,6 +503,7 @@ describe('DevModeManager', () => { { op: 'remove', path: '/spec/template/spec/volumes/0' }, { op: 'remove', path: '/spec/replicas' }, { op: 'remove', path: '/spec/template/spec/nodeSelector' }, + { op: 'replace', path: '/spec/template/spec/securityContext', value: { fsGroup: 2000 } }, ], undefined, undefined, @@ -536,6 +582,7 @@ describe('DevModeManager', () => { volumeMounts: [{ name: 'config-volume', mountPath: '/config' }], volumes: [{ name: 'config-volume', emptyDir: {} }], nodeSelector: { 'app-long': 'deployments-m7i' }, + securityContext: { fsGroup: 2000 }, }), 'lifecycle.goodrx.com/dev-mode-hpa-snapshot': JSON.stringify({ hpaName: 'grpc-echo-hpa', @@ -548,6 +595,7 @@ describe('DevModeManager', () => { replicas: 1, template: { spec: { + securityContext: { fsGroup: 1000, fsGroupChangePolicy: 'OnRootMismatch' }, nodeSelector: { 'app-long': 'deployments-m7i', 'kubernetes.io/hostname': 'agent-node-a', @@ -630,6 +678,11 @@ describe('DevModeManager', () => { path: '/spec/template/spec/nodeSelector', value: { 'app-long': 'deployments-m7i' }, }, + { + op: 'replace', + path: '/spec/template/spec/securityContext', + value: { fsGroup: 2000 }, + }, ], undefined, undefined, @@ -669,6 +722,7 @@ describe('DevModeManager', () => { replicas: 1, template: { spec: { + securityContext: { fsGroup: 1000, fsGroupChangePolicy: 'OnRootMismatch' }, containers: [ { name: 'grpc-echo', @@ -703,6 +757,7 @@ describe('DevModeManager', () => { { op: 'remove', path: '/spec/template/spec/containers/0/volumeMounts/0' }, { op: 'remove', path: '/spec/template/spec/volumes/0' }, { op: 'remove', path: '/spec/template/spec/nodeSelector' }, + { op: 'remove', path: '/spec/template/spec/securityContext' }, ], undefined, undefined, @@ -726,6 +781,7 @@ describe('DevModeManager', () => { replicas: 1, template: { spec: { + securityContext: { fsGroup: 1000, fsGroupChangePolicy: 'OnRootMismatch' }, containers: [ { name: 'lc-apps', @@ -777,6 +833,7 @@ describe('DevModeManager', () => { volumeMounts: null, volumes: null, nodeSelector: { 'app-long': 'deployments-m7i' }, + securityContext: { fsGroup: 2000 }, }, service: { serviceName: 'grpc-echo-service', @@ -809,6 +866,11 @@ describe('DevModeManager', () => { path: '/spec/template/spec/nodeSelector', value: { 'app-long': 'deployments-m7i' }, }, + { + op: 'replace', + path: '/spec/template/spec/securityContext', + value: { fsGroup: 2000 }, + }, ], undefined, undefined, diff --git a/src/server/lib/agentSession/__tests__/podFactory.test.ts b/src/server/lib/agentSession/__tests__/podFactory.test.ts index 162e3b64..a01bd0a3 100644 --- a/src/server/lib/agentSession/__tests__/podFactory.test.ts +++ b/src/server/lib/agentSession/__tests__/podFactory.test.ts @@ -147,13 +147,12 @@ describe('podFactory', () => { describe('buildSessionWorkspacePodSpec', () => { it('creates a pod with init and main containers', () => { const pod = buildSessionWorkspacePodSpec(baseOpts); - expect(pod.spec!.initContainers).toHaveLength(4); + expect(pod.spec!.initContainers).toHaveLength(3); expect(pod.spec!.containers).toHaveLength(2); expect(pod.spec!.initContainers!.map((container) => container.name)).toEqual([ 'prepare-workspace', 'init-workspace', 'seed-runtime-config', - 'prepare-editor-workspace', ]); expect(pod.spec!.containers!.map((container) => container.name)).toEqual(['editor', 'workspace-gateway']); }); @@ -819,23 +818,21 @@ describe('podFactory', () => { ], }); - expect(getInitContainer(pod, 'prepare-editor-workspace')).toEqual( + expect(getInitContainer(pod, 'seed-runtime-config')).toEqual( expect.objectContaining({ command: [ 'sh', '-c', expect.stringContaining(`cat > '${SESSION_WORKSPACE_EDITOR_PROJECT_FILE}' << 'WORKSPACE_EOF'`), ], - volumeMounts: [{ name: 'tmp', mountPath: '/tmp' }], + volumeMounts: expect.arrayContaining([{ name: 'tmp', mountPath: '/tmp' }]), }) ); - expect(getInitContainer(pod, 'prepare-editor-workspace').command?.[2]).toContain('"name": "org/repo"'); - expect(getInitContainer(pod, 'prepare-editor-workspace').command?.[2]).toContain( + expect(getInitContainer(pod, 'seed-runtime-config').command?.[2]).toContain('"name": "org/repo"'); + expect(getInitContainer(pod, 'seed-runtime-config').command?.[2]).toContain( '"path": "/workspace/repos/org/repo"' ); - expect(getInitContainer(pod, 'prepare-editor-workspace').command?.[2]).toContain( - '"path": "/workspace/repos/org/api"' - ); + expect(getInitContainer(pod, 'seed-runtime-config').command?.[2]).toContain('"path": "/workspace/repos/org/api"'); expect(getContainer(pod, 'workspace-gateway').env).toEqual( expect.arrayContaining([{ name: 'LIFECYCLE_SESSION_PRIMARY_REPO_PATH', value: '/workspace/repos/org/repo' }]) ); @@ -847,10 +844,7 @@ describe('podFactory', () => { skipWorkspaceBootstrap: true, }); - expect(pod.spec!.initContainers?.map((container) => container.name)).toEqual([ - 'seed-runtime-config', - 'prepare-editor-workspace', - ]); + expect(pod.spec!.initContainers?.map((container) => container.name)).toEqual(['seed-runtime-config']); }); it('does not set runtimeClassName when gVisor not requested', () => { diff --git a/src/server/lib/agentSession/__tests__/runtimeConfig.test.ts b/src/server/lib/agentSession/__tests__/runtimeConfig.test.ts index 3be87bb9..85910268 100644 --- a/src/server/lib/agentSession/__tests__/runtimeConfig.test.ts +++ b/src/server/lib/agentSession/__tests__/runtimeConfig.test.ts @@ -43,7 +43,7 @@ import { const DEFAULT_READINESS = { timeoutMs: 60000, - pollMs: 2000, + pollMs: 1000, }; const DEFAULT_RESOURCES = { @@ -321,19 +321,19 @@ describe('runtimeConfig', () => { it('merges direct readiness overrides over runtime defaults', () => { expect( - mergeAgentSessionReadiness(resolveAgentSessionReadinessFromDefaults({ timeoutMs: 60000, pollMs: 2000 }), { + mergeAgentSessionReadiness(resolveAgentSessionReadinessFromDefaults({ timeoutMs: 60000, pollMs: 1000 }), { timeoutMs: 120000, }) ).toEqual({ timeoutMs: 120000, - pollMs: 2000, + pollMs: 1000, }); }); it('merges service readiness overrides over runtime defaults', () => { expect( mergeAgentSessionReadinessForServices( - resolveAgentSessionReadinessFromDefaults({ timeoutMs: 60000, pollMs: 2000 }), + resolveAgentSessionReadinessFromDefaults({ timeoutMs: 60000, pollMs: 1000 }), [{ timeoutMs: 120000 }, { timeoutMs: 90000, pollMs: 500 }, undefined, { pollMs: 1000 }] ) ).toEqual({ diff --git a/src/server/lib/agentSession/devModeManager.ts b/src/server/lib/agentSession/devModeManager.ts index 1b49e983..3da88b41 100644 --- a/src/server/lib/agentSession/devModeManager.ts +++ b/src/server/lib/agentSession/devModeManager.ts @@ -24,6 +24,9 @@ const DEV_MODE_DEPLOYMENT_SNAPSHOT_ANNOTATION = 'lifecycle.goodrx.com/dev-mode-d const DEV_MODE_SERVICE_SNAPSHOT_ANNOTATION = 'lifecycle.goodrx.com/dev-mode-service-snapshot'; const DEV_MODE_HPA_SNAPSHOT_ANNOTATION = 'lifecycle.goodrx.com/dev-mode-hpa-snapshot'; const SAME_NODE_SELECTOR_KEY = 'kubernetes.io/hostname'; +const SHARED_WORKSPACE_FS_GROUP = 1000; +const SHARED_WORKSPACE_FS_GROUP_CHANGE_POLICY: NonNullable = + 'OnRootMismatch'; export interface DevModeOptions { namespace: string; @@ -45,6 +48,7 @@ export interface DevModeDeploymentSnapshot { volumeMounts: k8s.V1VolumeMount[] | null; volumes: k8s.V1Volume[] | null; nodeSelector: Record | null; + securityContext: k8s.V1PodSecurityContext | null; } export interface DevModeServiceSnapshot { @@ -70,6 +74,7 @@ interface AppliedDeploymentTemplate { template?: { spec?: { nodeSelector?: Record; + securityContext?: k8s.V1PodSecurityContext; containers?: Array<{ name?: string; command?: string[]; @@ -106,6 +111,23 @@ function selectorMatches( return Object.entries(expected).every(([key, value]) => selector[key] === value); } +function buildSharedWorkspacePodSecurityContext( + securityContext?: k8s.V1PodSecurityContext | null +): k8s.V1PodSecurityContext { + return { + ...(securityContext || {}), + fsGroup: SHARED_WORKSPACE_FS_GROUP, + fsGroupChangePolicy: SHARED_WORKSPACE_FS_GROUP_CHANGE_POLICY, + }; +} + +function hasSharedWorkspacePodSecurityContext(securityContext?: k8s.V1PodSecurityContext | null): boolean { + return ( + securityContext?.fsGroup === SHARED_WORKSPACE_FS_GROUP && + securityContext.fsGroupChangePolicy === SHARED_WORKSPACE_FS_GROUP_CHANGE_POLICY + ); +} + export class DevModeManager { private kc: k8s.KubeConfig; private appsApi: k8s.AppsV1Api; @@ -376,6 +398,7 @@ export class DevModeManager { const deploymentName = existing.metadata?.name || opts.deploymentName; const workDir = devConfig.workDir || '/workspace'; const existingContainerName = existing.spec?.template?.spec?.containers?.[0]?.name || deploymentName; + const existingPodSecurityContext = existing.spec?.template?.spec?.securityContext; const nodeSelector = requiredNodeName ? { ...(existing.spec?.template?.spec?.nodeSelector || {}), @@ -397,6 +420,7 @@ export class DevModeManager { template: { spec: { ...(nodeSelector ? { nodeSelector } : {}), + securityContext: buildSharedWorkspacePodSecurityContext(existingPodSecurityContext), volumes: [{ name: 'workspace', persistentVolumeClaim: { claimName: pvcName } }], containers: [ { @@ -590,6 +614,12 @@ export class DevModeManager { liveSpec?.nodeSelector, desiredTemplate.spec?.template?.spec?.nodeSelector ); + this.appendValuePatch( + patch, + '/spec/template/spec/securityContext', + liveSpec?.securityContext, + desiredTemplate.spec?.template?.spec?.securityContext + ); if (patch.length === 0) { return; @@ -696,6 +726,7 @@ export class DevModeManager { volumeMounts: liveContainer?.volumeMounts ? deepClone(liveContainer.volumeMounts) : null, volumes: liveSpec?.volumes ? deepClone(liveSpec.volumes) : null, nodeSelector: liveSpec?.nodeSelector ? deepClone(liveSpec.nodeSelector) : null, + securityContext: liveSpec?.securityContext ? deepClone(liveSpec.securityContext) : null, }; } @@ -824,6 +855,12 @@ export class DevModeManager { existing.spec?.template?.spec?.nodeSelector, snapshot.nodeSelector ); + this.appendValuePatch( + patch, + '/spec/template/spec/securityContext', + existing.spec?.template?.spec?.securityContext, + snapshot.securityContext + ); if (patch.length === 0) { return; @@ -992,6 +1029,17 @@ export class DevModeManager { } } + const liveSecurityContext = liveSpec?.securityContext; + if (isLikelyDevPatched && hasSharedWorkspacePodSecurityContext(liveSecurityContext)) { + const securityContextKeys = Object.keys(liveSecurityContext || {}); + if (securityContextKeys.length === 2) { + patch.push({ op: 'remove', path: '/spec/template/spec/securityContext' }); + } else { + patch.push({ op: 'remove', path: '/spec/template/spec/securityContext/fsGroup' }); + patch.push({ op: 'remove', path: '/spec/template/spec/securityContext/fsGroupChangePolicy' }); + } + } + return patch; } } diff --git a/src/server/lib/agentSession/editorServiceFactory.ts b/src/server/lib/agentSession/editorServiceFactory.ts index 04313c77..edd4ad79 100644 --- a/src/server/lib/agentSession/editorServiceFactory.ts +++ b/src/server/lib/agentSession/editorServiceFactory.ts @@ -69,7 +69,7 @@ export async function createSessionWorkspaceService( const { body: result } = await coreApi.createNamespacedService(namespace, service); logger.info( - `Session: workspace editor ready serviceName=${serviceName} namespace=${namespace} port=${SESSION_WORKSPACE_EDITOR_PORT}` + `Session: workspace service created serviceName=${serviceName} namespace=${namespace} port=${SESSION_WORKSPACE_EDITOR_PORT}` ); return result; } diff --git a/src/server/lib/agentSession/podFactory.ts b/src/server/lib/agentSession/podFactory.ts index 50cffb51..29c1d319 100644 --- a/src/server/lib/agentSession/podFactory.ts +++ b/src/server/lib/agentSession/podFactory.ts @@ -346,6 +346,13 @@ function generateEditorWorkspaceInitScript(workspaceRepos: AgentSessionWorkspace ].join('\n'); } +function stripShellScriptPreamble(script: string): string { + return script + .replace(/^#!\/bin\/sh\s*\n/, '') + .replace(/^set -e\s*\n/, '') + .trim(); +} + export function buildSessionWorkspacePodSpec(opts: SessionWorkspacePodOptions): k8s.V1Pod { const { podName, @@ -401,6 +408,9 @@ export function buildSessionWorkspacePodSpec(opts: SessionWorkspacePodOptions): }; const editorWorkspaceRepos = resolveEditorWorkspaceRepos(opts); const editorWorkspaceInitScript = generateEditorWorkspaceInitScript(editorWorkspaceRepos); + const sessionBootstrapScript = `${runtimeSeedScript.trim()}\n\n${stripShellScriptPreamble( + editorWorkspaceInitScript + )}\n`; const forwardedAgentEnv = opts.forwardedAgentEnv || {}; const forwardedAgentSecretEnv = buildPodEnvWithSecrets( forwardedAgentEnv, @@ -526,7 +536,7 @@ export function buildSessionWorkspacePodSpec(opts: SessionWorkspacePodOptions): name: 'seed-runtime-config', image: workspaceImage, imagePullPolicy: 'IfNotPresent', - command: ['sh', '-c', runtimeSeedScript], + command: ['sh', '-c', sessionBootstrapScript], resources, securityContext: { ...securityContext, @@ -551,29 +561,6 @@ export function buildSessionWorkspacePodSpec(opts: SessionWorkspacePodOptions): ], }); - initContainers.push({ - name: 'prepare-editor-workspace', - image: workspaceImage, - imagePullPolicy: 'IfNotPresent', - command: ['sh', '-c', editorWorkspaceInitScript], - resources, - securityContext: { - ...securityContext, - readOnlyRootFilesystem: false, - }, - volumeMounts: [ - { - name: 'tmp', - mountPath: '/tmp', - }, - ], - env: [ - { name: 'TMPDIR', value: '/tmp' }, - { name: 'TMP', value: '/tmp' }, - { name: 'TEMP', value: '/tmp' }, - ], - }); - const pod: k8s.V1Pod = { apiVersion: 'v1', kind: 'Pod', @@ -633,7 +620,7 @@ export function buildSessionWorkspacePodSpec(opts: SessionWorkspacePodOptions): port: SESSION_WORKSPACE_EDITOR_PORT, }, initialDelaySeconds: 1, - periodSeconds: 2, + periodSeconds: 1, }, volumeMounts: [ workspaceVolumeMount, @@ -683,7 +670,7 @@ export function buildSessionWorkspacePodSpec(opts: SessionWorkspacePodOptions): port: SESSION_WORKSPACE_GATEWAY_PORT, }, initialDelaySeconds: 1, - periodSeconds: 2, + periodSeconds: 1, }, volumeMounts: [ workspaceVolumeMount, @@ -784,7 +771,7 @@ async function waitForSessionWorkspacePodReady( const readyPollMs = normalizeNonNegativeInteger(readiness?.pollMs) ?? normalizeNonNegativeInteger(process.env.AGENT_SESSION_WORKSPACE_READY_POLL_MS) ?? - 2000; + 1000; const deadline = Date.now() + readyTimeoutMs; let lastObservedState = 'pending'; let lastPod: k8s.V1Pod | null = null; diff --git a/src/server/lib/agentSession/runtimeConfig.ts b/src/server/lib/agentSession/runtimeConfig.ts index 7140277c..b34b4624 100644 --- a/src/server/lib/agentSession/runtimeConfig.ts +++ b/src/server/lib/agentSession/runtimeConfig.ts @@ -74,7 +74,7 @@ export const DEFAULT_AGENT_SESSION_WORKSPACE_TOOL_EXECUTION_TIMEOUT_MS = 15000; export const DEFAULT_AGENT_SESSION_KEEP_ATTACHED_SERVICES_ON_SESSION_NODE = true; const DEFAULT_AGENT_READY_TIMEOUT_MS = 60000; -const DEFAULT_AGENT_READY_POLL_MS = 2000; +const DEFAULT_AGENT_READY_POLL_MS = 1000; const DEFAULT_WORKSPACE_RESOURCES: ResolvedAgentSessionResourceRequirements = { requests: { cpu: '500m', diff --git a/src/server/lib/kubernetes/JobMonitor.ts b/src/server/lib/kubernetes/JobMonitor.ts index eca281b1..8aba2b00 100644 --- a/src/server/lib/kubernetes/JobMonitor.ts +++ b/src/server/lib/kubernetes/JobMonitor.ts @@ -34,7 +34,7 @@ export interface MonitorOptions { export class JobMonitor { private static readonly DEFAULT_TIMEOUT = 1800; // 30 minutes - private static readonly POLL_INTERVAL = 2000; // 2 seconds + private static readonly POLL_INTERVAL = 1000; // 1 second // eslint-disable-next-line no-unused-vars constructor(private readonly jobName: string, private readonly namespace: string) {} diff --git a/src/server/services/__tests__/agentPrewarm.test.ts b/src/server/services/__tests__/agentPrewarm.test.ts index 3e8cdb6e..8cf5ade1 100644 --- a/src/server/services/__tests__/agentPrewarm.test.ts +++ b/src/server/services/__tests__/agentPrewarm.test.ts @@ -26,7 +26,7 @@ jest.mock('server/lib/agentSession/pvcFactory', () => ({ deleteAgentPvc: jest.fn(), })); jest.mock('../agentSessionCandidates', () => ({ - resolveAgentSessionServiceCandidates: jest.fn(), + resolveAgentSessionServiceCandidatesForBuild: jest.fn(), resolveRequestedAgentSessionServices: jest.fn(), })); jest.mock('server/lib/logger', () => ({ @@ -46,7 +46,10 @@ import Build from 'server/models/Build'; import { fetchLifecycleConfig } from 'server/models/yaml'; import AgentPrewarmService from 'server/services/agentPrewarm'; import { deleteAgentPvc } from 'server/lib/agentSession/pvcFactory'; -import { resolveAgentSessionServiceCandidates, resolveRequestedAgentSessionServices } from '../agentSessionCandidates'; +import { + resolveAgentSessionServiceCandidatesForBuild, + resolveRequestedAgentSessionServices, +} from '../agentSessionCandidates'; describe('AgentPrewarmService', () => { beforeEach(() => { @@ -82,6 +85,18 @@ describe('AgentPrewarmService', () => { revision: 'sha-123', configuredServiceNames: ['api', 'web'], services: [], + workspaceRepos: [ + { + repo: 'example-org/example-repo', + repoUrl: 'https://github.com/example-org/example-repo.git', + branch: 'sample-branch', + revision: 'sha-123', + mountPath: '/workspace', + primary: true, + }, + ], + serviceRefs: [], + skillPlan: { skills: [] }, }); (AgentPrewarm.query as jest.Mock).mockReturnValue({ where: jest.fn().mockReturnThis(), @@ -91,6 +106,17 @@ describe('AgentPrewarmService', () => { revision: 'sha-123', services: ['web', 'api'], status: 'ready', + workspaceRepos: [ + { + repo: 'example-org/example-repo', + repoUrl: 'https://github.com/example-org/example-repo.git', + branch: 'sample-branch', + revision: 'sha-123', + mountPath: '/workspace', + primary: true, + }, + ], + serviceRefs: [], }, ]), }); @@ -114,6 +140,18 @@ describe('AgentPrewarmService', () => { revision: 'sha-123', configuredServiceNames: ['api', 'web'], services: [], + workspaceRepos: [ + { + repo: 'example-org/example-repo', + repoUrl: 'https://github.com/example-org/example-repo.git', + branch: 'sample-branch', + revision: 'sha-123', + mountPath: '/workspace', + primary: true, + }, + ], + serviceRefs: [], + skillPlan: { skills: [] }, }); (AgentPrewarm.query as jest.Mock).mockReturnValue({ where: jest.fn().mockReturnThis(), @@ -128,10 +166,296 @@ describe('AgentPrewarmService', () => { buildUuid: 'build-123', requestId: 'req-123', }), + expect.objectContaining({ + jobId: expect.stringMatching(/^agent-prewarm:build-123:sha-123:api,web:/), + }) + ); + }); + + it('enqueues when an active prewarm has the same services but a different workspace identity', async () => { + const queueAdd = jest.fn().mockResolvedValue(undefined); + const queueManager = { + registerQueue: jest.fn().mockReturnValue({ add: queueAdd }), + }; + const service = new AgentPrewarmService({} as any, {} as any, {} as any, queueManager as any); + jest.spyOn(service as any, 'resolveBuildPrewarmPlan').mockResolvedValue({ + buildUuid: 'build-123', + namespace: 'env-sample', + repo: 'example-org/example-repo', + repoUrl: 'https://github.com/example-org/example-repo.git', + branch: 'sample-branch', + revision: 'sha-123', + configuredServiceNames: ['api', 'web'], + services: [], + workspaceRepos: [ + { + repo: 'example-org/example-repo', + repoUrl: 'https://github.com/example-org/example-repo.git', + branch: 'sample-branch', + revision: 'sha-123', + mountPath: '/workspace/repos/example-org/example-repo', + primary: true, + }, + { + repo: 'example-org/example-api', + repoUrl: 'https://github.com/example-org/example-api.git', + branch: 'api-branch', + revision: 'sha-api', + mountPath: '/workspace/repos/example-org/example-api', + primary: false, + }, + ], + serviceRefs: [ + { name: 'web', deployId: 1, repo: 'example-org/example-repo', branch: 'sample-branch' }, + { name: 'api', deployId: 2, repo: 'example-org/example-api', branch: 'api-branch' }, + ], + skillPlan: { skills: [] }, + }); + (AgentPrewarm.query as jest.Mock).mockReturnValue({ + where: jest.fn().mockReturnThis(), + whereIn: jest.fn().mockReturnThis(), + orderBy: jest.fn().mockResolvedValue([ + { + revision: 'sha-123', + services: ['web', 'api'], + workspaceRepos: [ + { + repo: 'example-org/example-repo', + repoUrl: 'https://github.com/example-org/example-repo.git', + branch: 'sample-branch', + revision: 'sha-123', + mountPath: '/workspace/repos/example-org/example-repo', + primary: false, + }, + { + repo: 'example-org/example-api', + repoUrl: 'https://github.com/example-org/example-api.git', + branch: 'api-branch', + revision: 'sha-api', + mountPath: '/workspace/repos/example-org/example-api', + primary: true, + }, + ], + serviceRefs: [ + { name: 'web', deployId: 1, repo: 'example-org/example-repo', branch: 'sample-branch' }, + { name: 'api', deployId: 2, repo: 'example-org/example-api', branch: 'api-branch' }, + ], + status: 'ready', + }, + ]), + }); + + await expect(service.queueBuildPrewarm('build-123')).resolves.toBe(true); + expect(queueAdd).toHaveBeenCalledTimes(1); + }); + }); + + describe('getCompatibleReadyPrewarm', () => { + it('matches a multi-repo prewarm when workspace layout and requested deploys align', async () => { + (AgentPrewarm.query as jest.Mock).mockReturnValue({ + where: jest.fn().mockReturnThis(), + orderBy: jest.fn().mockResolvedValue([ + { + uuid: 'prewarm-123', + status: 'ready', + revision: 'sha-repo', + services: ['web', 'api'], + workspaceRepos: [ + { + repo: 'example-org/example-repo', + repoUrl: 'https://github.com/example-org/example-repo.git', + branch: 'sample-branch', + revision: 'sha-repo', + mountPath: '/workspace/repos/example-org/example-repo', + primary: true, + }, + { + repo: 'example-org/example-api', + repoUrl: 'https://github.com/example-org/example-api.git', + branch: 'api-branch', + revision: 'sha-api', + mountPath: '/workspace/repos/example-org/example-api', + primary: false, + }, + ], + serviceRefs: [ + { name: 'web', deployId: 1, repo: 'example-org/example-repo', branch: 'sample-branch' }, + { name: 'api', deployId: 2, repo: 'example-org/example-api', branch: 'api-branch' }, + ], + }, + ]), + }); + const service = new AgentPrewarmService( + {} as any, + {} as any, + {} as any, + { + registerQueue: jest.fn().mockReturnValue({ add: jest.fn() }), + } as any + ); + + await expect( + service.getCompatibleReadyPrewarm({ + buildUuid: 'build-123', + requestedServices: ['api'], + revision: 'sha-repo', + workspaceRepos: [ + { + repo: 'example-org/example-repo', + repoUrl: 'https://github.com/example-org/example-repo.git', + branch: 'sample-branch', + revision: 'sha-repo', + mountPath: '/workspace/repos/example-org/example-repo', + primary: true, + }, + { + repo: 'example-org/example-api', + repoUrl: 'https://github.com/example-org/example-api.git', + branch: 'api-branch', + revision: 'sha-api', + mountPath: '/workspace/repos/example-org/example-api', + primary: false, + }, + ], + requestedServiceRefs: [ + { + name: 'api', + deployId: 2, + repo: 'example-org/example-api', + branch: 'api-branch', + workspacePath: '/workspace/repos/example-org/example-api', + }, + ], + }) + ).resolves.toEqual(expect.objectContaining({ uuid: 'prewarm-123' })); + }); + + it('rejects a ready prewarm when the requested workspace layout differs', async () => { + (AgentPrewarm.query as jest.Mock).mockReturnValue({ + where: jest.fn().mockReturnThis(), + orderBy: jest.fn().mockResolvedValue([ + { + uuid: 'prewarm-123', + status: 'ready', + revision: 'sha-repo', + services: ['web', 'api'], + workspaceRepos: [ + { + repo: 'example-org/example-repo', + repoUrl: 'https://github.com/example-org/example-repo.git', + branch: 'sample-branch', + revision: 'sha-repo', + mountPath: '/workspace/repos/example-org/example-repo', + primary: true, + }, + ], + serviceRefs: [{ name: 'web', deployId: 1, repo: 'example-org/example-repo', branch: 'sample-branch' }], + }, + ]), + }); + const service = new AgentPrewarmService( + {} as any, + {} as any, + {} as any, { - jobId: 'agent-prewarm:build-123:sha-123:api,web', - } + registerQueue: jest.fn().mockReturnValue({ add: jest.fn() }), + } as any ); + + await expect( + service.getCompatibleReadyPrewarm({ + buildUuid: 'build-123', + requestedServices: ['web'], + revision: 'sha-repo', + workspaceRepos: [ + { + repo: 'example-org/example-repo', + repoUrl: 'https://github.com/example-org/example-repo.git', + branch: 'sample-branch', + revision: 'sha-repo', + mountPath: '/workspace', + primary: true, + }, + ], + requestedServiceRefs: [ + { + name: 'web', + deployId: 1, + repo: 'example-org/example-repo', + branch: 'sample-branch', + workspacePath: '/workspace', + }, + ], + }) + ).resolves.toBeNull(); + }); + + it('rejects a ready prewarm when the requested deploy identity differs', async () => { + (AgentPrewarm.query as jest.Mock).mockReturnValue({ + where: jest.fn().mockReturnThis(), + orderBy: jest.fn().mockResolvedValue([ + { + uuid: 'prewarm-123', + status: 'ready', + revision: 'sha-repo', + services: ['api'], + workspaceRepos: [ + { + repo: 'example-org/example-repo', + repoUrl: 'https://github.com/example-org/example-repo.git', + branch: 'sample-branch', + revision: 'sha-repo', + mountPath: '/workspace/repos/example-org/example-repo', + primary: true, + }, + ], + serviceRefs: [ + { + name: 'api', + deployId: 2, + repo: 'example-org/example-repo', + branch: 'sample-branch', + workspacePath: '/workspace/repos/example-org/example-repo', + }, + ], + }, + ]), + }); + const service = new AgentPrewarmService( + {} as any, + {} as any, + {} as any, + { + registerQueue: jest.fn().mockReturnValue({ add: jest.fn() }), + } as any + ); + + await expect( + service.getCompatibleReadyPrewarm({ + buildUuid: 'build-123', + requestedServices: ['api'], + revision: 'sha-repo', + workspaceRepos: [ + { + repo: 'example-org/example-repo', + repoUrl: 'https://github.com/example-org/example-repo.git', + branch: 'sample-branch', + revision: 'sha-repo', + mountPath: '/workspace/repos/example-org/example-repo', + primary: true, + }, + ], + requestedServiceRefs: [ + { + name: 'api', + deployId: 3, + repo: 'example-org/example-repo', + branch: 'sample-branch', + workspacePath: '/workspace/repos/example-org/example-repo', + }, + ], + }) + ).resolves.toBeNull(); }); }); @@ -167,7 +491,7 @@ describe('AgentPrewarmService', () => { }, }, }); - (resolveAgentSessionServiceCandidates as jest.Mock).mockReturnValue([ + (resolveAgentSessionServiceCandidatesForBuild as jest.Mock).mockResolvedValue([ { name: 'web', deployId: 1, devConfig: { command: 'pnpm dev' } }, ]); (resolveRequestedAgentSessionServices as jest.Mock).mockReturnValue([ @@ -179,6 +503,87 @@ describe('AgentPrewarmService', () => { expect(plan?.revision).toBe('0123456789abcdef0123456789abcdef01234567'); expect(plan?.revision).not.toBe('1b9337'); }); + + it('builds a prewarm plan for multi-repo service selections', async () => { + const queueManager = { + registerQueue: jest.fn().mockReturnValue({ add: jest.fn() }), + }; + const service = new AgentPrewarmService({} as any, {} as any, {} as any, queueManager as any); + const build = { + kind: 'environment', + namespace: 'env-sample', + pullRequest: { + fullName: 'example-org/example-repo', + branchName: 'sample-branch', + latestCommit: 'sha-repo', + }, + deploys: [], + }; + + (Build.query as jest.Mock).mockReturnValue({ + findOne: jest.fn().mockReturnValue({ + withGraphFetched: jest.fn().mockResolvedValue(build), + }), + }); + (fetchLifecycleConfig as jest.Mock).mockResolvedValue({ + environment: { + agentSession: { + prewarm: { + services: ['api', 'web'], + }, + }, + }, + }); + (resolveAgentSessionServiceCandidatesForBuild as jest.Mock).mockResolvedValue([ + { name: 'web', deployId: 1, devConfig: { command: 'pnpm dev' } }, + { name: 'api', deployId: 2, devConfig: { command: 'pnpm dev' } }, + ]); + (resolveRequestedAgentSessionServices as jest.Mock).mockReturnValue([ + { + name: 'api', + deployId: 2, + devConfig: { command: 'pnpm dev' }, + repo: 'example-org/example-api', + branch: 'api-branch', + revision: 'sha-api', + }, + { + name: 'web', + deployId: 1, + devConfig: { command: 'pnpm dev' }, + repo: 'example-org/example-repo', + branch: 'sample-branch', + revision: 'sha-repo', + }, + ]); + + const plan = await (service as any).resolveBuildPrewarmPlan('build-123'); + + expect(plan?.workspaceRepos).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + repo: 'example-org/example-repo', + branch: 'sample-branch', + mountPath: '/workspace/repos/example-org/example-repo', + primary: true, + }), + expect.objectContaining({ + repo: 'example-org/example-api', + branch: 'api-branch', + mountPath: '/workspace/repos/example-org/example-api', + primary: false, + }), + ]) + ); + expect(plan?.serviceRefs).toEqual( + expect.arrayContaining([ + expect.objectContaining({ name: 'web', deployId: 1, repo: 'example-org/example-repo' }), + expect.objectContaining({ name: 'api', deployId: 2, repo: 'example-org/example-api' }), + ]) + ); + expect(plan?.repo).toBe('example-org/example-repo'); + expect(plan?.revision).toBe('sha-repo'); + }); }); describe('getReadyPrewarmByPvc', () => { diff --git a/src/server/services/__tests__/agentSession.test.ts b/src/server/services/__tests__/agentSession.test.ts index 22d557b1..7d727d91 100644 --- a/src/server/services/__tests__/agentSession.test.ts +++ b/src/server/services/__tests__/agentSession.test.ts @@ -590,11 +590,30 @@ describe('AgentSessionService', () => { await AgentSessionService.createSession(optsWithServices); - expect(mockGetCompatibleReadyPrewarm).toHaveBeenCalledWith({ - buildUuid: 'build-123', - requestedServices: ['web'], - revision: undefined, - }); + expect(mockGetCompatibleReadyPrewarm).toHaveBeenCalledWith( + expect.objectContaining({ + buildUuid: 'build-123', + requestedServices: ['web'], + revision: undefined, + workspaceRepos: [ + expect.objectContaining({ + repo: 'example-org/example-repo', + branch: 'feature/example-session', + mountPath: '/workspace', + primary: true, + }), + ], + requestedServiceRefs: [ + expect.objectContaining({ + name: 'web', + deployId: 1, + repo: 'example-org/example-repo', + branch: 'feature/example-session', + workspacePath: '/workspace', + }), + ], + }) + ); expect(createAgentPvc).not.toHaveBeenCalled(); expect(createSessionWorkspacePod).toHaveBeenCalledWith( expect.objectContaining({ @@ -618,11 +637,30 @@ describe('AgentSessionService', () => { await AgentSessionService.createSession(optsWithServices); - expect(mockGetCompatibleReadyPrewarm).toHaveBeenCalledWith({ - buildUuid: 'build-123', - requestedServices: ['api'], - revision: undefined, - }); + expect(mockGetCompatibleReadyPrewarm).toHaveBeenCalledWith( + expect.objectContaining({ + buildUuid: 'build-123', + requestedServices: ['api'], + revision: undefined, + workspaceRepos: [ + expect.objectContaining({ + repo: 'example-org/example-repo', + branch: 'feature/example-session', + mountPath: '/workspace', + primary: true, + }), + ], + requestedServiceRefs: [ + expect.objectContaining({ + name: 'api', + deployId: 2, + repo: 'example-org/example-repo', + branch: 'feature/example-session', + workspacePath: '/workspace', + }), + ], + }) + ); expect(createAgentPvc).toHaveBeenCalledWith('test-ns', 'agent-pvc-aaaaaaaa', '10Gi', 'build-123'); expect(createSessionWorkspacePod).toHaveBeenCalledWith( expect.objectContaining({ @@ -633,6 +671,89 @@ describe('AgentSessionService', () => { ); }); + it('reuses a compatible multi-repo prewarm when workspace layout matches', async () => { + mockGetCompatibleReadyPrewarm.mockResolvedValue({ + uuid: 'prewarm-2', + pvcName: 'agent-prewarm-pvc-5678', + services: ['api', 'web'], + status: 'ready', + }); + + const optsWithServices: CreateSessionOptions = { + ...baseOpts, + buildUuid: 'build-123', + services: [ + { + name: 'web', + deployId: 1, + devConfig: { image: 'node:20', command: 'pnpm dev' }, + }, + { + name: 'api', + deployId: 2, + devConfig: { image: 'node:20', command: 'pnpm dev' }, + repo: 'org/api', + branch: 'feature/api', + revision: 'sha-api', + }, + ], + }; + + await AgentSessionService.createSession(optsWithServices); + + expect(mockGetCompatibleReadyPrewarm).toHaveBeenCalledWith( + expect.objectContaining({ + buildUuid: 'build-123', + requestedServices: ['web', 'api'], + workspaceRepos: expect.arrayContaining([ + expect.objectContaining({ + repo: 'example-org/example-repo', + branch: 'feature/example-session', + mountPath: '/workspace/repos/example-org/example-repo', + primary: true, + }), + expect.objectContaining({ + repo: 'org/api', + branch: 'feature/api', + revision: 'sha-api', + mountPath: '/workspace/repos/org/api', + primary: false, + }), + ]), + requestedServiceRefs: expect.arrayContaining([ + expect.objectContaining({ + name: 'web', + deployId: 1, + repo: 'example-org/example-repo', + branch: 'feature/example-session', + workspacePath: '/workspace/repos/example-org/example-repo', + }), + expect.objectContaining({ + name: 'api', + deployId: 2, + repo: 'org/api', + branch: 'feature/api', + workspacePath: '/workspace/repos/org/api', + }), + ]), + }) + ); + expect(createAgentPvc).not.toHaveBeenCalled(); + expect(createSessionWorkspacePod).toHaveBeenCalledWith( + expect.objectContaining({ + pvcName: 'agent-prewarm-pvc-5678', + skipWorkspaceBootstrap: true, + workspaceRepos: expect.arrayContaining([ + expect.objectContaining({ + repo: 'example-org/example-repo', + mountPath: '/workspace/repos/example-org/example-repo', + }), + expect.objectContaining({ repo: 'org/api', mountPath: '/workspace/repos/org/api' }), + ]), + }) + ); + }); + it('passes resolved agent-session resources through to pod creation when provided', async () => { const optsWithResources: CreateSessionOptions = { ...baseOpts, diff --git a/src/server/services/agentPrewarm.ts b/src/server/services/agentPrewarm.ts index 5fde7ee8..41eb0268 100644 --- a/src/server/services/agentPrewarm.ts +++ b/src/server/services/agentPrewarm.ts @@ -16,6 +16,7 @@ import 'server/lib/dependencies'; import * as k8s from '@kubernetes/client-node'; +import { createHash } from 'crypto'; import { v4 as uuid } from 'uuid'; import BaseService from './_service'; import AgentPrewarm from 'server/models/AgentPrewarm'; @@ -48,7 +49,10 @@ import type { DevConfig } from 'server/models/yaml/YamlService'; import { BuildKind } from 'shared/constants'; import { QUEUE_NAMES } from 'shared/config'; import GlobalConfigService from './globalConfig'; -import { resolveAgentSessionServiceCandidates, resolveRequestedAgentSessionServices } from './agentSessionCandidates'; +import { + resolveAgentSessionServiceCandidatesForBuild, + resolveRequestedAgentSessionServices, +} from './agentSessionCandidates'; const logger = () => getLogger(); const AGENT_PREWARM_ERROR_MESSAGE_MAX_LENGTH = 4000; @@ -104,6 +108,142 @@ export function canReusePrewarm(prewarmServices: string[], requestedServices: st return requested.every((serviceName) => available.has(serviceName)); } +function normalizeWorkspaceRepoSignature(repo: AgentSessionWorkspaceRepo): string { + return JSON.stringify({ + repo: repo.repo, + repoUrl: repo.repoUrl, + branch: repo.branch, + revision: repo.revision || null, + mountPath: repo.mountPath, + primary: Boolean(repo.primary), + }); +} + +function hasExactWorkspaceRepoMatch( + prewarmWorkspaceRepos: AgentSessionWorkspaceRepo[], + requestedWorkspaceRepos: AgentSessionWorkspaceRepo[] +): boolean { + if (prewarmWorkspaceRepos.length !== requestedWorkspaceRepos.length) { + return false; + } + + const requestedRepoSet = new Set(requestedWorkspaceRepos.map((repo) => normalizeWorkspaceRepoSignature(repo))); + return ( + requestedRepoSet.size === prewarmWorkspaceRepos.length && + prewarmWorkspaceRepos.every((repo) => requestedRepoSet.has(normalizeWorkspaceRepoSignature(repo))) + ); +} + +function hasServiceRefCoverage( + prewarmServiceRefs: AgentSessionSelectedService[], + requestedServiceRefs: AgentSessionSelectedService[] +): boolean { + if (requestedServiceRefs.length === 0) { + return true; + } + + const availableDeployIds = new Set(prewarmServiceRefs.map((serviceRef) => serviceRef.deployId)); + return requestedServiceRefs.every((serviceRef) => availableDeployIds.has(serviceRef.deployId)); +} + +function normalizeSelectedServiceSignature(serviceRef: AgentSessionSelectedService): string { + return JSON.stringify({ + name: serviceRef.name, + deployId: serviceRef.deployId, + repo: serviceRef.repo || null, + branch: serviceRef.branch || null, + revision: serviceRef.revision || null, + resourceName: serviceRef.resourceName || null, + workspacePath: serviceRef.workspacePath || null, + workDir: serviceRef.workDir || null, + }); +} + +function hasExactServiceRefMatch( + prewarmServiceRefs: AgentSessionSelectedService[], + requestedServiceRefs: AgentSessionSelectedService[] +): boolean { + if (prewarmServiceRefs.length !== requestedServiceRefs.length) { + return false; + } + + const requestedServiceRefSet = new Set( + requestedServiceRefs.map((serviceRef) => normalizeSelectedServiceSignature(serviceRef)) + ); + return ( + requestedServiceRefSet.size === prewarmServiceRefs.length && + prewarmServiceRefs.every((serviceRef) => requestedServiceRefSet.has(normalizeSelectedServiceSignature(serviceRef))) + ); +} + +type PrewarmMatchOptions = { + revision?: string; + requestedServices: string[]; + workspaceRepos?: AgentSessionWorkspaceRepo[]; + requestedServiceRefs?: AgentSessionSelectedService[]; + serviceMatchMode: 'coverage' | 'exact'; + serviceRefMatchMode?: 'coverage' | 'exact'; +}; + +function matchesPrewarm( + prewarm: Pick, + opts: PrewarmMatchOptions +): boolean { + const sameRevision = !opts.revision || !prewarm.revision || prewarm.revision === opts.revision; + if (!sameRevision) { + return false; + } + + if (opts.workspaceRepos?.length && !hasExactWorkspaceRepoMatch(prewarm.workspaceRepos || [], opts.workspaceRepos)) { + return false; + } + + if (opts.requestedServiceRefs?.length) { + const serviceRefMatches = + opts.serviceRefMatchMode === 'exact' + ? hasExactServiceRefMatch(prewarm.serviceRefs || [], opts.requestedServiceRefs) + : hasServiceRefCoverage(prewarm.serviceRefs || [], opts.requestedServiceRefs); + if (!serviceRefMatches) { + return false; + } + } + + return opts.serviceMatchMode === 'exact' + ? hasExactServiceMatch(prewarm.services || [], opts.requestedServices) + : canReusePrewarm(prewarm.services || [], opts.requestedServices); +} + +function matchesBuildPrewarmPlan( + prewarm: Pick, + plan: Pick +): boolean { + return matchesPrewarm(prewarm, { + revision: plan.revision, + requestedServices: plan.configuredServiceNames, + workspaceRepos: plan.workspaceRepos, + requestedServiceRefs: plan.serviceRefs, + serviceMatchMode: 'exact', + serviceRefMatchMode: 'exact', + }); +} + +function buildPrewarmIdentityToken(plan: { + workspaceRepos?: AgentSessionWorkspaceRepo[]; + serviceRefs?: AgentSessionSelectedService[]; +}): string { + const normalizedWorkspaceRepos = (plan.workspaceRepos || []) + .map((repo) => normalizeWorkspaceRepoSignature(repo)) + .sort(); + const normalizedServiceRefs = (plan.serviceRefs || []) + .map((serviceRef) => normalizeSelectedServiceSignature(serviceRef)) + .sort(); + + return createHash('sha1') + .update(JSON.stringify({ workspaceRepos: normalizedWorkspaceRepos, serviceRefs: normalizedServiceRefs })) + .digest('hex') + .slice(0, 12); +} + function truncateErrorMessage(message: string): string { if (message.length <= AGENT_PREWARM_ERROR_MESSAGE_MAX_LENGTH) { return message; @@ -134,6 +274,8 @@ export default class AgentPrewarmService extends BaseService { buildUuid: string; requestedServices: string[]; revision?: string; + workspaceRepos?: AgentSessionWorkspaceRepo[]; + requestedServiceRefs?: AgentSessionSelectedService[]; }): Promise { const prewarms = await AgentPrewarm.query() .where({ @@ -143,13 +285,16 @@ export default class AgentPrewarmService extends BaseService { .orderBy('updatedAt', 'desc'); return ( - prewarms.find((prewarm) => { - if (params.revision && prewarm.revision && prewarm.revision !== params.revision) { - return false; - } - - return canReusePrewarm(prewarm.services || [], params.requestedServices); - }) || null + prewarms.find((prewarm) => + matchesPrewarm(prewarm, { + revision: params.revision, + requestedServices: params.requestedServices, + workspaceRepos: params.workspaceRepos, + requestedServiceRefs: params.requestedServiceRefs, + serviceMatchMode: 'coverage', + serviceRefMatchMode: 'coverage', + }) + ) || null ); } @@ -178,10 +323,7 @@ export default class AgentPrewarmService extends BaseService { .where({ buildUuid: plan.buildUuid }) .whereIn('status', ['queued', 'running', 'ready']) .orderBy('updatedAt', 'desc'); - const matchingPrewarm = activePrewarms.find((prewarm) => { - const sameRevision = !plan.revision || !prewarm.revision || prewarm.revision === plan.revision; - return sameRevision && hasExactServiceMatch(prewarm.services || [], plan.configuredServiceNames); - }); + const matchingPrewarm = activePrewarms.find((prewarm) => matchesBuildPrewarmPlan(prewarm, plan)); if (matchingPrewarm) { return false; @@ -194,7 +336,9 @@ export default class AgentPrewarmService extends BaseService { ...extractContextForQueue(), }, { - jobId: `agent-prewarm:${plan.buildUuid}:${plan.revision || 'head'}:${plan.configuredServiceNames.join(',')}`, + jobId: `agent-prewarm:${plan.buildUuid}:${plan.revision || 'head'}:${plan.configuredServiceNames.join( + ',' + )}:${buildPrewarmIdentityToken(plan)}`, } ); @@ -217,10 +361,7 @@ export default class AgentPrewarmService extends BaseService { .where({ buildUuid: plan.buildUuid }) .whereIn('status', ['running', 'ready']) .orderBy('updatedAt', 'desc'); - const matchingPrewarm = existingPrewarm.find((prewarm) => { - const sameRevision = !plan.revision || !prewarm.revision || prewarm.revision === plan.revision; - return sameRevision && hasExactServiceMatch(prewarm.services || [], plan.configuredServiceNames); - }); + const matchingPrewarm = existingPrewarm.find((prewarm) => matchesBuildPrewarmPlan(prewarm, plan)); if (matchingPrewarm) { return matchingPrewarm; } @@ -374,7 +515,7 @@ export default class AgentPrewarmService extends BaseService { private async resolveBuildPrewarmPlan(buildUuid: string): Promise { const build = await Build.query() .findOne({ uuid: buildUuid }) - .withGraphFetched('[pullRequest, deploys.[deployable]]'); + .withGraphFetched('[pullRequest, deploys.[deployable, repository, service]]'); if ( !build || build.kind !== BuildKind.ENVIRONMENT || @@ -395,14 +536,14 @@ export default class AgentPrewarmService extends BaseService { ); } - private resolvePrewarmPlanFromConfig( + private async resolvePrewarmPlanFromConfig( buildUuid: string, namespace: string, repositoryFullName: string, revision: string | undefined, lifecycleConfig: LifecycleConfig | null, build: Build - ): ResolvedBuildPrewarmPlan | null { + ): Promise { const configuredServiceNames = normalizeServiceNames( lifecycleConfig?.environment?.agentSession?.prewarm?.services || [] ); @@ -410,7 +551,7 @@ export default class AgentPrewarmService extends BaseService { return null; } - const candidates = resolveAgentSessionServiceCandidates(build.deploys || [], lifecycleConfig as LifecycleConfig); + const candidates = await resolveAgentSessionServiceCandidatesForBuild(build); const services = resolveRequestedAgentSessionServices(candidates, configuredServiceNames).map((service) => ({ name: service.name, deployId: service.deployId, @@ -424,13 +565,19 @@ export default class AgentPrewarmService extends BaseService { workspaceRepos, services: resolvedServices, selectedServices, - } = resolveAgentSessionServicePlan({}, services); - if (workspaceRepos.length !== 1) { + } = resolveAgentSessionServicePlan( + { + repoUrl: `https://github.com/${repositoryFullName}.git`, + branch: build.pullRequest?.branchName, + revision, + }, + services + ); + const workspaceRepo = workspaceRepos.find((repo) => repo.primary) || workspaceRepos[0]; + if (!workspaceRepo) { return null; } - const [workspaceRepo] = workspaceRepos; - return { buildUuid, namespace, diff --git a/src/server/services/agentSession.ts b/src/server/services/agentSession.ts index d8c82586..0e0f4111 100644 --- a/src/server/services/agentSession.ts +++ b/src/server/services/agentSession.ts @@ -584,7 +584,13 @@ async function deleteAgentRuntimeResources( ]); } -async function resolveCompatiblePrewarm(buildUuid: string | undefined, requestedServices: string[], revision?: string) { +async function resolveCompatiblePrewarm( + buildUuid: string | undefined, + requestedServices: string[], + revision?: string, + workspaceRepos?: AgentSessionWorkspaceRepo[], + requestedServiceRefs?: AgentSessionSelectedService[] +) { if (!buildUuid) { return null; } @@ -593,6 +599,8 @@ async function resolveCompatiblePrewarm(buildUuid: string | undefined, requested buildUuid, requestedServices, revision, + workspaceRepos, + requestedServiceRefs, }); } @@ -888,13 +896,13 @@ export default class AgentSessionService { opts.userIdentity || null ) : Promise.resolve([]), - workspaceRepos.length === 1 - ? resolveCompatiblePrewarm( - opts.buildUuid, - resolvedServiceNames, - primaryWorkspaceRepo?.revision || opts.revision - ) - : Promise.resolve(null), + resolveCompatiblePrewarm( + opts.buildUuid, + resolvedServiceNames, + primaryWorkspaceRepo?.revision || opts.revision, + workspaceRepos, + selectedServices + ), resolveForwardedAgentEnv(resolvedServices, opts.namespace, sessionUuid, opts.buildUuid), ]); const sessionPodMcpConfigJson = serializeSessionWorkspaceGatewayServers(sessionPodServers); @@ -954,6 +962,8 @@ export default class AgentSessionService { ), ensureAgentSessionServiceAccount(opts.namespace), isGvisorAvailable(), + createSessionWorkspaceService(opts.namespace, podName, opts.buildUuid), + ensureAgentNetworkPolicy(opts.namespace), ]); const infraSetupMs = elapsedMs(infraSetupStartedAt); @@ -1030,10 +1040,6 @@ export default class AgentSessionService { persistedDevModeDeployIds.push(...enabledServices.map((service) => service.deployId)); const finalizeStartedAt = Date.now(); - await Promise.all([ - createSessionWorkspaceService(opts.namespace, podName, opts.buildUuid), - ensureAgentNetworkPolicy(opts.namespace), - ]); await Promise.all([ redis.setex( `${SESSION_REDIS_PREFIX}${sessionUuid}`,