import type { Context, Span, SpanOptions } from "@opentelemetry/api";
import { ROOT_CONTEXT, SpanKind, SpanStatusCode, context, trace } from "@opentelemetry/api";
import type {
  AckCallbackResult,
  MachinePreset,
  V3ProdTaskRunExecution,
  V3ProdTaskRunExecutionPayload,
  TaskRunError,
  TaskRunExecution,
  TaskRunExecutionLazyAttemptPayload,
  TaskRunExecutionResult,
  TaskRunFailedExecutionResult,
  TaskRunSuccessfulExecutionResult,
  serverWebsocketMessages,
} from "@trigger.dev/core/v3";
import { TaskRunErrorCodes, parsePacket, SemanticInternalAttributes } from "@trigger.dev/core/v3";
import type { ZodMessageSender } from "@trigger.dev/core/v3/zodMessageHandler";
import type {
  BackgroundWorker,
  BackgroundWorkerTask,
  Prisma,
  TaskRunStatus,
} from "@trigger.dev/database";
import { z } from "zod";
import { $replica, prisma } from "~/db.server";
import { env } from "~/env.server";
import { findEnvironmentById } from "~/models/runtimeEnvironment.server";
import { findQueueInEnvironment, sanitizeQueueName } from "~/models/taskQueue.server";
import { generateJWTTokenForEnvironment } from "~/services/apiAuth.server";
import { logger } from "~/services/logger.server";
import { singleton } from "~/utils/singleton";
import { marqs } from "~/v3/marqs/index.server";
import type { RuntimeEnvironmentForEnvRepo } from "../environmentVariables/environmentVariablesRepository.server";
import {
  RuntimeEnvironmentForEnvRepoPayload,
  resolveVariablesForEnvironment,
} from "../environmentVariables/environmentVariablesRepository.server";
import type { EnvironmentVariable } from "../environmentVariables/repository";
import { FailedTaskRunService } from "../failedTaskRun.server";
import { generateFriendlyId } from "../friendlyIdentifiers";
import { socketIo } from "../handleSocketIo.server";
import { machinePresetFromConfig, machinePresetFromRun } from "../machinePresets.server";
import {
  findCurrentWorkerDeployment,
  getWorkerDeploymentFromWorker,
  getWorkerDeploymentFromWorkerTask,
} from "../models/workerDeployment.server";
import { CrashTaskRunService } from "../services/crashTaskRun.server";
import { CreateTaskRunAttemptService } from "../services/createTaskRunAttempt.server";
import { RestoreCheckpointService } from "../services/restoreCheckpoint.server";
import {
  FINAL_ATTEMPT_STATUSES,
  FINAL_RUN_STATUSES,
  isFinalAttemptStatus,
  isFinalRunStatus,
} from "../taskStatus";
import { tracer } from "../tracer.server";
import { getMaxDuration } from "../utils/maxDuration";
import type { MessagePayload } from "./types";

const WithTraceContext = z.object({
  traceparent: z.string().optional(),
  tracestate: z.string().optional(),
});

export const SharedQueueExecuteMessageBody = WithTraceContext.extend({
  type: z.literal("EXECUTE"),
  taskIdentifier: z.string(),
  checkpointEventId: z.string().optional(),
  retryCheckpointsDisabled: z.boolean().optional(),
});

export type SharedQueueExecuteMessageBody = z.infer<typeof SharedQueueExecuteMessageBody>;

export const SharedQueueResumeMessageBody = WithTraceContext.extend({
  type: z.literal("RESUME"),
  completedAttemptIds: z.string().array(),
  resumableAttemptId: z.string(),
  checkpointEventId: z.string().optional(),
});

export type SharedQueueResumeMessageBody = z.infer<typeof SharedQueueResumeMessageBody>;

export const SharedQueueResumeAfterDurationMessageBody = WithTraceContext.extend({
  type: z.literal("RESUME_AFTER_DURATION"),
  resumableAttemptId: z.string(),
  checkpointEventId: z.string(),
});

export type SharedQueueResumeAfterDurationMessageBody = z.infer<
  typeof SharedQueueResumeAfterDurationMessageBody
>;

export const SharedQueueFailMessageBody = WithTraceContext.extend({
  type: z.literal("FAIL"),
  reason: z.string(),
});

export type SharedQueueFailMessageBody = z.infer<typeof SharedQueueFailMessageBody>;

export const SharedQueueMessageBody = z.discriminatedUnion("type", [
  SharedQueueExecuteMessageBody,
  SharedQueueResumeMessageBody,
  SharedQueueResumeAfterDurationMessageBody,
  SharedQueueFailMessageBody,
]);

export type SharedQueueMessageBody = z.infer<typeof SharedQueueMessageBody>;

type BackgroundWorkerWithTasks = BackgroundWorker & { tasks: BackgroundWorkerTask[] };

export type SharedQueueConsumerOptions = {
  maximumItemsPerTrace?: number;
  traceTimeoutSeconds?: number;
  nextTickInterval?: number;
  interval?: number;
};

type HandleMessageAction = "ack_and_do_more_work" | "nack" | "nack_and_do_more_work" | "noop";

type DoWorkInternalResult = {
  reason: string;
  outcome: "execution" | "retry_with_nack" | "fail_with_ack" | "noop";
  attrs?: Record<string, string | number | boolean | undefined>;
  error?: Error | string;
  interval?: number;
  action?: HandleMessageAction;
};

type HandleMessageResult = {
  action: HandleMessageAction;
  interval?: number;
  retryInMs?: number;
  reason?: string;
  attrs?: Record<string, string | number | boolean | undefined>;
  error?: Error | string;
};

export class SharedQueueConsumer {
  private _backgroundWorkers: Map<string, BackgroundWorkerWithTasks> = new Map();
  private _deprecatedWorkers: Map<string, BackgroundWorkerWithTasks> = new Map();
  private _enabled = false;
  private _options: Required<SharedQueueConsumerOptions>;
  private _perTraceCountdown: number | undefined;
  private _traceStartedAt: Date | undefined;
  private _currentSpanContext: Context | undefined;
  private _reasonStats: Record<string, number> = {};
  private _actionStats: Record<string, number> = {};
  private _outcomeStats: Record<string, number> = {};
  private _currentSpan: Span | undefined;
  private _endSpanInNextIteration = false;
  private _tasks = sharedQueueTasks;
  private _id: string;
  private _connectedAt: Date;
  private _iterationsCount = 0;
  private _totalIterationsCount = 0;
  private _runningDurationInMs = 0;
  private _currentMessage: MessagePayload | undefined;
  private _currentMessageData: SharedQueueMessageBody | undefined;
  private _stopWorkerQueueConsumer?: () => void;

  constructor(
    private _providerSender: ZodMessageSender<typeof serverWebsocketMessages>,
    options: SharedQueueConsumerOptions = {}
  ) {
    this._options = {
      maximumItemsPerTrace: options.maximumItemsPerTrace ?? 500,
      traceTimeoutSeconds: options.traceTimeoutSeconds ?? 1,
      nextTickInterval: options.nextTickInterval ?? 1000, // 1 second
      interval: options.interval ?? 100, // 100ms
    };

    this._id = generateFriendlyId("shared-queue", 6);
    this._connectedAt = new Date();
  }

  // This method is called when a background worker is deprecated and will no longer be used unless a run is locked to it
  public async deprecateBackgroundWorker(id: string) {
    const backgroundWorker = this._backgroundWorkers.get(id);

    if (!backgroundWorker) {
      return;
    }

    this._deprecatedWorkers.set(id, backgroundWorker);
    this._backgroundWorkers.delete(id);
  }

  public async registerBackgroundWorker(id: string, envId?: string) {
    if (!envId) {
      logger.error("Environment ID is required for background worker registration", {
        backgroundWorkerId: id,
      });
      return;
    }

    const backgroundWorker = await prisma.backgroundWorker.findFirst({
      where: {
        friendlyId: id,
        runtimeEnvironmentId: envId,
      },
      include: {
        tasks: true,
      },
    });

    if (!backgroundWorker) {
      return;
    }

    this._backgroundWorkers.set(backgroundWorker.id, backgroundWorker);

    logger.debug("Registered background worker", { backgroundWorker: backgroundWorker.id });

    // Start reading from the queue if we haven't already
    this.#enable();
  }

