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

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

import type { 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 completeRef = makeFunctionReference<
  "mutation",
  {
    attemptId: Id<"workAttempts">;
    classification: string;
    owner: string;
    summary: string;
  },
  unknown
>("workExecution:completeAttempt");
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 } : {}),
  });
};

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 createdAt = now();
    const runId = await ctx.db.insert("workRuns", {
      workId: work._id,
      sliceId: args.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(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 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 };
  },
});

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");
    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 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;
  },
});

export const completeAttempt = internalMutation({
  args: {
    attemptId: v.id("workAttempts"),
    owner: v.string(),
    classification: v.string(),
    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 false;
    const endedAt = now();
    await ctx.db.patch(attempt._id, {
      status: "terminal",
      endedAt,
      classification: args.classification,
      summary: args.summary,
      leaseExpiresAt: undefined,
    });
    const run = await ctx.db.get(attempt.runId);
    if (!run) return false;
    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: args.classification === "NeedsInput" ? "needs-input" : "ready",
        updatedAt: endedAt,
      });
    if (work)
      await ctx.db.insert("workEvents", {
        workId: work._id,
        kind: "attempt.completed",
        idempotencyKey: `attempt-completed:${attempt._id}`,
        referenceId: String(attempt._id),
        payloadJson: JSON.stringify({
          classification: args.classification,
          summary: args.summary,
        }),
        createdAt: endedAt,
      });
    return true;
  },
});

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,
        owner,
        sequence: item.sequence,
        kind: item.kind,
        message: item.message,
        metadataJson: JSON.stringify(item.metadata),
      });
    await ctx.runMutation(completeRef, {
      attemptId: args.attemptId,
      owner,
      classification: outcome.classification,
      summary: outcome.summary,
    });
    return null;
  },
});

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: "queued",
          leaseOwner: undefined,
          leaseExpiresAt: undefined,
        });
        const run = await ctx.db.get(attempt.runId);
        if (run)
          await ctx.scheduler.runAfter(0, executeRef, {
            attemptId: attempt._id,
            scenario: run.scenario,
          });
        reconciled += 1;
      }
    }
    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();
  },
});