pikasTech-agentrun/src/mgr/queue-retry.ts

import { AgentRunError } from "../common/errors.js";
import { redactText } from "../common/redaction.js";
import type { CommandRecord, FailureKind, JsonRecord, JsonValue, QueueAttemptRecord, QueueRetryResult, RunRecord } from "../common/types.js";
import { createKubernetesRunnerJob, type RunnerJobDefaults } from "./kubernetes-runner-job.js";
import { buildQueueRunInput, buildQueueRunnerJobInput, queueStateFromCore } from "./queue-dispatch.js";
import type { AgentRunStore } from "./store.js";

interface QueueRetryRequest {
  idempotencyKey: string;
  reason: string;
  dryRun: boolean;
}

export interface RetryQueueTaskOptions {
  store: AgentRunStore;
  taskId: string;
  input: JsonRecord;
  defaults: RunnerJobDefaults;
}

export async function retryQueueTask(options: RetryQueueTaskOptions): Promise<QueueRetryResult> {
  const request = validateQueueRetryRequest(options.input);
  const reservation = await options.store.reserveQueueRetryAttempt(options.taskId, request);
  if (request.dryRun) return queueRetryResult({ reservation, mutation: false });
  const attempt = reservation.attempt;
  if (!attempt) throw new AgentRunError("infra-failed", "queue retry reservation did not return an attempt", { httpStatus: 500 });
  if (attempt.state !== "reserved") return queueRetryReplayResult(options.store, reservation);
  let run: RunRecord;
  let command: CommandRecord;
  let runnerJob: JsonRecord;
  try {
    if (!attempt.runId || !attempt.runnerJobId) throw new AgentRunError("infra-failed", `queue retry attempt ${attempt.attemptId} is missing reserved execution identities`, { httpStatus: 500 });
    run = await options.store.createRun(buildQueueRunInput(reservation.task, {}), { id: attempt.runId });
    command = await options.store.createCommand(run.id, {
      type: "turn",
      payload: reservation.task.payload,
      idempotencyKey: `queue:${reservation.task.id}:attempt:${attempt.attemptId}:command`,
    });
    runnerJob = await createKubernetesRunnerJob({
      store: options.store,
      runId: run.id,
      input: buildQueueRunnerJobInput(reservation.task, command.id, {
        attemptId: attempt.attemptId,
        runnerJobId: attempt.runnerJobId,
        idempotencyKey: `queue:${reservation.task.id}:attempt:${attempt.attemptId}:runner-job`,
      }),
      defaults: options.defaults,
    });
  } catch (error) {
    await options.store.recordQueueRetryAttemptFailure(reservation.task.id, attempt.attemptId, materializationFailure(error));
    throw error;
  }
  const state = queueStateFromCore(run, command);
  const sessionId = run.sessionRef?.sessionId ?? attempt.sessionId;
  const sessionPath = sessionId ? `/api/v1/sessions/${encodeURIComponent(sessionId)}` : attempt.sessionPath;
  const activation = await options.store.activateQueueRetryAttempt(reservation.task.id, attempt.attemptId, {
    commandId: command.id,
    sessionId,
    sessionPath,
    state: state === "pending" ? "running" : state,
    failureKind: run.failureKind,
    failureMessage: run.failureMessage,
  });
  if (activation.activated) {
    await options.store.appendEvent(run.id, "backend_status", {
      phase: "queue-retried",
      taskId: reservation.task.id,
      queue: reservation.task.queue,
      lane: reservation.task.lane,
      attemptId: attempt.attemptId,
      retryIndex: attempt.retryIndex,
      previousAttemptId: attempt.previousAttemptId,
      runnerJobId: attempt.runnerJobId,
      sessionPath,
      reason: attempt.reason,
    });
  }
  return queueRetryResult({
    reservation,
    mutation: activation.activated || runnerJob.mutation === true,
    task: activation.task,
    attempt: activation.attempt,
    run,
    command,
    runnerJob,
  });
}