  public async start() {
    this.#enable();
  }

  public async stop(reason: string = "Provider disconnected") {
    if (!this._enabled) {
      return;
    }

    console.log("❌ Stopping the SharedQueueConsumer");

    this._stopWorkerQueueConsumer?.();

    logger.debug("Stopping shared queue consumer");
    this._enabled = false;

    if (this._currentSpan) {
      this._currentSpan.end();
    }
  }

  #enable() {
    if (this._enabled) {
      return;
    }

    this._enabled = true;
    this._perTraceCountdown = this._options.maximumItemsPerTrace;
    this._traceStartedAt = new Date();
    this._reasonStats = {};
    this._actionStats = {};
    this._outcomeStats = {};
    this._stopWorkerQueueConsumer = marqs?.startSharedWorkerQueueConsumer(this._id);

    console.log("✅ Started the SharedQueueConsumer");

    this.#doWork().finally(() => {});
  }

  #endCurrentSpan() {
    if (this._currentSpan) {
      for (const [reason, count] of Object.entries(this._reasonStats)) {
        this._currentSpan.setAttribute(`reasons_${reason}`, count);
      }

      for (const [action, count] of Object.entries(this._actionStats)) {
        this._currentSpan.setAttribute(`actions_${action}`, count);
      }

      for (const [outcome, count] of Object.entries(this._outcomeStats)) {
        this._currentSpan.setAttribute(`outcomes_${outcome}`, count);
      }

      this._currentSpan.end();
    }
  }

  async #doWork() {
    if (!this._enabled) {
      this.#endCurrentSpan();
      return;
    }

    // Check if the trace has expired
    if (
      this._perTraceCountdown === 0 ||
      Date.now() - this._traceStartedAt!.getTime() > this._options.traceTimeoutSeconds * 1000 ||
      this._currentSpanContext === undefined ||
      this._endSpanInNextIteration
    ) {
      this.#endCurrentSpan();

      const traceDurationInMs = this._traceStartedAt
        ? Date.now() - this._traceStartedAt.getTime()
        : undefined;
      const iterationsPerSecond = traceDurationInMs
        ? this._iterationsCount / (traceDurationInMs / 1000)
        : undefined;

      // Create a new trace
      this._currentSpan = tracer.startSpan(
        "SharedQueueConsumer.doWork()",
        {
          kind: SpanKind.CONSUMER,
          attributes: {
            id: this._id,
            iterations: this._iterationsCount,
            total_iterations: this._totalIterationsCount,
            options_maximumItemsPerTrace: this._options.maximumItemsPerTrace,
            options_nextTickInterval: this._options.nextTickInterval,
            options_interval: this._options.interval,
            connected_at: this._connectedAt.toISOString(),
            consumer_age_in_seconds: (Date.now() - this._connectedAt.getTime()) / 1000,
            do_work_internal_per_second: this._iterationsCount / (this._runningDurationInMs / 1000),
            running_duration_ms: this._runningDurationInMs,
            trace_timeout_in_seconds: this._options.traceTimeoutSeconds,
            trace_duration_ms: traceDurationInMs,
            iterations_per_second: iterationsPerSecond,
            iterations_per_minute: iterationsPerSecond ? iterationsPerSecond * 60 : undefined,
          },
        },
        ROOT_CONTEXT
      );

      logger.debug("SharedQueueConsumer starting new trace", {
        reasonStats: this._reasonStats,
        actionStats: this._actionStats,
        outcomeStats: this._outcomeStats,
        iterationCount: this._iterationsCount,
        consumerId: this._id,
      });

      // Get the span trace context
      this._currentSpanContext = trace.setSpan(ROOT_CONTEXT, this._currentSpan);

      this._perTraceCountdown = this._options.maximumItemsPerTrace;
      this._traceStartedAt = new Date();
      this._reasonStats = {};
      this._actionStats = {};
      this._outcomeStats = {};
      this._iterationsCount = 0;
      this._runningDurationInMs = 0;
      this._endSpanInNextIteration = false;
    }

    return context.with(this._currentSpanContext ?? ROOT_CONTEXT, async () => {
      await tracer.startActiveSpan("doWorkInternal()", async (span) => {
        let nextInterval = this._options.interval;

        span.setAttributes({
          id: this._id,
          total_iterations: this._totalIterationsCount,
          iterations: this._iterationsCount,
        });

        const startAt = performance.now();

        try {
          const result = await this.#doWorkInternal();

          if (result.reason !== "no_message_dequeued") {
            logger.debug("SharedQueueConsumer doWorkInternal result", { result });
          }

          this._reasonStats[result.reason] = (this._reasonStats[result.reason] ?? 0) + 1;
          this._outcomeStats[result.outcome] = (this._outcomeStats[result.outcome] ?? 0) + 1;

          if (result.action) {
            this._actionStats[result.action] = (this._actionStats[result.action] ?? 0) + 1;
          }

          span.setAttribute("reason", result.reason);

          if (result.attrs) {
            for (const [key, value] of Object.entries(result.attrs)) {
              if (value) {
                span.setAttribute(key, value);
              }
            }
          }

          if (result.error) {
            span.recordException(result.error);
            span.setStatus({ code: SpanStatusCode.ERROR });
            this._currentSpan?.recordException(result.error);
            this._currentSpan?.setStatus({ code: SpanStatusCode.ERROR });
            this._endSpanInNextIteration = true;
          }

          if (typeof result.interval === "number") {
            nextInterval = Math.max(result.interval, 0); // Cannot be negative
          }

          span.setAttribute("nextInterval", nextInterval);
        } catch (error) {
          if (error instanceof Error) {
            this._currentSpan?.recordException(error);
          } else {
            this._currentSpan?.recordException(new Error(String(error)));
          }

          this._endSpanInNextIteration = true;
        } finally {
          this._runningDurationInMs = this._runningDurationInMs + (performance.now() - startAt);
          this._iterationsCount++;
          this._totalIterationsCount++;
          this._perTraceCountdown = this._perTraceCountdown! - 1;

          span.end();

          setTimeout(() => {
            this.#doWork().finally(() => {});
          }, nextInterval);
        }
      });
    });
  }

  async #doWorkInternal(): Promise<DoWorkInternalResult> {
    // Attempt to dequeue a message from the shared queue
    // If no message is available, reschedule the worker to run again in 1 second
    // If a message is available, find the BackgroundWorkerTask that matches the message's taskIdentifier
    // If no matching task is found, nack the message and reschedule the worker to run again in 1 second
    // If the matching task is found, create the task attempt and lock the task run, then send the task run to the client
    // Store the message as a processing message
    // If the websocket connection disconnects before the task run is completed, nack the message
    // When the task run completes, ack the message
    // Using a heartbeat mechanism, if the client keeps responding with a heartbeat, we'll keep the message processing and increase the visibility timeout.

    this._currentMessage = undefined;
    this._currentMessageData = undefined;

    const message = await marqs?.dequeueMessageFromSharedWorkerQueue(this._id);

    if (!message) {
      return {
        reason: "no_message_dequeued",
        outcome: "noop",
        interval: this._options.nextTickInterval,
      };
    }

    const dequeuedAt = new Date();

    logger.log("dequeueMessageInSharedQueue()", { queueMessage: message });

    const messageBody = SharedQueueMessageBody.safeParse(message.data);

    if (!messageBody.success) {
      logger.error("Failed to parse message", {
        queueMessage: message.data,
        error: messageBody.error,
      });

      await this.#ack(message.messageId);

      return {
        reason: "failed_to_parse_message",
        outcome: "fail_with_ack",
        attrs: { message_id: message.messageId, message_version: message.version },
        error: messageBody.error,
      };
    }

    const hydrateAttributes = (attrs: Record<string, string | number | boolean | undefined>) => {
      return {
        ...attrs,
        message_id: message.messageId,
        message_version: message.version,
        run_id: message.messageId,
        message_type: messageBody.data.type,
      };
    };

    this._currentMessage = message;
    this._currentMessageData = messageBody.data;

    const messageResult = await this.#handleMessage(message, messageBody.data, dequeuedAt);

    switch (messageResult.action) {
      case "noop": {
        return {
          reason: messageResult.reason ?? "none_specified",
          outcome: "execution",
          attrs: hydrateAttributes(messageResult.attrs ?? {}),
          error: messageResult.error,
          interval: messageResult.interval,
          action: "noop",
        };
      }
      case "ack_and_do_more_work": {
        await this.#ack(message.messageId);

        return {
          reason: messageResult.reason ?? "none_specified",
          outcome: "fail_with_ack",
          attrs: hydrateAttributes(messageResult.attrs ?? {}),
          error: messageResult.error,
          interval: messageResult.interval,
          action: "ack_and_do_more_work",
        };
      }
      case "nack_and_do_more_work": {
        await this.#nack(message.messageId, messageResult.retryInMs);

        return {
          reason: messageResult.reason ?? "none_specified",
          outcome: "retry_with_nack",
          attrs: hydrateAttributes(messageResult.attrs ?? {}),
          error: messageResult.error,
          interval: messageResult.interval,
          action: "nack_and_do_more_work",
        };
      }
      case "nack": {
        await marqs?.nackMessage(message.messageId);

        return {
          reason: messageResult.reason ?? "none_specified",
          outcome: "retry_with_nack",
          attrs: hydrateAttributes(messageResult.attrs ?? {}),
          error: messageResult.error,
          action: "nack",
        };
      }
    }
  }

  async #handleMessage(
    message: MessagePayload,
    data: SharedQueueMessageBody,
    dequeuedAt: Date
  ): Promise<HandleMessageResult> {
    return await this.#startActiveSpan("handleMessage()", async (span) => {
      // TODO: For every ACK, decide what should be done with the existing run and attempts. Make sure to check the current statuses first.
      switch (data.type) {
        // MARK: EXECUTE
        case "EXECUTE": {
          return await this.#handleExecuteMessage(message, data, dequeuedAt);
        }
        // MARK: DEP RESUME
        // Resume after dependency completed with no remaining retries
        case "RESUME": {
          return await this.#handleResumeMessage(message, data, dequeuedAt);
        }
        // MARK: DURATION RESUME
        // Resume after duration-based wait
        case "RESUME_AFTER_DURATION": {
          return await this.#handleResumeAfterDurationMessage(message, data, dequeuedAt);
        }
        // MARK: FAIL
        // Fail for whatever reason, usually runs that have been resumed but stopped heartbeating
        case "FAIL": {
          return await this.#handleFailMessage(message, data, dequeuedAt);
        }
      }
    });
  }

  async #handleExecuteMessage(
    message: MessagePayload,
    data: SharedQueueExecuteMessageBody,
    dequeuedAt: Date
  ): Promise<HandleMessageResult> {
    const existingTaskRun = await prisma.taskRun.findFirst({
      where: {
        id: message.messageId,
      },
    });

    if (!existingTaskRun) {
      logger.error("No existing task run", {
        queueMessage: message.data,
        messageId: message.messageId,
      });

      // INFO: There used to be a race condition where tasks could be triggered, but execute messages could be dequeued before the run finished being created in the DB
      //       This should not be happening anymore. In case it does, consider reqeueuing here with a brief delay while limiting total retries.
      return { action: "ack_and_do_more_work", reason: "no_existing_task_run" };
    }

    const retryingFromCheckpoint = !!data.checkpointEventId;

    const EXECUTABLE_RUN_STATUSES = {
      fromCheckpoint: ["WAITING_TO_RESUME"] satisfies TaskRunStatus[],
      withoutCheckpoint: ["PENDING", "RETRYING_AFTER_FAILURE"] satisfies TaskRunStatus[],
    };

    if (
      (retryingFromCheckpoint &&
        !EXECUTABLE_RUN_STATUSES.fromCheckpoint.includes(existingTaskRun.status)) ||
      (!retryingFromCheckpoint &&
        !EXECUTABLE_RUN_STATUSES.withoutCheckpoint.includes(existingTaskRun.status))
    ) {
      logger.warn("Task run has invalid status for execution. Going to ack", {
        queueMessage: message.data,
        messageId: message.messageId,
        taskRun: existingTaskRun.id,
        status: existingTaskRun.status,
        retryingFromCheckpoint,
      });

      return {
        action: "ack_and_do_more_work",
        reason: "invalid_run_status",
        attrs: { status: existingTaskRun.status, retryingFromCheckpoint },
      };
    }

    // Check if the task run is locked to a specific worker, if not, use the current worker deployment
    const deployment = await this.#startActiveSpan("findCurrentWorkerDeployment", async (span) => {
      return existingTaskRun.lockedById
        ? await getWorkerDeploymentFromWorkerTask(existingTaskRun.lockedById)
        : existingTaskRun.lockedToVersionId
          ? await getWorkerDeploymentFromWorker(existingTaskRun.lockedToVersionId)
          : await findCurrentWorkerDeployment({
              environmentId: existingTaskRun.runtimeEnvironmentId,
              type: "V1",
            });
    });

    const worker = deployment?.worker;

    if (!deployment || !worker) {
      // This happens when a run is "WAITING_FOR_DEPLOY" and is expected
      logger.info("No matching deployment found for task run", {
        queueMessage: message.data,
        messageId: message.messageId,
      });

      await this.#markRunAsWaitingForDeploy(existingTaskRun.id);

      return {
        action: "ack_and_do_more_work",
        reason: "no_matching_deployment",
        attrs: {
          run_id: existingTaskRun.id,
          locked_by_id: existingTaskRun.lockedById ?? undefined,
          locked_to_version_id: existingTaskRun.lockedToVersionId ?? undefined,
          environment_id: existingTaskRun.runtimeEnvironmentId,
        },
      };
    }

    const imageReference = deployment.imageReference;

    if (!imageReference) {
      logger.error("Deployment is missing an image reference", {
        queueMessage: message.data,
        messageId: message.messageId,
        deployment: deployment.id,
      });

      await this.#markRunAsWaitingForDeploy(existingTaskRun.id);

      return {
        action: "ack_and_do_more_work",
        reason: "missing_image_reference",
        attrs: {
          deployment_id: deployment.id,
        },
      };
    }

    const backgroundTask = worker.tasks.find(
      (task) => task.slug === existingTaskRun.taskIdentifier
    );

    if (!backgroundTask) {
      const nonCurrentTask = await prisma.backgroundWorkerTask.findFirst({
        where: {
          slug: existingTaskRun.taskIdentifier,
          projectId: existingTaskRun.projectId,
          runtimeEnvironmentId: existingTaskRun.runtimeEnvironmentId,
        },
        include: {
          worker: {
            include: {
              deployment: {
                include: {},
              },
            },
          },
        },
      });

      if (nonCurrentTask) {
        logger.warn("Task for this run exists but is not part of the current deploy", {
          taskRun: existingTaskRun.id,
          taskIdentifier: existingTaskRun.taskIdentifier,
        });
      } else {
        logger.warn("Task for this run has never been deployed", {
          taskRun: existingTaskRun.id,
          taskIdentifier: existingTaskRun.taskIdentifier,
        });
      }

      await this.#markRunAsWaitingForDeploy(existingTaskRun.id);

      // If this task is ever deployed, a new message will be enqueued after successful indexing
      return {
        action: "ack_and_do_more_work",
        reason: "task_not_deployed",
        attrs: {
          run_id: existingTaskRun.id,
          task_identifier: existingTaskRun.taskIdentifier,
        },
      };
    }

    const lockedAt = new Date();
    const machinePreset =
      existingTaskRun.machinePreset ??
      machinePresetFromConfig(backgroundTask.machineConfig ?? {}).name;
    const maxDurationInSeconds = getMaxDuration(
      existingTaskRun.maxDurationInSeconds,
      backgroundTask.maxDurationInSeconds
    );
    const startedAt = existingTaskRun.startedAt ?? dequeuedAt;
    const baseCostInCents = env.CENTS_PER_RUN;

    const lockedTaskRun = await prisma.taskRun.update({
      where: {
        id: message.messageId,
      },
      data: {
        lockedAt,
        lockedById: backgroundTask.id,
        lockedToVersionId: worker.id,
        taskVersion: worker.version,
        sdkVersion: worker.sdkVersion,
        cliVersion: worker.cliVersion,
        startedAt: startedAt,
        baseCostInCents: baseCostInCents,
        machinePreset: machinePreset,
        maxDurationInSeconds,
      },
      include: {
        runtimeEnvironment: true,
        attempts: {
          take: 1,
          orderBy: { number: "desc" },
        },
        checkpoints: {
          take: 1,
          orderBy: {
            createdAt: "desc",
          },
        },
        lockedBy: true,
      },
    });

    if (!lockedTaskRun) {
      logger.warn("Failed to lock task run", {
        taskRun: existingTaskRun.id,
        taskIdentifier: existingTaskRun.taskIdentifier,
        deployment: deployment.id,
        backgroundWorker: worker.id,
        messageId: message.messageId,
      });

      return {
        action: "ack_and_do_more_work",
        reason: "failed_to_lock_task_run",
        attrs: {
          run_id: existingTaskRun.id,
          task_identifier: existingTaskRun.taskIdentifier,
          deployment_id: deployment.id,
          background_worker_id: worker.id,
          message_id: message.messageId,
        },
      };
    }

    const queue = await findQueueInEnvironment(
      lockedTaskRun.queue,
      lockedTaskRun.runtimeEnvironmentId,
      lockedTaskRun.lockedById ?? undefined,
      backgroundTask
    );

    if (!queue) {
      logger.debug("SharedQueueConsumer queue not found, so acking message", {
        queueMessage: message,
        taskRunQueue: lockedTaskRun.queue,
        runtimeEnvironmentId: lockedTaskRun.runtimeEnvironmentId,
      });

      return {
        action: "ack_and_do_more_work",
        reason: "queue_not_found",
        attrs: {
          queue_name: sanitizeQueueName(lockedTaskRun.queue),
          runtime_environment_id: lockedTaskRun.runtimeEnvironmentId,
        },
        interval: this._options.nextTickInterval,
      };
    }

    if (!this._enabled) {
      logger.debug("SharedQueueConsumer not enabled, so nacking message", {
        queueMessage: message,
      });

      return {
        action: "nack",
        reason: "not_enabled",
        attrs: {
          message_id: message.messageId,
        },
      };
    }

    const nextAttemptNumber = lockedTaskRun.attempts[0] ? lockedTaskRun.attempts[0].number + 1 : 1;

    const isRetry =
      nextAttemptNumber > 1 &&
      (lockedTaskRun.status === "WAITING_TO_RESUME" ||
        lockedTaskRun.status === "RETRYING_AFTER_FAILURE");

    try {
      if (data.checkpointEventId) {
        const restoreService = new RestoreCheckpointService();

        const checkpoint = await restoreService.call({
          eventId: data.checkpointEventId,
          isRetry,
        });

        if (!checkpoint) {
          logger.error("Failed to restore checkpoint", {
            queueMessage: message.data,
            messageId: message.messageId,
            runStatus: lockedTaskRun.status,
            isRetry,
          });

          return {
            action: "ack_and_do_more_work",
            reason: "failed_to_restore_checkpoint",
            attrs: {
              run_status: lockedTaskRun.status,
              is_retry: isRetry,
              checkpoint_event_id: data.checkpointEventId,
            },
          };
        }

        return {
          action: "noop",
          reason: "restored_checkpoint",
          attrs: {
            checkpoint_event_id: data.checkpointEventId,
          },
        };
      }

      if (!worker.supportsLazyAttempts) {
        try {
          const service = new CreateTaskRunAttemptService();
          await service.call({
            runId: lockedTaskRun.id,
            setToExecuting: false,
          });
        } catch (error) {
          logger.error("Failed to create task run attempt for outdated worker", {
            error,
            taskRun: lockedTaskRun.id,
          });

          const service = new CrashTaskRunService();
          await service.call(lockedTaskRun.id, {
            errorCode: TaskRunErrorCodes.OUTDATED_SDK_VERSION,
          });

          return {
            action: "ack_and_do_more_work",
            reason: "failed_to_create_attempt_for_outdated_worker",
            attrs: {
              message_id: message.messageId,
              run_id: lockedTaskRun.id,
            },
            error: error instanceof Error ? error : String(error),
          };
        }
      }

      if (isRetry && !data.retryCheckpointsDisabled) {
        socketIo.coordinatorNamespace.emit("READY_FOR_RETRY", {
          version: "v1",
          runId: lockedTaskRun.id,
        });

        // Retries for workers with disabled retry checkpoints will be handled just like normal attempts
        return {
          action: "noop",
          reason: "retry_checkpoints_disabled",
        };
      }

      const machine =
        machinePresetFromRun(lockedTaskRun) ??
        machinePresetFromConfig(lockedTaskRun.lockedBy?.machineConfig ?? {});

      return await this.#startActiveSpan("scheduleAttemptOnProvider", async (span) => {
        span.setAttributes({
          run_id: lockedTaskRun.id,
        });

        if (await this._providerSender.validateCanSendMessage()) {
          await this._providerSender.send("BACKGROUND_WORKER_MESSAGE", {
            backgroundWorkerId: worker.friendlyId,
            data: {
              type: "SCHEDULE_ATTEMPT",
              image: imageReference,
              version: deployment.version,
              machine,
              nextAttemptNumber,
              // identifiers
              id: "placeholder", // TODO: Remove this completely in a future release
              envId: lockedTaskRun.runtimeEnvironment.id,
              envType: lockedTaskRun.runtimeEnvironment.type,
              orgId: lockedTaskRun.runtimeEnvironment.organizationId,
              projectId: lockedTaskRun.runtimeEnvironment.projectId,
              runId: lockedTaskRun.id,
              dequeuedAt: dequeuedAt.getTime(),
            },
          });

          return {
            action: "noop",
            reason: "scheduled_attempt",
            attrs: {
              next_attempt_number: nextAttemptNumber,
            },
          };
        } else {
          return {
            action: "nack_and_do_more_work",
            reason: "provider_not_connected",
            attrs: {
              run_id: lockedTaskRun.id,
            },
            interval: this._options.nextTickInterval,
            retryInMs: 5_000,
          };
        }
      });
    } catch (e) {
      // We now need to unlock the task run and delete the task run attempt
      await prisma.$transaction([
        prisma.taskRun.update({
          where: {
            id: lockedTaskRun.id,
          },
          data: {
            lockedAt: null,
            lockedById: null,
            status: lockedTaskRun.status,
            startedAt: existingTaskRun.startedAt,
          },
        }),
      ]);

      logger.error("SharedQueueConsumer errored, so nacking message", {
        queueMessage: message,
        error: e instanceof Error ? { name: e.name, message: e.message, stack: e.stack } : e,
      });

      return {
        action: "nack_and_do_more_work",
        reason: "failed_to_schedule_attempt",
        error: e instanceof Error ? e : String(e),
        interval: this._options.nextTickInterval,
        retryInMs: 5_000,
      };
    }
  }

  async #handleResumeMessage(
    message: MessagePayload,
    data: SharedQueueResumeMessageBody,
    dequeuedAt: Date
  ): Promise<HandleMessageResult> {
    if (data.checkpointEventId) {
      try {
        const restoreService = new RestoreCheckpointService();

        const checkpoint = await restoreService.call({
          eventId: data.checkpointEventId,
        });

        if (!checkpoint) {
          logger.error("Failed to restore checkpoint", {
            queueMessage: message.data,
            messageId: message.messageId,
          });

          return {
            action: "ack_and_do_more_work",
            reason: "failed_to_restore_checkpoint",
            attrs: {
              checkpoint_event_id: data.checkpointEventId,
            },
          };
        }

        return {
          action: "noop",
          reason: "restored_checkpoint",
          attrs: {
            checkpoint_event_id: data.checkpointEventId,
          },
        };
      } catch (e) {
        return {
          action: "nack_and_do_more_work",
          reason: "failed_to_restore_checkpoint",
          error: e instanceof Error ? e : String(e),
        };
      }
    }

    const resumableRun = await prisma.taskRun.findFirst({
      where: {
        id: message.messageId,
        status: {
          notIn: FINAL_RUN_STATUSES,
        },
      },
    });

    if (!resumableRun) {
      logger.error("Resumable run not found", {
        queueMessage: message.data,
        messageId: message.messageId,
      });

      return {
        action: "ack_and_do_more_work",
        reason: "run_not_found",
      };
    }

    if (resumableRun.status !== "EXECUTING") {
      logger.warn("Run is not executing, will try to resume anyway", {
        queueMessage: message.data,
        messageId: message.messageId,
        runStatus: resumableRun.status,
      });
    }

    const resumableAttempt = await prisma.taskRunAttempt.findFirst({
      where: {
        id: data.resumableAttemptId,
      },
      include: {
        checkpoints: {
          take: 1,
          orderBy: {
            createdAt: "desc",
          },
        },
      },
    });

    if (!resumableAttempt) {
      logger.error("Resumable attempt not found", {
        queueMessage: message.data,
        messageId: message.messageId,
      });

      return {
        action: "ack_and_do_more_work",
        reason: "resumable_attempt_not_found",
        attrs: {
          attempt_id: data.resumableAttemptId,
        },
      };
    }

    const queue = await findQueueInEnvironment(
      resumableRun.queue,
      resumableRun.runtimeEnvironmentId,
      resumableRun.lockedById ?? undefined
    );

    if (!queue) {
      logger.debug("SharedQueueConsumer queue not found, so nacking message", {
        queueName: sanitizeQueueName(resumableRun.queue),
        attempt: resumableAttempt,
      });

      return {
        action: "ack_and_do_more_work",
        reason: "queue_not_found",
        attrs: {
          queue_name: sanitizeQueueName(resumableRun.queue),
        },
        interval: this._options.nextTickInterval,
      };
    }

    if (!this._enabled) {
      return {
        action: "nack",
        reason: "not_enabled",
        attrs: {
          message_id: message.messageId,
        },
      };
    }

    try {
      const { completions, executions } = await this.#resolveCompletedAttemptsForResumeMessage(
        data.completedAttemptIds
      );

      const resumeMessage = {
        version: "v1" as const,
        runId: resumableAttempt.taskRunId,
        attemptId: resumableAttempt.id,
        attemptFriendlyId: resumableAttempt.friendlyId,
        completions,
        executions,
      };

      logger.debug("Broadcasting RESUME_AFTER_DEPENDENCY_WITH_ACK", {
        resumeMessage,
        message,
        resumableRun,
      });

      // The attempt should still be running so we can broadcast to all coordinators to resume immediately
      const responses = await this.#startActiveSpan(
        "emitResumeAfterDependencyWithAck",
        async (span) => {
          try {
            span.setAttribute("attempt_id", resumableAttempt.id);
            span.setAttribute(
              "timeout_in_ms",
              env.SHARED_QUEUE_CONSUMER_EMIT_RESUME_DEPENDENCY_TIMEOUT_MS
            );

            const responses = await socketIo.coordinatorNamespace
              .timeout(env.SHARED_QUEUE_CONSUMER_EMIT_RESUME_DEPENDENCY_TIMEOUT_MS)
              .emitWithAck("RESUME_AFTER_DEPENDENCY_WITH_ACK", resumeMessage);

            span.setAttribute("response_count", responses.length);

            const hasSuccess = responses.some((response) => response.success);

            span.setAttribute("has_success", hasSuccess);
            span.setAttribute("is_timeout", false);

            return responses;
          } catch (e) {
            if (e instanceof Error && "responses" in e && Array.isArray(e.responses)) {
              span.setAttribute("is_timeout", false);

              const responses = e.responses as AckCallbackResult[];

              span.setAttribute("response_count", responses.length);

              const hasSuccess = responses.some(
                (response) => "success" in response && response.success
              );

              span.setAttribute("has_success", hasSuccess);

              return responses;
            }

            throw e;
          }
        }
      );

      logger.debug("RESUME_AFTER_DEPENDENCY_WITH_ACK received", {
        resumeMessage,
        responses,
        message,
      });

      if (responses.length === 0) {
        logger.error("RESUME_AFTER_DEPENDENCY_WITH_ACK no response", {
          resumeMessage,
          message,
        });

        return {
          action: "nack_and_do_more_work",
          reason: "resume_after_dependency_with_ack_no_response",
          attrs: {
            resume_message: "RESUME_AFTER_DEPENDENCY_WITH_ACK",
          },
          interval: this._options.nextTickInterval,
          retryInMs: 5_000,
        };
      }

      const hasSuccess = responses.some((response) => response.success);

      if (hasSuccess) {
        return {
          action: "noop",
          reason: "resume_after_dependency_with_ack_success",
        };
      }

      // No coordinator was able to resume the run
      logger.warn("RESUME_AFTER_DEPENDENCY_WITH_ACK failed", {
        resumeMessage,
        responses,
        message,
      });

      // Let's check if the run is frozen
      if (resumableRun.status === "WAITING_TO_RESUME") {
        logger.debug("RESUME_AFTER_DEPENDENCY_WITH_ACK run is waiting to be restored", {
          queueMessage: message.data,
          messageId: message.messageId,
        });

        try {
          const restoreService = new RestoreCheckpointService();

          const checkpointEvent = await restoreService.getLastCheckpointEventIfUnrestored(
            resumableRun.id
          );

          if (checkpointEvent) {
            // The last checkpoint hasn't been restored yet, so restore it
            const checkpoint = await restoreService.call({
              eventId: checkpointEvent.id,
            });

            if (!checkpoint) {
              logger.debug("RESUME_AFTER_DEPENDENCY_WITH_ACK failed to restore checkpoint", {
                queueMessage: message.data,
                messageId: message.messageId,
              });

              return {
                action: "ack_and_do_more_work",
                reason: "failed_to_restore_checkpoint",
                attrs: {
                  checkpoint_event_id: checkpointEvent.id,
                },
              };
            }

            logger.debug("RESUME_AFTER_DEPENDENCY_WITH_ACK restored checkpoint", {
              queueMessage: message.data,
              messageId: message.messageId,
              checkpoint,
            });

            return {
              action: "noop",
              reason: "restored_checkpoint",
              attrs: {
                checkpoint_event_id: data.checkpointEventId,
              },
            };
          } else {
            logger.debug(
              "RESUME_AFTER_DEPENDENCY_WITH_ACK run is frozen without last checkpoint event",
              {
                queueMessage: message.data,
                messageId: message.messageId,
              }
            );

            return {
              action: "noop",
              reason: "resume_after_dependency_with_ack_frozen",
            };
          }
        } catch (e) {
          return {
            action: "nack_and_do_more_work",
            reason: "waiting_to_resume_threw",
            error: e instanceof Error ? e : String(e),
            interval: this._options.nextTickInterval,
            retryInMs: 5_000,
          };
        }
      }

      logger.debug("RESUME_AFTER_DEPENDENCY_WITH_ACK retrying", {
        queueMessage: message.data,
        messageId: message.messageId,
      });

      return {
        action: "nack_and_do_more_work",
        reason: "resume_after_dependency_with_ack_retrying",
        attrs: {
          message_id: message.messageId,
        },
        interval: this._options.nextTickInterval,
        retryInMs: 5_000,
      };
    } catch (e) {
      if (e instanceof ResumePayloadAttemptsNotFoundError) {
        return {
          action: "ack_and_do_more_work",
          reason: "failed_to_get_resume_payloads_for_attempts",
          attrs: {
            attempt_ids: e.attemptIds.join(","),
          },
        };
      } else if (e instanceof ResumePayloadExecutionNotFoundError) {
        return {
          action: "ack_and_do_more_work",
          reason: "failed_to_get_resume_payloads_missing_execution",
          attrs: {
            attempt_id: e.attemptId,
          },
        };
      } else if (e instanceof ResumePayloadCompletionNotFoundError) {
        return {
          action: "ack_and_do_more_work",
          reason: "failed_to_get_resume_payloads_missing_completion",
          attrs: {
            attempt_id: e.attemptId,
          },
        };
      }

      logger.error("RESUME_AFTER_DEPENDENCY_WITH_ACK threw, nacking with delay", {
        message,
        error: e,
      });

      return {
        action: "nack_and_do_more_work",
        reason: "resume_after_dependency_with_ack_threw",
        error: e instanceof Error ? e : String(e),
        interval: this._options.nextTickInterval,
        retryInMs: 5_000,
      };
    }
  }

  async #handleResumeAfterDurationMessage(
    message: MessagePayload,
    data: SharedQueueResumeAfterDurationMessageBody,
    dequeuedAt: Date
  ): Promise<HandleMessageResult> {
    try {
      const restoreService = new RestoreCheckpointService();

      const checkpoint = await restoreService.call({
        eventId: data.checkpointEventId,
      });

      if (!checkpoint) {
        logger.error("Failed to restore checkpoint", {
          queueMessage: message.data,
          messageId: message.messageId,
        });

        return {
          action: "ack_and_do_more_work",
          reason: "failed_to_restore_checkpoint",
          attrs: {
            checkpoint_event_id: data.checkpointEventId,
          },
        };
      }

      return {
        action: "noop",
        reason: "restored_checkpoint",
        attrs: {
          checkpoint_event_id: data.checkpointEventId,
        },
      };
    } catch (e) {
      return {
        action: "nack_and_do_more_work",
        reason: "restoring_checkpoint_threw",
        error: e instanceof Error ? e : String(e),
      };
    }
  }

  async #handleFailMessage(
    message: MessagePayload,
    data: SharedQueueFailMessageBody,
    dequeuedAt: Date
  ): Promise<HandleMessageResult> {
    const existingTaskRun = await prisma.taskRun.findFirst({
      where: {
        id: message.messageId,
      },
    });

    if (!existingTaskRun) {
      logger.error("No existing task run to fail", {
        queueMessage: data,
        messageId: message.messageId,
      });

      return {
        action: "ack_and_do_more_work",
        reason: "no_existing_task_run",
      };
    }

    // TODO: Consider failing the attempt and retrying instead. This may not be a good idea, as dequeued FAIL messages tend to point towards critical, persistent errors.
    const service = new CrashTaskRunService();
    await service.call(existingTaskRun.id, {
      crashAttempts: true,
      reason: data.reason,
    });

    return {
      action: "ack_and_do_more_work",
      reason: "message_failed",
    };
  }

  async #ack(messageId: string) {
    await marqs?.acknowledgeMessage(messageId, "Acking and doing more work in SharedQueueConsumer");
  }

  async #nack(messageId: string, nackRetryInMs?: number) {
    const retryAt = nackRetryInMs ? Date.now() + nackRetryInMs : undefined;
    await marqs?.nackMessage(messageId, retryAt);
  }

  async #markRunAsWaitingForDeploy(runId: string) {
    logger.debug("Marking run as waiting for deploy", { runId });

    const _run = await prisma.taskRun.update({
      where: {
        id: runId,
      },
      data: {
        status: "WAITING_FOR_DEPLOY",
      },
    });
  }

  async #resolveCompletedAttemptsForResumeMessage(
    completedAttemptIds: string[]
  ): Promise<{ completions: TaskRunExecutionResult[]; executions: TaskRunExecution[] }> {
    return await this.#startActiveSpan("resolveCompletedAttemptsForResumeMessage", async (span) => {
      span.setAttribute("completed_attempt_count", completedAttemptIds.length);

      // Chunk the completedAttemptIds into chunks of 10
      const chunkedCompletedAttemptIds = chunk(
        completedAttemptIds,
        env.SHARED_QUEUE_CONSUMER_RESOLVE_PAYLOADS_BATCH_SIZE
      );

      span.setAttribute("chunk_count", chunkedCompletedAttemptIds.length);
      span.setAttribute("chunk_size", env.SHARED_QUEUE_CONSUMER_RESOLVE_PAYLOADS_BATCH_SIZE);

      const allResumePayloads = await this.#startActiveSpan(
        "resolveAllResumePayloads",
        async (span) => {
          span.setAttribute("chunk_count", chunkedCompletedAttemptIds.length);
          span.setAttribute("chunk_size", env.SHARED_QUEUE_CONSUMER_RESOLVE_PAYLOADS_BATCH_SIZE);
          span.setAttribute("completed_attempt_count", completedAttemptIds.length);

          return await Promise.all(
            chunkedCompletedAttemptIds.map(async (attemptIds) => {
              const payloads = await this.#startActiveSpan("getResumePayloads", async (span) => {
                span.setAttribute("attempt_ids", attemptIds.join(","));
                span.setAttribute("attempt_count", attemptIds.length);

                const payloads = await this._tasks.getResumePayloads(attemptIds);

                span.setAttribute("payload_count", payloads.length);

                return payloads;
              });

              return {
                completions: payloads.map((payload) => payload.completion),
                executions: payloads.map((payload) => payload.execution),
              };
            })
          );
        }
      );

      return {
        completions: allResumePayloads.flatMap((payload) => payload.completions),
        executions: allResumePayloads.flatMap((payload) => payload.executions),
      };
    });
  }

  async #startActiveSpan<T>(
    name: string,
    fn: (span: Span) => Promise<T>,
    options?: SpanOptions
  ): Promise<T> {
    return await tracer.startActiveSpan(name, options ?? {}, async (span) => {
      if (this._currentMessage) {
        span.setAttribute("message_id", this._currentMessage.messageId);
        span.setAttribute("run_id", this._currentMessage.messageId);
        span.setAttribute("message_version", this._currentMessage.version);
      }

      if (this._currentMessageData) {
        span.setAttribute("message_type", this._currentMessageData.type);
      }

      try {
        return await fn(span);
      } catch (error) {
        if (error instanceof Error) {
          span.recordException(error);
        } else {
          span.recordException(String(error));
        }

        span.setStatus({ code: SpanStatusCode.ERROR });

        throw error;
      } finally {
        span.end();
      }
    });
  }
}

