diff --git a/apps/supervisor/src/index.ts b/apps/supervisor/src/index.ts index 0e274b3039..024affde32 100644 --- a/apps/supervisor/src/index.ts +++ b/apps/supervisor/src/index.ts @@ -267,6 +267,7 @@ class ManagedSupervisor { snapshotId: message.snapshot.id, snapshotFriendlyId: message.snapshot.friendlyId, placementTags: message.placementTags, + envVars: message.envVars, }); // Disabled for now diff --git a/apps/supervisor/src/workloadManager/docker.ts b/apps/supervisor/src/workloadManager/docker.ts index d6651d325a..d2a3b70482 100644 --- a/apps/supervisor/src/workloadManager/docker.ts +++ b/apps/supervisor/src/workloadManager/docker.ts @@ -110,6 +110,12 @@ export class DockerWorkloadManager implements WorkloadManager { }); } + if (opts.envVars) { + Object.entries(opts.envVars).forEach(([key, value]) => { + envVars.push(`${key}=${value}`); + }); + } + const hostConfig: Docker.HostConfig = { AutoRemove: !!this.opts.dockerAutoremove, }; diff --git a/apps/supervisor/src/workloadManager/kubernetes.ts b/apps/supervisor/src/workloadManager/kubernetes.ts index 96fbd7a274..de4cbc1199 100644 --- a/apps/supervisor/src/workloadManager/kubernetes.ts +++ b/apps/supervisor/src/workloadManager/kubernetes.ts @@ -225,6 +225,12 @@ export class KubernetesWorkloadManager implements WorkloadManager { value: value, })) : []), + ...(opts.envVars + ? Object.entries(opts.envVars).map(([key, value]) => ({ + name: key, + value: value, + })) + : []), ], }, ], diff --git a/apps/supervisor/src/workloadManager/types.ts b/apps/supervisor/src/workloadManager/types.ts index 90b6195779..094f8a133b 100644 --- a/apps/supervisor/src/workloadManager/types.ts +++ b/apps/supervisor/src/workloadManager/types.ts @@ -24,6 +24,7 @@ export interface WorkloadManagerCreateOptions { nextAttemptNumber?: number; dequeuedAt: Date; placementTags?: PlacementTag[]; + envVars?: Record; // identifiers envId: string; envType: EnvironmentType; diff --git a/apps/webapp/app/v3/services/worker/workerGroupTokenService.server.ts b/apps/webapp/app/v3/services/worker/workerGroupTokenService.server.ts index befe2a0a89..c997b6fc5e 100644 --- a/apps/webapp/app/v3/services/worker/workerGroupTokenService.server.ts +++ b/apps/webapp/app/v3/services/worker/workerGroupTokenService.server.ts @@ -370,12 +370,44 @@ export class AuthenticatedWorkerInstance extends WithRunEngine { } async dequeue({ runnerId }: { runnerId?: string }): Promise { - return await this._engine.dequeueFromWorkerQueue({ + const messages = await this._engine.dequeueFromWorkerQueue({ consumerId: this.workerInstanceId, workerQueue: this.masterQueue, workerId: this.workerInstanceId, runnerId, }); + + // Fetch and inject environment variables for each message + const messagesWithEnvVars = await Promise.all( + messages.map(async (message) => { + const defaultMachinePreset = machinePresetFromName(defaultMachine); + + const environment = await this._prisma.runtimeEnvironment.findFirst({ + where: { + id: message.environment.id, + }, + include: { + parentEnvironment: true, + }, + }); + + const envVars = environment + ? await this.getEnvVars( + environment, + message.run.id, + message.run.machine ?? defaultMachinePreset, + environment.parentEnvironment ?? undefined + ) + : {}; + + return { + ...message, + envVars, + }; + }) + ); + + return messagesWithEnvVars; } async heartbeatWorkerInstance() { diff --git a/packages/core/src/v3/isomorphic/friendlyId.ts b/packages/core/src/v3/isomorphic/friendlyId.ts index 90fa31bd57..f910d30e0b 100644 --- a/packages/core/src/v3/isomorphic/friendlyId.ts +++ b/packages/core/src/v3/isomorphic/friendlyId.ts @@ -36,6 +36,11 @@ export function fromFriendlyId(friendlyId: string, expectedEntityName?: string): const parts = friendlyId.split("_"); + // If there's no underscore, assume it's already an internal ID + if (parts.length === 1) { + return friendlyId; + } + if (parts.length !== 2) { throw new Error("Invalid friendly ID format"); } diff --git a/packages/core/src/v3/schemas/runEngine.ts b/packages/core/src/v3/schemas/runEngine.ts index 376a8522de..4dcd997338 100644 --- a/packages/core/src/v3/schemas/runEngine.ts +++ b/packages/core/src/v3/schemas/runEngine.ts @@ -270,5 +270,6 @@ export const DequeuedMessage = z.object({ id: z.string(), }), placementTags: z.array(PlacementTag).optional(), + envVars: z.record(z.string()).optional(), }); export type DequeuedMessage = z.infer; diff --git a/packages/trigger-sdk/package.json b/packages/trigger-sdk/package.json index 86a0a8a5f2..1311344d50 100644 --- a/packages/trigger-sdk/package.json +++ b/packages/trigger-sdk/package.json @@ -128,4 +128,4 @@ "main": "./dist/commonjs/v3/index.js", "types": "./dist/commonjs/v3/index.d.ts", "module": "./dist/esm/v3/index.js" -} \ No newline at end of file +}