zopu-code/packages/backend/convex/workExecution.ts

import {
  FakeHarnessLive,
  type FakeScenario,
} from "@code/primitives/harness-runtime";
import {
  type AttemptClassification,
  defaultCodingKitV0,
  resolveOutcome,
} from "@code/primitives/resolver";
import { makeFunctionReference } from "convex/server";
import { ConvexError, v } from "convex/values";
import { Effect } from "effect";

import type { Doc, Id } from "./_generated/dataModel";
import {
  internalAction,
  internalMutation,
  mutation,
  query,
} from "./_generated/server";
import { requireProjectMember } from "./authz";

const executeRef = makeFunctionReference<
  "action",
  { attemptId: Id<"workAttempts">; scenario: FakeScenario }
>("workExecution:executeFakeAttempt");
const claimRef = makeFunctionReference<
  "mutation",
  { attemptId: Id<"workAttempts">; leaseMs: number; owner: string },
  { number: number } | null
>("workExecution:claimAttempt");
const checkpointRef = makeFunctionReference<
  "mutation",
  {
    attemptId: Id<"workAttempts">;
    kind: string;
    message: string;
    metadataJson: string;
    owner: string;
    sequence: number;
  },
  unknown
>("workExecution:checkpointAttempt");
const finishRef = makeFunctionReference<
  "mutation",
  {
    attemptId: Id<"workAttempts">;
    classification: string;
    owner: string;
    retryable: boolean;
    summary: string;
  },
  unknown
>("workExecution:finishAttempt");
const now = () => Date.now();

const requireWorkForMember = async (ctx: any, workId: Id<"works">) => {
  const work = await ctx.db.get(workId);
  if (!work) throw new ConvexError("Work not found");
  await requireProjectMember(ctx, work.projectId);
  return work;
};

const appendWorkEvent = async (
  ctx: any,
  workId: Id<"works">,
  kind: any,
  idempotencyKey: string,
  referenceId?: string,
  payloadJson?: string
) => {
  const existing = await ctx.db
    .query("workEvents")
    .withIndex("by_work_and_idempotencyKey", (q: any) =>
      q.eq("workId", workId).eq("idempotencyKey", idempotencyKey)
    )
    .unique();
  if (existing) return;
  await ctx.db.insert("workEvents", {
    workId,
    kind,
    idempotencyKey,
    createdAt: now(),
    ...(referenceId ? { referenceId } : {}),
    ...(payloadJson ? { payloadJson } : {}),
  });
};

// Slice the Run targets for the currently approved Design. `sliceId` is
// optional only to let the resolver default to the first ready slice.
const resolveRunSlice = async (
  ctx: any,
  work: Doc<"works">,
  sliceId?: string
) => {
  const currentDesignVersion = work.designVersion;
  if (currentDesignVersion === undefined)
    throw new ConvexError("Work has no approved Design to execute");
  const slices = await ctx.db
    .query("workSlices")
    .withIndex("by_work_and_designVersion", (q: any) =>
      q.eq("workId", work._id).eq("designVersion", currentDesignVersion)
    )
    .collect();
  if (slices.length === 0)
    throw new ConvexError("No slices found for the current approved Design");
  const slice = sliceId
    ? slices.find((row: any) => row.sliceId === sliceId)
    : slices.sort((a: any, b: any) => a.ordinal - b.ordinal)[0];
  if (!slice)
    throw new ConvexError(
      "Requested slice does not belong to the current approved Design"
    );
  return slice as Doc<"workSlices">;
};

const setSliceStatus = async (
  ctx: any,
  workId: Id<"works">,
  sliceId: string | undefined,
  status: any
) => {
  if (!sliceId) return;
  const slice = (
    await ctx.db
      .query("workSlices")
      .withIndex("by_work_and_designVersion", (q: any) =>
        q.eq("workId", workId)
      )
      .collect()
  ).find((row: any) => row.sliceId === sliceId);
  if (slice) await ctx.db.patch(slice._id, { status });
};

const RETRYABLE_WORK_STATUSES = ["failed", "needs-input", "blocked", "ready"];