export function validateQueueRetryRequest(input: JsonRecord): QueueRetryRequest {
  const allowed = new Set(["idempotencyKey", "reason", "dryRun"]);
  const unexpected = Object.keys(input).filter((key) => !allowed.has(key));
  if (unexpected.length > 0) throw new AgentRunError("schema-invalid", `queue retry does not accept task or dispatch patches: ${unexpected.join(", ")}`, { httpStatus: 400, details: { unexpected, allowed: [...allowed], valuesPrinted: false } });
  const idempotencyKey = requiredBoundedString(input.idempotencyKey, "idempotencyKey", 200);
  const reason = requiredBoundedString(input.reason, "reason", 1_000);
  if (input.dryRun !== undefined && typeof input.dryRun !== "boolean") throw new AgentRunError("schema-invalid", "dryRun must be a boolean", { httpStatus: 400 });
  return { idempotencyKey, reason, dryRun: input.dryRun === true };
}

async function queueRetryReplayResult(store: AgentRunStore, reservation: Awaited<ReturnType<AgentRunStore["reserveQueueRetryAttempt"]>>): Promise<QueueRetryResult> {
  const attempt = reservation.attempt;
  if (!attempt) return queueRetryResult({ reservation, mutation: false });
  const run = attempt.runId ? await store.getRun(attempt.runId) : null;
  const command = attempt.commandId ? await store.getCommand(attempt.commandId) : null;
  const runnerJob = run && command
    ? (await store.listRunnerJobs(run.id, command.id)).find((item) => item.id === attempt.runnerJobId || item.attemptId === attempt.attemptId)?.result ?? null
    : null;
  return queueRetryResult({ reservation, mutation: false, run, command, runnerJob });
}

function queueRetryResult(input: {
  reservation: Awaited<ReturnType<AgentRunStore["reserveQueueRetryAttempt"]>>;
  mutation: boolean;
  task?: QueueRetryResult["task"];
  attempt?: QueueAttemptRecord | null;
  run?: RunRecord | null;
  command?: CommandRecord | null;
  runnerJob?: JsonRecord | null;
}): QueueRetryResult {
  const task = input.task ?? input.reservation.task;
  const attempt = input.attempt === undefined ? input.reservation.attempt : input.attempt;
  const run = input.run ?? null;
  const command = input.command ?? null;
  const runnerJob = input.runnerJob ?? null;
  return {
    action: "queue-retry",
    mutation: input.mutation,
    idempotentReplay: input.reservation.idempotentReplay,
    task,
    attempt,
    previousAttempt: input.reservation.previousAttempt,
    run,
    command,
    runnerJob,
    envImage: jsonRecordOrNull(runnerJob?.envImage),
    workReady: jsonRecordOrNull(runnerJob?.workReady),
    allowedTransition: input.reservation.allowedTransition,
    pollActions: attempt ? retryPollActions(task.id, attempt, run, command) : [],
  };
}

function retryPollActions(taskId: string, attempt: QueueAttemptRecord, run: RunRecord | null, command: CommandRecord | null): JsonRecord[] {
  return [
    actionDescriptor("inspect-task", "describe", "task", taskId),
    actionDescriptor("list-attempts", "get", "attempt", attempt.attemptId),
    ...(run ? [actionDescriptor("inspect-run", "describe", "run", run.id, run.id)] : []),
    ...(run && command ? [
      actionDescriptor("inspect-command", "describe", "command", command.id, run.id, command.id),
      { ...actionDescriptor("poll-events", "events", "run", run.id, run.id, command.id), afterSeq: 0, limit: 100 },
    ] : []),
  ];
}

function actionDescriptor(action: string, operation: string, resourceKind: string, resourceName: string, runId: string | null = null, commandId: string | null = null): JsonRecord {
  return { action, operation, resourceKind, resourceName, runId, commandId, valuesPrinted: false };
}

function requiredBoundedString(value: JsonValue | undefined, field: string, maxLength: number): string {
  if (typeof value !== "string" || value.trim().length === 0) throw new AgentRunError("schema-invalid", `${field} is required`, { httpStatus: 400 });
  const normalized = value.trim();
  if (normalized.length > maxLength) throw new AgentRunError("schema-invalid", `${field} must be at most ${maxLength} characters`, { httpStatus: 400 });
  return normalized;
}

function materializationFailure(error: unknown): { failureKind: FailureKind; failureMessage: string } {
  const failureKind = error instanceof AgentRunError ? error.failureKind : "infra-failed";
  const rawMessage = error instanceof Error ? error.message : String(error);
  return { failureKind, failureMessage: redactText(rawMessage).trim().slice(0, 500) || "queue retry materialization failed" };
}

function jsonRecordOrNull(value: JsonValue | undefined): JsonRecord | null {
  return value && typeof value === "object" && !Array.isArray(value) ? value as JsonRecord : null;
}