Skip to content

Commit f62ff10

Browse files
vmelikyanvmelikyan
committed
add prewarm for agent session
1 parent 42219c3 commit f62ff10

21 files changed

Lines changed: 1716 additions & 87 deletions

File tree

‎docs/schema/yaml/1.0.0.yaml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,12 @@ environment:
9797
# @param environment.agentSession.resources.editor.limits
9898
limits:
9999

100+
# @param environment.agentSession.prewarm
101+
prewarm:
102+
# @param environment.agentSession.prewarm.services
103+
services:
104+
# @param environment.agentSession.prewarm.services[]
105+
- ''
100106
# @section services
101107
services:
102108
# @param services[]

‎src/server/db/migrations/016_add_agent_prewarms.ts‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
/**
2+
* Copyright 2026 GoodRx, Inc.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
import { Knex } from 'knex';
18+
19+
export async function up(knex: Knex): Promise<void> {
20+
await knex.schema.createTable('agent_prewarms', (table) => {
21+
table.increments('id').primary();
22+
table.uuid('uuid').notNullable().unique().defaultTo(knex.raw('gen_random_uuid()'));
23+
table.string('buildUuid').notNullable().index();
24+
table.string('namespace').notNullable();
25+
table.string('repo').nullable();
26+
table.string('branch').nullable();
27+
table.string('revision').nullable().index();
28+
table.string('pvcName').notNullable();
29+
table.string('jobName').notNullable();
30+
table.string('status').notNullable().defaultTo('queued').index();
31+
table.jsonb('services').notNullable().defaultTo('[]');
32+
table.text('errorMessage').nullable();
33+
table.timestamp('completedAt').nullable();
34+
table.timestamp('createdAt').notNullable().defaultTo(knex.fn.now());
35+
table.timestamp('updatedAt').notNullable().defaultTo(knex.fn.now());
36+
table.unique(['buildUuid', 'revision', 'pvcName']);
37+
table.index(['buildUuid', 'status', 'updatedAt']);
38+
});
39+
}
40+
41+
export async function down(knex: Knex): Promise<void> {
42+
await knex.schema.dropTableIfExists('agent_prewarms');
43+
}

‎src/server/jobs/agentSessionPrewarm.ts‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
/**
2+
* Copyright 2026 GoodRx, Inc.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
import { Job } from 'bullmq';
18+
import { getLogger } from 'server/lib/logger';
19+
import AgentPrewarmService, { AgentPrewarmQueueJob } from 'server/services/agentPrewarm';
20+
21+
const logger = getLogger();
22+
23+
export async function processAgentSessionPrewarm(job: Job<AgentPrewarmQueueJob>): Promise<void> {
24+
const { buildUuid } = job.data;
25+
await new AgentPrewarmService().prepareBuildPrewarm(buildUuid).catch((error) => {
26+
logger.error({ error, buildUuid }, 'Agent prewarm job failed');
27+
throw error;
28+
});
29+
}

‎src/server/jobs/index.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ import RedisClient from 'server/lib/redisClient';
2121
import QueueManager from 'server/lib/queueManager';
2222
import { MAX_GITHUB_API_REQUEST, GITHUB_API_REQUEST_INTERVAL, QUEUE_NAMES } from 'shared/config';
2323
import { processAgentSessionCleanup } from './agentSessionCleanup';
24+
import { processAgentSessionPrewarm } from './agentSessionPrewarm';
2425
import { processAgentSandboxSessionLaunch } from './agentSandboxSessionLaunch';
2526

2627
let isBootstrapped = false;
@@ -110,6 +111,11 @@ export default function bootstrapJobs(services: IServices) {
110111
concurrency: 1,
111112
});
112113

114+
queueManager.registerWorker(QUEUE_NAMES.AGENT_SESSION_PREWARM, processAgentSessionPrewarm, {
115+
connection: redisClient.getConnection(),
116+
concurrency: 2,
117+
});
118+
113119
queueManager.registerWorker(QUEUE_NAMES.AGENT_SANDBOX_SESSION_LAUNCH, processAgentSandboxSessionLaunch, {
114120
connection: redisClient.getConnection(),
115121
concurrency: 5,

‎src/server/lib/agentSession/podFactory.ts‎

Lines changed: 65 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,7 @@ export interface AgentPodOpts {
138138
timeoutMs: number;
139139
pollMs: number;
140140
};
141+
skipWorkspaceBootstrap?: boolean;
141142
resources?: {
142143
agent?: k8s.V1ResourceRequirements;
143144
editor?: k8s.V1ResourceRequirements;
@@ -251,6 +252,7 @@ export function buildAgentPodSpec(opts: AgentPodOpts): k8s.V1Pod {
251252
claudePrAttribution,
252253
useGvisor,
253254
userIdentity,
255+
skipWorkspaceBootstrap,
254256
} = opts;
255257

256258
const initScriptOpts: InitScriptOpts = {
@@ -319,65 +321,69 @@ export function buildAgentPodSpec(opts: AgentPodOpts): k8s.V1Pod {
319321
type: 'RuntimeDefault',
320322
},
321323
},
322-
initContainers: [
323-
{
324-
name: 'prepare-workspace',
325-
image,
326-
imagePullPolicy: 'IfNotPresent',
327-
command: ['sh', '-c', `mkdir -p "${AGENT_WORKSPACE_VOLUME_ROOT}/${AGENT_WORKSPACE_SUBPATH}"`],
328-
resources,
329-
securityContext: {
330-
...securityContext,
331-
readOnlyRootFilesystem: false,
332-
},
333-
volumeMounts: [
334-
{
335-
name: 'workspace',
336-
mountPath: AGENT_WORKSPACE_VOLUME_ROOT,
337-
},
338-
{
339-
name: 'tmp',
340-
mountPath: '/tmp',
341-
},
342-
],
343-
env: [
344-
{ name: 'TMPDIR', value: '/tmp' },
345-
{ name: 'TMP', value: '/tmp' },
346-
{ name: 'TEMP', value: '/tmp' },
347-
],
348-
},
349-
{
350-
name: 'init-workspace',
351-
image,
352-
imagePullPolicy: 'IfNotPresent',
353-
command: ['sh', '-c', initScript],
354-
resources,
355-
securityContext: {
356-
...securityContext,
357-
readOnlyRootFilesystem: false,
358-
},
359-
volumeMounts: [
360-
workspaceVolumeMount,
361-
{
362-
name: 'claude-config',
363-
mountPath: '/home/claude/.claude',
364-
},
365-
{
366-
name: 'tmp',
367-
mountPath: '/tmp',
368-
},
369-
],
370-
env: [
371-
{ name: 'HOME', value: '/home/claude/.claude' },
372-
{ name: 'TMPDIR', value: '/tmp' },
373-
{ name: 'TMP', value: '/tmp' },
374-
{ name: 'TEMP', value: '/tmp' },
375-
...forwardedAgentSecretEnv,
376-
...githubTokenEnv,
377-
...userEnv,
378-
],
379-
},
380-
],
324+
...(skipWorkspaceBootstrap
325+
? {}
326+
: {
327+
initContainers: [
328+
{
329+
name: 'prepare-workspace',
330+
image,
331+
imagePullPolicy: 'IfNotPresent',
332+
command: ['sh', '-c', `mkdir -p "${AGENT_WORKSPACE_VOLUME_ROOT}/${AGENT_WORKSPACE_SUBPATH}"`],
333+
resources,
334+
securityContext: {
335+
...securityContext,
336+
readOnlyRootFilesystem: false,
337+
},
338+
volumeMounts: [
339+
{
340+
name: 'workspace',
341+
mountPath: AGENT_WORKSPACE_VOLUME_ROOT,
342+
},
343+
{
344+
name: 'tmp',
345+
mountPath: '/tmp',
346+
},
347+
],
348+
env: [
349+
{ name: 'TMPDIR', value: '/tmp' },
350+
{ name: 'TMP', value: '/tmp' },
351+
{ name: 'TEMP', value: '/tmp' },
352+
],
353+
},
354+
{
355+
name: 'init-workspace',
356+
image,
357+
imagePullPolicy: 'IfNotPresent',
358+
command: ['sh', '-c', initScript],
359+
resources,
360+
securityContext: {
361+
...securityContext,
362+
readOnlyRootFilesystem: false,
363+
},
364+
volumeMounts: [
365+
workspaceVolumeMount,
366+
{
367+
name: 'claude-config',
368+
mountPath: '/home/claude/.claude',
369+
},
370+
{
371+
name: 'tmp',
372+
mountPath: '/tmp',
373+
},
374+
],
375+
env: [
376+
{ name: 'HOME', value: '/home/claude/.claude' },
377+
{ name: 'TMPDIR', value: '/tmp' },
378+
{ name: 'TMP', value: '/tmp' },
379+
{ name: 'TEMP', value: '/tmp' },
380+
...forwardedAgentSecretEnv,
381+
...githubTokenEnv,
382+
...userEnv,
383+
],
384+
},
385+
],
386+
}),
381387
containers: [
382388
{
383389
name: 'agent',

0 commit comments

Comments
 (0)