export const startSimulatedExecution = mutation({
  args: {
    workId: v.id("works"),
    scenario: v.union(
      v.literal("success"),
      v.literal("transient-failure-then-success"),
      v.literal("needs-input"),
      v.literal("permanent-failure"),
      v.literal("cancelled")
    ),
    sliceId: v.optional(v.string()),
  },
  handler: async (ctx, args) => {
    const work = await requireWorkForMember(ctx, args.workId);
    if (work.status !== "ready")
      throw new ConvexError("Work must be Ready before simulated execution");
    if (
      work.definitionApprovalVersion !== work.definitionVersion ||
      work.designApprovalVersion !== work.designVersion
    )
      throw new ConvexError(
        "Execution requires exact approved Definition and Design versions"
      );
    const slice = await resolveRunSlice(ctx, work, args.sliceId);
    const createdAt = now();
    const runId = await ctx.db.insert("workRuns", {
      workId: work._id,
      sliceId: slice.sliceId,
      status: "ready",
      scenario: args.scenario,
      kitId: defaultCodingKitV0.id,
      kitVersion: defaultCodingKitV0.version,
      createdAt,
    });
    const attemptId = await ctx.db.insert("workAttempts", {
      runId,
      workId: work._id,
      number: 1,
      status: "queued",
    });
    await ctx.db.patch(work._id, {
      status: "executing",
      updatedAt: createdAt,
    });
    await ctx.db.patch(slice._id, { status: "running" });
    await ctx.db.patch(runId, { status: "running", startedAt: createdAt });
    await appendWorkEvent(
      ctx,
      work._id,
      "run.started",
      `run-started:${runId}`,
      String(runId)
    );
    await ctx.scheduler.runAfter(0, executeRef, {
      attemptId,
      scenario: args.scenario,
    });
    return { runId, attemptId };
  },
});

export const cancelSimulatedExecution = mutation({
  args: { runId: v.id("workRuns") },
  handler: async (ctx, args) => {
    const run = await ctx.db.get(args.runId);
    if (!run) throw new ConvexError("Run not found");
    const work = await requireWorkForMember(ctx, run.workId);
    if (run.status === "terminal" || run.status === "cancelled")
      return { cancelled: false };
    await ctx.db.patch(run._id, {
      status: "cancelled",
      endedAt: now(),
      terminalClassification: "Cancelled",
      terminalSummary: "Simulation cancelled",
    });
    const attempts = await ctx.db
      .query("workAttempts")
      .withIndex("by_run_and_number", (q: any) => q.eq("runId", run._id))
      .collect();
    for (const attempt of attempts)
      if (attempt.status !== "terminal")
        await ctx.db.patch(attempt._id, {
          status: "terminal",
          endedAt: now(),
          classification: "Cancelled",
          summary: "Simulation cancelled",
        });
    await setSliceStatus(ctx, work._id, run.sliceId, "ready");
    await ctx.db.patch(work._id, { status: "ready", updatedAt: now() });
    await appendWorkEvent(
      ctx,
      work._id,
      "run.cancelled",
      `run-cancelled:${run._id}`,
      String(run._id)
    );
    return { cancelled: true };
  },
});

// User-initiated restart of a terminal Run. Unlike the automatic resolver
// retry, this is an explicit decision and may exceed the kit budget.
export const retrySimulatedExecution = mutation({
  args: { runId: v.id("workRuns") },
  handler: async (ctx, args) => {
    const run = await ctx.db.get(args.runId);
    if (!run) throw new ConvexError("Run not found");
    const work = await requireWorkForMember(ctx, run.workId);
    if (run.status !== "terminal")
      throw new ConvexError("Only terminal Runs can be retried");
    if (!RETRYABLE_WORK_STATUSES.includes(work.status))
      throw new ConvexError("Work is not in a retryable state");
    const attempts = await ctx.db
      .query("workAttempts")
      .withIndex("by_run_and_number", (q: any) => q.eq("runId", run._id))
      .collect();
    const number = attempts.length + 1;
    const attemptId = await ctx.db.insert("workAttempts", {
      runId: run._id,
      workId: work._id,
      number,
      status: "queued",
    });
    await ctx.db.patch(run._id, {
      status: "running",
      startedAt: now(),
      endedAt: undefined,
      terminalClassification: undefined,
      terminalSummary: undefined,
    });
    await ctx.db.patch(work._id, { status: "executing", updatedAt: now() });
    await setSliceStatus(ctx, work._id, run.sliceId, "running");
    await ctx.scheduler.runAfter(0, executeRef, {
      attemptId,
      scenario: run.scenario,
    });
    return { attemptId, number };
  },
});