class ResumePayloadAttemptsNotFoundError extends Error {
  constructor(public readonly attemptIds: string[]) {
    super(`Resume payload attempts not found for attempts ${attemptIds.join(", ")}`);
  }
}

class ResumePayloadExecutionNotFoundError extends Error {
  constructor(public readonly attemptId: string) {
    super(`Resume payload execution not found for attempt ${attemptId}`);
  }
}

class ResumePayloadCompletionNotFoundError extends Error {
  constructor(public readonly attemptId: string) {
    super(`Resume payload completion not found for attempt ${attemptId}`);
  }
}

function chunk<T>(arr: T[], chunkSize: number): T[][] {
  return Array.from({ length: Math.ceil(arr.length / chunkSize) }, (_, i) =>
    arr.slice(i * chunkSize, i * chunkSize + chunkSize)
  );
}

export const AttemptForCompletionGetPayload = {
  select: {
    status: true,
    output: true,
    outputType: true,
    error: true,
    taskRun: {
      select: {
        taskIdentifier: true,
        friendlyId: true,
      },
    },
  },
} as const;

type AttemptForCompletion = Prisma.TaskRunAttemptGetPayload<typeof AttemptForCompletionGetPayload>;

export const AttemptForExecutionGetPayload = {
  select: {
    id: true,
    friendlyId: true,
    taskRunId: true,
    number: true,
    startedAt: true,
    createdAt: true,
    backgroundWorkerId: true,
    backgroundWorkerTaskId: true,
    backgroundWorker: {
      select: {
        contentHash: true,
        version: true,
      },
    },
    backgroundWorkerTask: {
      select: {
        machineConfig: true,
        slug: true,
        filePath: true,
        exportName: true,
      },
    },
    status: true,
    runtimeEnvironment: {
      select: {
        ...RuntimeEnvironmentForEnvRepoPayload.select,
        organization: {
          select: {
            id: true,
            slug: true,
            title: true,
          },
        },
        project: {
          select: {
            id: true,
            externalRef: true,
            slug: true,
            name: true,
          },
        },
      },
    },
    taskRun: {
      select: {
        id: true,
        status: true,
        traceContext: true,
        machinePreset: true,
        friendlyId: true,
        payload: true,
        payloadType: true,
        context: true,
        createdAt: true,
        startedAt: true,
        isTest: true,
        replayedFromTaskRunFriendlyId: true,
        metadata: true,
        metadataType: true,
        idempotencyKey: true,
        usageDurationMs: true,
        costInCents: true,
        baseCostInCents: true,
        maxDurationInSeconds: true,
        runTags: true,
        taskEventStore: true,
      },
    },
    queue: {
      select: {
        name: true,
        friendlyId: true,
      },
    },
  },
} as const;

