diff --git a/apps/webapp/app/v3/marqs/devQueueConsumer.server.ts b/apps/webapp/app/v3/marqs/devQueueConsumer.server.ts index 35c313004..fc1a5ed83 100644 --- a/apps/webapp/app/v3/marqs/devQueueConsumer.server.ts +++ b/apps/webapp/app/v3/marqs/devQueueConsumer.server.ts @@ -413,6 +413,7 @@ export class DevQueueConsumer { data: { lockedAt: new Date(), lockedById: backgroundTask.id, + status: "EXECUTING", lockedToVersionId: backgroundWorker.id, }, include: { diff --git a/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts b/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts index 924430702..0f8f5f6c0 100644 --- a/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts +++ b/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts @@ -26,7 +26,11 @@ import { marqs, sanitizeQueueName } from "~/v3/marqs/index.server"; import { EnvironmentVariablesRepository } from "../environmentVariables/environmentVariablesRepository.server"; import { generateFriendlyId } from "../friendlyIdentifiers"; import { socketIo } from "../handleSocketIo.server"; -import { findCurrentWorkerDeployment } from "../models/workerDeployment.server"; +import { + findCurrentWorkerDeployment, + getWorkerDeploymentFromWorker, + getWorkerDeploymentFromWorkerTask, +} from "../models/workerDeployment.server"; import { RestoreCheckpointService } from "../services/restoreCheckpoint.server"; import { SEMINTATTRS_FORCE_RECORDING, tracer } from "../tracer.server"; import { CrashTaskRunService } from "../services/crashTaskRun.server"; @@ -317,11 +321,11 @@ export class SharedQueueConsumer { return; } - const deployment = existingTaskRun.lockedToVersion?.deployment - ? { - ...existingTaskRun.lockedToVersion.deployment, - worker: existingTaskRun.lockedToVersion, - } + // Check if the task run is locked to a specific worker, if not, use the current worker deployment + const deployment = existingTaskRun.lockedById + ? await getWorkerDeploymentFromWorkerTask(existingTaskRun.lockedById) + : existingTaskRun.lockedToVersionId + ? await getWorkerDeploymentFromWorker(existingTaskRun.lockedToVersionId) : await findCurrentWorkerDeployment(existingTaskRun.runtimeEnvironmentId); if (!deployment || !deployment.worker) { diff --git a/apps/webapp/app/v3/models/workerDeployment.server.ts b/apps/webapp/app/v3/models/workerDeployment.server.ts index e9e19c1b5..211a20f5c 100644 --- a/apps/webapp/app/v3/models/workerDeployment.server.ts +++ b/apps/webapp/app/v3/models/workerDeployment.server.ts @@ -1,16 +1,30 @@ import type { Prettify } from "@trigger.dev/core"; import { CURRENT_DEPLOYMENT_LABEL } from "~/consts"; -import { prisma } from "~/db.server"; +import { Prisma, prisma } from "~/db.server"; -export type CurrentWorkerDeployment = Prettify>>>; +export type CurrentWorkerDeployment = Prettify< + NonNullable>> +>; -export async function findCurrentWorkerDeployment(environmentId: string) { +type WorkerDeploymentWithWorkerTasks = Prisma.WorkerDeploymentGetPayload<{ + include: { + worker: { + include: { + tasks: true; + }; + }; + }; +}>; + +export async function findCurrentWorkerDeployment( + environmentId: string +): Promise { const promotion = await prisma.workerDeploymentPromotion.findUnique({ where: { environmentId_label: { environmentId, label: CURRENT_DEPLOYMENT_LABEL, - } + }, }, include: { deployment: { @@ -20,10 +34,64 @@ export async function findCurrentWorkerDeployment(environmentId: string) { tasks: true, }, }, - } - } - } + }, + }, + }, }); return promotion?.deployment; -} \ No newline at end of file +} + +export async function getWorkerDeploymentFromWorker( + workerId: string +): Promise { + const worker = await prisma.backgroundWorker.findUnique({ + where: { + id: workerId, + }, + include: { + deployment: true, + tasks: true, + }, + }); + + if (!worker?.deployment) { + return; + } + + const { deployment, ...workerWithoutDeployment } = worker; + + return { + ...deployment, + worker: workerWithoutDeployment, + }; +} + +export async function getWorkerDeploymentFromWorkerTask( + workerTaskId: string +): Promise { + const workerTask = await prisma.backgroundWorkerTask.findUnique({ + where: { + id: workerTaskId, + }, + include: { + worker: { + include: { + deployment: true, + tasks: true, + }, + }, + }, + }); + + if (!workerTask?.worker.deployment) { + return; + } + + const { deployment, ...workerWithoutDeployment } = workerTask.worker; + + return { + ...deployment, + worker: workerWithoutDeployment, + }; +}