export const claimAttempt = internalMutation({
  args: {
    attemptId: v.id("workAttempts"),
    owner: v.string(),
    leaseMs: v.number(),
  },
  handler: async (ctx, args) => {
    const attempt = await ctx.db.get(args.attemptId);
    if (!attempt || attempt.status === "terminal") return null;
    const claimedAt = now();
    await ctx.db.patch(attempt._id, {
      status: "claimed",
      leaseOwner: args.owner,
      leaseExpiresAt: claimedAt + args.leaseMs,
      startedAt: attempt.startedAt ?? claimedAt,
    });
    await ctx.db.patch(attempt.runId, {
      status: "running",
      startedAt: claimedAt,
    });
    return { ...attempt, status: "claimed" as const, leaseOwner: args.owner };
  },
});

export const checkpointAttempt = internalMutation({
  args: {
    attemptId: v.id("workAttempts"),
    owner: v.string(),
    sequence: v.number(),
    kind: v.string(),
    message: v.string(),
    metadataJson: v.string(),
  },
  handler: async (ctx, args) => {
    const attempt = await ctx.db.get(args.attemptId);
    if (
      !attempt ||
      attempt.leaseOwner !== args.owner ||
      attempt.status === "terminal"
    )
      return false;
    const existing = await ctx.db
      .query("workAttemptEvents")
      .withIndex("by_attempt_and_sequence", (q: any) =>
        q.eq("attemptId", attempt._id).eq("sequence", args.sequence)
      )
      .unique();
    if (!existing)
      await ctx.db.insert("workAttemptEvents", {
        attemptId: attempt._id,
        sequence: args.sequence,
        kind: args.kind,
        message: args.message,
        metadataJson: args.metadataJson,
        occurredAt: now(),
      });
    await ctx.db.patch(attempt._id, {
      status: "running",
      leaseExpiresAt: now() + 60_000,
    });
    return true;
  },
});

// Single resolver mutation: record the attempt outcome, then either schedule
// the next attempt within the kit retry policy or settle the Run/Work at the
// terminal classification. Replaces the old completeAttempt that always reset
// Work to "ready" and ignored the retry policy.
export const finishAttempt = internalMutation({
  args: {
    attemptId: v.id("workAttempts"),
    owner: v.string(),
    classification: v.string(),
    retryable: v.boolean(),
    summary: v.string(),
  },
  handler: async (ctx, args) => {
    const attempt = await ctx.db.get(args.attemptId);
    if (
      !attempt ||
      attempt.leaseOwner !== args.owner ||
      attempt.status === "terminal"
    )
      return null;
    const endedAt = now();
    await ctx.db.patch(attempt._id, {
      status: "terminal",
      endedAt,
      classification: args.classification,
      summary: args.summary,
      leaseExpiresAt: undefined,
    });
    const outcome = {
      classification: args.classification as AttemptClassification,
      retryable: args.retryable,
      summary: args.summary,
    };
    const resolution = resolveOutcome(
      outcome,
      attempt.number,
      defaultCodingKitV0.retryPolicy
    );
    const run = await ctx.db.get(attempt.runId);
    if (!run) return null;
    await appendWorkEvent(
      ctx,
      attempt.workId,
      "attempt.completed",
      `attempt-completed:${attempt._id}`,
      String(attempt._id),
      JSON.stringify({
        classification: args.classification,
        retried: resolution.kind === "retry",
        summary: args.summary,
      })
    );
    if (resolution.kind === "retry") {
      // Keep Run running and Work executing; spin the next attempt.
      const nextAttemptId = await ctx.db.insert("workAttempts", {
        runId: run._id,
        workId: attempt.workId,
        number: attempt.number + 1,
        status: "queued",
      });
      await ctx.scheduler.runAfter(0, executeRef, {
        attemptId: nextAttemptId,
        scenario: run.scenario,
      });
      return { nextAttemptId, retried: true };
    }
    await ctx.db.patch(run._id, {
      status: "terminal",
      endedAt,
      terminalClassification: args.classification,
      terminalSummary: args.summary,
    });
    const work = await ctx.db.get(run.workId);
    if (work)
      await ctx.db.patch(work._id, {
        status: resolution.workStatus,
        updatedAt: endedAt,
      });
    await setSliceStatus(
      ctx,
      attempt.workId,
      run.sliceId,
      args.classification === "Succeeded" ? "completed" : "ready"
    );
    return { retried: false };
  },
});