type AttemptForExecution = Prisma.TaskRunAttemptGetPayload<typeof AttemptForExecutionGetPayload>;

class SharedQueueTasks {
  private _completionPayloadFromAttempt(attempt: AttemptForCompletion): TaskRunExecutionResult {
    const ok = attempt.status === "COMPLETED";

    if (ok) {
      const success: TaskRunSuccessfulExecutionResult = {
        ok,
        id: attempt.taskRun.friendlyId,
        output: attempt.output ?? undefined,
        outputType: attempt.outputType,
        taskIdentifier: attempt.taskRun.taskIdentifier,
      };
      return success;
    } else {
      const failure: TaskRunFailedExecutionResult = {
        ok,
        id: attempt.taskRun.friendlyId,
        error: attempt.error as TaskRunError,
        taskIdentifier: attempt.taskRun.taskIdentifier,
      };
      return failure;
    }
  }

  private async _executionFromAttempt(
    attempt: AttemptForExecution,
    machinePreset?: MachinePreset
  ): Promise<V3ProdTaskRunExecution> {
    const { backgroundWorkerTask, taskRun, queue } = attempt;

    if (!machinePreset) {
      machinePreset =
        machinePresetFromRun(attempt.taskRun) ??
        machinePresetFromConfig(backgroundWorkerTask.machineConfig ?? {});
    }

    const metadata = await parsePacket({
      data: taskRun.metadata ?? undefined,
      dataType: taskRun.metadataType,
    });

    const execution: V3ProdTaskRunExecution = {
      task: {
        id: backgroundWorkerTask.slug,
        filePath: backgroundWorkerTask.filePath,
        exportName: backgroundWorkerTask.exportName ?? backgroundWorkerTask.slug,
      },
      attempt: {
        id: attempt.friendlyId,
        number: attempt.number,
        startedAt: attempt.startedAt ?? attempt.createdAt,
        backgroundWorkerId: attempt.backgroundWorkerId,
        backgroundWorkerTaskId: attempt.backgroundWorkerTaskId,
        status: "EXECUTING" as const,
      },
      run: {
        id: taskRun.friendlyId,
        payload: taskRun.payload,
        payloadType: taskRun.payloadType,
        context: taskRun.context,
        createdAt: taskRun.createdAt,
        startedAt: taskRun.startedAt ?? taskRun.createdAt,
        tags: taskRun.runTags ?? [],
        isTest: taskRun.isTest,
        isReplay: !!taskRun.replayedFromTaskRunFriendlyId,
        idempotencyKey: taskRun.idempotencyKey ?? undefined,
        durationMs: taskRun.usageDurationMs,
        costInCents: taskRun.costInCents,
        baseCostInCents: taskRun.baseCostInCents,
        metadata,
        maxDuration: taskRun.maxDurationInSeconds ?? undefined,
      },
      queue: {
        id: queue.friendlyId,
        name: queue.name,
      },
      environment: {
        id: attempt.runtimeEnvironment.id,
        slug: attempt.runtimeEnvironment.slug,
        type: attempt.runtimeEnvironment.type,
      },
      organization: {
        id: attempt.runtimeEnvironment.organization.id,
        slug: attempt.runtimeEnvironment.organization.slug,
        name: attempt.runtimeEnvironment.organization.title,
      },
      project: {
        id: attempt.runtimeEnvironment.project.id,
        ref: attempt.runtimeEnvironment.project.externalRef,
        slug: attempt.runtimeEnvironment.project.slug,
        name: attempt.runtimeEnvironment.project.name,
      },
      batch: undefined, // TODO: Removing this for now until we can do it more efficiently
      worker: {
        id: attempt.backgroundWorkerId,
        contentHash: attempt.backgroundWorker.contentHash,
        version: attempt.backgroundWorker.version,
      },
      machine: machinePreset,
    };

    return execution;
  }