export const executeFakeAttempt = internalAction({
  args: {
    attemptId: v.id("workAttempts"),
    scenario: v.union(
      v.literal("success"),
      v.literal("transient-failure-then-success"),
      v.literal("needs-input"),
      v.literal("permanent-failure"),
      v.literal("cancelled")
    ),
  },
  handler: async (ctx, args) => {
    const owner = `fake-worker:${args.attemptId}`;
    const claimed = await ctx.runMutation(claimRef, {
      attemptId: args.attemptId,
      owner,
      leaseMs: 60_000,
    });
    if (!claimed) return null;
    const [events, outcome] = await Effect.runPromise(
      FakeHarnessLive().run({
        attemptNumber: claimed.number,
        scenario: args.scenario,
      })
    );
    for (const item of events)
      await ctx.runMutation(checkpointRef, {
        attemptId: args.attemptId,
        kind: item.kind,
        message: item.message,
        metadataJson: JSON.stringify(item.metadata),
        owner,
        sequence: item.sequence,
      });
    await ctx.runMutation(finishRef, {
      attemptId: args.attemptId,
      classification: outcome.classification,
      owner,
      retryable: outcome.retryable,
      summary: outcome.summary,
    });
    return null;
  },
});

// Bounded recovery: a claimed/running attempt whose lease expired means the
// worker died mid-flight. Terminate that attempt, then either resume within
// the kit budget or settle the Run/Work so nothing stays "running" forever.
export const reconcileExpiredAttempts = internalMutation({
  args: {},
  handler: async (ctx) => {
    const active = await ctx.db.query("workAttempts").collect();
    let reconciled = 0;
    for (const attempt of active) {
      if (
        (attempt.status === "claimed" || attempt.status === "running") &&
        attempt.leaseExpiresAt !== undefined &&
        attempt.leaseExpiresAt < now()
      ) {
        await ctx.db.patch(attempt._id, {
          status: "terminal",
          endedAt: now(),
          classification: "Blocked",
          summary: "Attempt lease expired and was reconciled",
          leaseOwner: undefined,
          leaseExpiresAt: undefined,
        });
        const run = await ctx.db.get(attempt.runId);
        await appendWorkEvent(
          ctx,
          attempt.workId,
          "attempt.reconciled",
          `attempt-reconciled:${attempt._id}`,
          String(attempt._id),
          JSON.stringify({ attemptNumber: attempt.number })
        );
        reconciled += 1;
        if (!run) continue;
        const attempts = await ctx.db
          .query("workAttempts")
          .withIndex("by_run_and_number", (q: any) => q.eq("runId", run._id))
          .collect();
        const withinBudget =
          attempts.length < defaultCodingKitV0.retryPolicy.maxAttempts;
        if (run.status === "running" && withinBudget) {
          const nextAttemptId = await ctx.db.insert("workAttempts", {
            runId: run._id,
            workId: attempt.workId,
            number: attempts.length + 1,
            status: "queued",
          });
          await ctx.scheduler.runAfter(0, executeRef, {
            attemptId: nextAttemptId,
            scenario: run.scenario,
          });
        } else {
          await ctx.db.patch(run._id, {
            status: "terminal",
            endedAt: now(),
            terminalClassification: "Blocked",
            terminalSummary: "Run reconciled after expired lease",
          });
          const work = await ctx.db.get(run.workId);
          if (work)
            await ctx.db.patch(work._id, {
              status: "blocked",
              updatedAt: now(),
            });
          await setSliceStatus(ctx, attempt.workId, run.sliceId, "ready");
        }
      }
    }
    return { reconciled };
  },
});

export const listRunEvents = query({
  args: { attemptId: v.id("workAttempts") },
  handler: async (ctx, args) => {
    const attempt = await ctx.db.get(args.attemptId);
    if (!attempt) return [];
    const work = await ctx.db.get(attempt.workId);
    if (!work) return [];
    await requireProjectMember(ctx, work.projectId);
    return await ctx.db
      .query("workAttemptEvents")
      .withIndex("by_attempt_and_sequence", (q: any) =>
        q.eq("attemptId", args.attemptId)
      )
      .order("asc")
      .collect();
  },
});