  async getCompletionPayloadFromAttempt(id: string): Promise<TaskRunExecutionResult | undefined> {
    const attempt = await prisma.taskRunAttempt.findFirst({
      where: {
        id,
        status: {
          in: FINAL_ATTEMPT_STATUSES,
        },
      },
      ...AttemptForCompletionGetPayload,
    });

    if (!attempt) {
      logger.error("No completed attempt found", { id });
      return;
    }

    return this._completionPayloadFromAttempt(attempt);
  }

  async getExecutionPayloadFromAttempt({
    id,
    setToExecuting,
    isRetrying,
    skipStatusChecks,
  }: {
    id: string;
    setToExecuting?: boolean;
    isRetrying?: boolean;
    skipStatusChecks?: boolean;
  }): Promise<V3ProdTaskRunExecutionPayload | undefined> {
    const attempt = await prisma.taskRunAttempt.findFirst({
      where: {
        id,
      },
      ...AttemptForExecutionGetPayload,
    });

    if (!attempt) {
      logger.error("getExecutionPayloadFromAttempt: No attempt found", { id });
      return;
    }

    if (!skipStatusChecks) {
      switch (attempt.status) {
        case "CANCELED":
        case "EXECUTING": {
          logger.error(
            "getExecutionPayloadFromAttempt: Invalid attempt status for execution payload retrieval",
            {
              attemptId: id,
              status: attempt.status,
            }
          );
          return;
        }
      }

      switch (attempt.taskRun.status) {
        case "CANCELED":
        case "EXECUTING":
        case "INTERRUPTED": {
          logger.error(
            "getExecutionPayloadFromAttempt: Invalid run status for execution payload retrieval",
            {
              attemptId: id,
              runId: attempt.taskRunId,
              status: attempt.taskRun.status,
            }
          );
          return;
        }
      }
    }

    if (setToExecuting) {
      if (isFinalAttemptStatus(attempt.status) || isFinalRunStatus(attempt.taskRun.status)) {
        logger.error("getExecutionPayloadFromAttempt: Status already in final state", {
          attempt: {
            id: attempt.id,
            status: attempt.status,
          },
          run: {
            id: attempt.taskRunId,
            status: attempt.taskRun.status,
          },
        });
        return;
      }

      await prisma.taskRunAttempt.update({
        where: {
          id,
        },
        data: {
          status: "EXECUTING",
          taskRun: {
            update: {
              data: {
                status: isRetrying ? "RETRYING_AFTER_FAILURE" : "EXECUTING",
              },
            },
          },
        },
      });
    }

    const { backgroundWorkerTask, taskRun } = attempt;

    const machinePreset =
      machinePresetFromRun(attempt.taskRun) ??
      machinePresetFromConfig(backgroundWorkerTask.machineConfig ?? {});

    const execution = await this._executionFromAttempt(attempt, machinePreset);
    const variables = await this.#buildEnvironmentVariables(
      attempt.runtimeEnvironment,
      taskRun.id,
      machinePreset,
      taskRun.taskEventStore ?? undefined
    );

    const payload: V3ProdTaskRunExecutionPayload = {
      execution,
      traceContext: taskRun.traceContext as Record<string, unknown>,
      environment: variables.reduce((acc: Record<string, string>, curr) => {
        acc[curr.key] = curr.value;
        return acc;
      }, {}),
    };

    return payload;
  }

  async getResumePayload(attemptId: string): Promise<
    | {
        execution: V3ProdTaskRunExecution;
        completion: TaskRunExecutionResult;
      }
    | undefined
  > {
    const attempt = await prisma.taskRunAttempt.findFirst({
      where: {
        id: attemptId,
      },
      select: {
        ...AttemptForExecutionGetPayload.select,
        error: true,
        output: true,
        outputType: true,
        taskRun: {
          select: {
            ...AttemptForExecutionGetPayload.select.taskRun.select,
            taskIdentifier: true,
          },
        },
      },
    });

    if (!attempt) {
      logger.error("getResumePayload: No attempt found", { id: attemptId });
      return;
    }

    const execution = await this._executionFromAttempt(attempt);
    const completion = this._completionPayloadFromAttempt(attempt);

    return {
      execution,
      completion,
    };
  }

  async getResumePayloads(attemptIds: string[]): Promise<
    Array<{
      execution: V3ProdTaskRunExecution;
      completion: TaskRunExecutionResult;
    }>
  > {
    const attempts = await prisma.taskRunAttempt.findMany({
      where: {
        id: {
          in: attemptIds,
        },
      },
      select: {
        ...AttemptForExecutionGetPayload.select,
        error: true,
        output: true,
        outputType: true,
        taskRun: {
          select: {
            ...AttemptForExecutionGetPayload.select.taskRun.select,
            taskIdentifier: true,
          },
        },
      },
    });

    if (attempts.length !== attemptIds.length) {
      logger.error("getResumePayloads: Not all attempts found", { attemptIds });

      throw new ResumePayloadAttemptsNotFoundError(attemptIds);
    }

    const payloads = await Promise.all(
      attempts.map(async (attempt) => {
        const execution = await this._executionFromAttempt(attempt);

        if (!execution) {
          throw new ResumePayloadExecutionNotFoundError(attempt.id);
        }

        const completion = this._completionPayloadFromAttempt(attempt);

        if (!completion) {
          throw new ResumePayloadCompletionNotFoundError(attempt.id);
        }

        return {
          execution,
          completion,
        };
      })
    );

    return payloads;
  }

  async getLatestExecutionPayloadFromRun(
    id: string,
    setToExecuting?: boolean,
    isRetrying?: boolean
  ): Promise<V3ProdTaskRunExecutionPayload | undefined> {
    const run = await prisma.taskRun.findFirst({
      where: {
        id,
      },
      include: {
        attempts: {
          take: 1,
          orderBy: {
            createdAt: "desc",
          },
        },
      },
    });

    const latestAttempt = run?.attempts[0];

    if (!latestAttempt) {
      logger.error("No attempts for run", { id });
      return;
    }

    return this.getExecutionPayloadFromAttempt({
      id: latestAttempt.id,
      setToExecuting,
      isRetrying,
    });
  }

  async getLazyAttemptPayload(
    envId: string,
    runId: string
  ): Promise<TaskRunExecutionLazyAttemptPayload | undefined> {
    const environment = await findEnvironmentById(envId);

    if (!environment) {
      logger.error("getLazyAttemptPayload: Environment not found", { runId, envId });
      return;
    }

    const run = await prisma.taskRun.findFirst({
      where: {
        id: runId,
      },
      select: {
        id: true,
        traceContext: true,
        friendlyId: true,
        isTest: true,
        replayedFromTaskRunFriendlyId: true,
        lockedBy: {
          select: {
            machineConfig: true,
          },
        },
        machinePreset: true,
        taskEventStore: true,
      },
    });

    if (!run) {
      logger.error("getLazyAttemptPayload: Run not found", { runId, envId });
      return;
    }

    const attemptCount = await prisma.taskRunAttempt.count({
      where: {
        taskRunId: run.id,
      },
    });

    logger.debug("getLazyAttemptPayload: Getting lazy attempt payload for run", {
      run,
      attemptCount,
    });

    const machinePreset =
      machinePresetFromRun(run) ?? machinePresetFromConfig(run.lockedBy?.machineConfig ?? {});

    const variables = await this.#buildEnvironmentVariables(
      environment,
      run.id,
      machinePreset,
      run.taskEventStore ?? undefined
    );

    return {
      traceContext: run.traceContext as Record<string, unknown>,
      environment: variables.reduce((acc: Record<string, string>, curr) => {
        acc[curr.key] = curr.value;
        return acc;
      }, {}),
      runId: run.friendlyId,
      messageId: run.id,
      isTest: run.isTest,
      isReplay: !!run.replayedFromTaskRunFriendlyId,
      attemptCount,
      metrics: [],
    } satisfies TaskRunExecutionLazyAttemptPayload;
  }

  async taskHeartbeat(attemptFriendlyId: string) {
    logger.debug("[SharedQueueConsumer] taskHeartbeat()", { id: attemptFriendlyId });

    const taskRunAttempt = await prisma.taskRunAttempt.findFirst({
      where: { friendlyId: attemptFriendlyId },
    });

    if (!taskRunAttempt) {
      return;
    }

    await this.#heartbeat(taskRunAttempt.taskRunId);
  }

  async taskRunHeartbeat(runId: string) {
    logger.debug("[SharedQueueConsumer] taskRunHeartbeat()", { runId });

    await this.#heartbeat(runId);
  }

  public async taskRunFailed(completion: TaskRunFailedExecutionResult) {
    logger.debug("[SharedQueueConsumer] taskRunFailed()", { completion });

    const service = new FailedTaskRunService();

    await service.call(completion.id, completion);
  }

  async #heartbeat(runId: string) {
    await marqs?.heartbeatMessage(runId);

    try {
      // There can be a lot of calls per minute and the data doesn't have to be accurate, so use the read replica
      const taskRun = await $replica.taskRun.findFirst({
        where: {
          id: runId,
        },
        select: {
          id: true,
          status: true,
          runtimeEnvironment: {
            select: {
              type: true,
            },
          },
          lockedToVersion: {
            select: {
              supportsLazyAttempts: true,
            },
          },
        },
      });

      if (!taskRun) {
        logger.error("SharedQueueTasks.#heartbeat: Task run not found", {
          runId,
        });

        return;
      }

      if (taskRun.runtimeEnvironment.type === "DEVELOPMENT") {
        return;
      }

      if (isFinalRunStatus(taskRun.status)) {
        logger.debug("SharedQueueTasks.#heartbeat: Task run is in final status", {
          runId,
          status: taskRun.status,
        });

        // Signal to exit any leftover containers
        socketIo.coordinatorNamespace.emit("REQUEST_RUN_CANCELLATION", {
          version: "v1",
          runId: taskRun.id,
          // Give the run a few seconds to exit to complete any flushing etc
          delayInMs: taskRun.lockedToVersion?.supportsLazyAttempts ? 5_000 : undefined,
        });
        return;
      }
    } catch (error) {
      logger.error("SharedQueueTasks.#heartbeat: Error signaling run cancellation", {
        runId,
        error: error instanceof Error ? error.message : error,
      });
    }
  }

  async #buildEnvironmentVariables(
    environment: RuntimeEnvironmentForEnvRepo,
    runId: string,
    machinePreset: MachinePreset,
    taskEventStore?: string
  ): Promise<Array<EnvironmentVariable>> {
    const variables = await resolveVariablesForEnvironment(environment);

    const jwt = await generateJWTTokenForEnvironment(environment, {
      run_id: runId,
      machine_preset: machinePreset.name,
    });

    if (taskEventStore) {
      const resourceAttributes = JSON.stringify({
        [SemanticInternalAttributes.TASK_EVENT_STORE]: taskEventStore,
      });

      variables.push(...[{ key: "OTEL_RESOURCE_ATTRIBUTES", value: resourceAttributes }]);
    }

    return [
      ...variables,
      ...[
        { key: "TRIGGER_JWT", value: jwt },
        { key: "TRIGGER_RUN_ID", value: runId },
        {
          key: "TRIGGER_MACHINE_PRESET",
          value: machinePreset.name,
        },
      ],
    ];
  }
}

export const sharedQueueTasks = singleton("sharedQueueTasks", () => new SharedQueueTasks());