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

import { defaultCodingKitV0 } from "@code/primitives/resolver";
import { WorkflowManager } from "@convex-dev/workflow";
import type { WorkflowId } from "@convex-dev/workflow";
import { ConvexError, v } from "convex/values";

import { components, internal } from "./_generated/api";
import type { Doc, Id } from "./_generated/dataModel";
import { internalMutation, internalQuery, mutation } from "./_generated/server";
import type { MutationCtx } from "./_generated/server";
import { requireProjectMember } from "./authz";

export const workflow = new WorkflowManager(components.workflow);

const resolveReadySlice = async (
  ctx: MutationCtx,
  work: Doc<"works">,
  sliceId?: string
) => {
  if (work.designVersion === undefined) {
    throw new ConvexError("Work has no approved Design to execute");
  }
  const slices = await ctx.db
    .query("workSlices")
    .withIndex("by_workId_and_designVersion", (q) =>
      q.eq("workId", work._id).eq("designVersion", work.designVersion!)
    )
    .collect();
  const slice = sliceId
    ? slices.find((candidate) => candidate.sliceId === sliceId)
    : slices
        .sort((left, right) => left.ordinal - right.ordinal)
        .find((candidate) => candidate.status === "ready");
  if (!slice || slice.status !== "ready") {
    throw new ConvexError("Only the next ready slice can be executed");
  }
  return slice;
};

export const execute = workflow
  .define({ args: { attemptId: v.id("workAttempts") } })
  .handler(async (step, args): Promise<void> => {
    try {
      const result = await step.runAction(
        internal.workExecutionAgent.executeAttempt,
        args,
        { retry: true }
      );
      await step.runMutation(internal.workExecutionWorkflow.completeAttempt, {
        attemptId: args.attemptId,
        result: {
          ...result,
          changedFiles: [...result.changedFiles],
          events: result.events.map((item) => ({
            ...item,
            metadata: { ...item.metadata },
          })),
        },
      });
    } catch (error) {
      await step.runMutation(internal.workExecutionWorkflow.failAttempt, {
        attemptId: args.attemptId,
        summary: error instanceof Error ? error.message : "Execution failed",
      });
    }
  });

export const startExecution = mutation({
  args: {
    sliceId: v.optional(v.string()),
    workId: v.id("works"),
  },
  handler: async (
    ctx,
    args
  ): Promise<{
    attemptId: Id<"workAttempts">;
    runId: Id<"workRuns">;
    workflowId: WorkflowId;
  }> => {
    const work = await ctx.db.get(args.workId);
    if (!work) {
      throw new ConvexError("Work not found");
    }
    await requireProjectMember(ctx, work.projectId);
    if (work.status !== "ready") {
      throw new ConvexError("Work must be Ready before execution");
    }
    if (
      work.definitionApprovalVersion !== work.definitionVersion ||
      work.designApprovalVersion !== work.designVersion
    ) {
      throw new ConvexError("Execution requires exact approved versions");
    }
    const project = await ctx.db.get(work.projectId);
    if (!project?.gitConnectionId) {
      throw new ConvexError("Connect Git credentials to this project first");
    }
    const slice = await resolveReadySlice(ctx, work, args.sliceId);
    const createdAt = Date.now();
    const runId = await ctx.db.insert("workRuns", {
      createdAt,
      designVersion: slice.designVersion,
      executionKind: "real",
      kitId: defaultCodingKitV0.id,
      kitVersion: defaultCodingKitV0.version,
      scenario: "success",
      sliceId: slice.sliceId,
      sliceRowId: slice._id,
      status: "running",
      workId: work._id,
    });
    const workspaceKey = `work-${work._id}-run-${runId}`;
    const attemptId = await ctx.db.insert("workAttempts", {
      number: 1,
      runId,
      status: "queued",
      workId: work._id,
      workspaceKey,
    });
    const workflowId: WorkflowId = await workflow.start(
      ctx,
      internal.workExecutionWorkflow.execute,
      { attemptId }
    );
    await ctx.db.patch(runId, { startedAt: createdAt, workflowId });
    await ctx.db.patch(slice._id, { status: "running" });
    await ctx.db.patch(work._id, { status: "executing", updatedAt: createdAt });
    await ctx.db.insert("workEvents", {
      createdAt,
      idempotencyKey: `real-run-started:${runId}`,
      kind: "run.started",
      referenceId: String(runId),
      workId: work._id,
    });
    return { attemptId, runId, workflowId };
  },
});

export const executionContext = internalQuery({
  args: { attemptId: v.id("workAttempts") },
  handler: async (ctx, args) => {
    const attempt = await ctx.db.get(args.attemptId);
    if (!attempt) {
      throw new ConvexError("Attempt not found");
    }
    const run = await ctx.db.get(attempt.runId);
    const work = await ctx.db.get(attempt.workId);
    if (!run || !work || !run.sliceRowId) {
      throw new ConvexError("Execution records are incomplete");
    }
    const [slice, project] = await Promise.all([
      ctx.db.get(run.sliceRowId),
      ctx.db.get(work.projectId),
    ]);
    if (!slice || !project?.gitConnectionId) {
      throw new ConvexError("Project execution configuration is incomplete");
    }
    const connection = await ctx.db.get(project.gitConnectionId);
    if (!connection) {
      throw new ConvexError("Git connection not found");
    }
    return {
      attempt,
      connection,
      project,
      prompt: [
        `Implement this approved Zopu slice: ${slice.title}`,
        `Objective: ${slice.objective}`,
        `Observable behavior: ${slice.observableBehavior}`,
        "Inspect the repository instructions first. Make focused changes and run relevant checks.",
      ].join("\n\n"),
      run,
      work,
    };
  },
});

export const markAttemptRunning = internalMutation({
  args: { attemptId: v.id("workAttempts") },
  handler: async (ctx, args) => {
    const attempt = await ctx.db.get(args.attemptId);
    if (!attempt || attempt.status !== "queued") {
      return false;
    }
    await ctx.db.patch(attempt._id, {
      startedAt: Date.now(),
      status: "running",
    });
    return true;
  },
});

const settleSliceAndWork = async (
  ctx: MutationCtx,
  run: Doc<"workRuns">,
  succeeded: boolean
) => {
  if (!run.sliceRowId) {
    return;
  }
  const slice = await ctx.db.get(run.sliceRowId);
  if (!slice) {
    return;
  }
  await ctx.db.patch(slice._id, { status: succeeded ? "completed" : "ready" });
  const work = await ctx.db.get(run.workId);
  if (!work) {
    return;
  }
  let status: Doc<"works">["status"] = succeeded ? "completed" : "failed";
  if (succeeded) {
    const slices = await ctx.db
      .query("workSlices")
      .withIndex("by_workId_and_designVersion", (q) =>
        q.eq("workId", work._id).eq("designVersion", slice.designVersion)
      )
      .collect();
    const next = slices
      .sort((left, right) => left.ordinal - right.ordinal)
      .find((candidate) => candidate.status === "planned");
    if (next) {
      await ctx.db.patch(next._id, { status: "ready" });
      status = "ready";
    }
  }
  await ctx.db.patch(work._id, { status, updatedAt: Date.now() });
};

export const completeAttempt = internalMutation({
  args: {
    attemptId: v.id("workAttempts"),
    result: v.object({
      baseRevision: v.string(),
      candidateRevision: v.string(),
      changedFiles: v.array(v.string()),
      diff: v.string(),
      environmentId: v.string(),
      events: v.array(
        v.object({
          kind: v.string(),
          message: v.string(),
          metadata: v.record(v.string(), v.string()),
          occurredAt: v.number(),
          sequence: v.number(),
        })
      ),
      summary: v.string(),
    }),
  },
  handler: async (ctx, args) => {
    const attempt = await ctx.db.get(args.attemptId);
    if (!attempt || attempt.status === "terminal") {
      return;
    }
    const run = await ctx.db.get(attempt.runId);
    const work = await ctx.db.get(attempt.workId);
    if (!run || !work) {
      throw new ConvexError("Execution records not found");
    }
    for (const item of args.result.events) {
      await ctx.db.insert("workAttemptEvents", {
        attemptId: attempt._id,
        kind: item.kind,
        message: item.message,
        metadataJson: JSON.stringify(item.metadata),
        occurredAt: item.occurredAt,
        sequence: item.sequence,
      });
    }
    const endedAt = Date.now();
    await ctx.db.patch(attempt._id, {
      classification: "Succeeded",
      endedAt,
      status: "terminal",
      summary: args.result.summary,
    });
    await ctx.db.patch(run._id, {
      baseRevision: args.result.baseRevision,
      candidateRevision: args.result.candidateRevision,
      endedAt,
      environmentId: args.result.environmentId,
      status: "terminal",
      terminalClassification: "Succeeded",
      terminalSummary: args.result.summary,
    });
    await ctx.db.insert("workArtifacts", {
      attemptId: attempt._id,
      createdAt: endedAt,
      designVersion: run.designVersion,
      environmentId: args.result.environmentId,
      idempotencyKey: `real-diff:${attempt._id}`,
      kind: "diff",
      metadataJson: JSON.stringify({ changedFiles: args.result.changedFiles }),
      organizationId: work.organizationId,
      producer: "agentos-codex",
      projectId: work.projectId,
      provenanceJson: JSON.stringify({
        baseRevision: args.result.baseRevision,
      }),
      runId: run._id,
      sliceId: run.sliceId,
      sourceRevision: args.result.candidateRevision,
      title: "Implementation diff",
      verificationStatus: "unverified",
      workId: work._id,
      ...(args.result.diff.length > 0
        ? {
            uri: `data:text/plain;charset=utf-8,${encodeURIComponent(args.result.diff.slice(0, 50_000))}`,
          }
        : {}),
    });
    await settleSliceAndWork(ctx, run, true);
    await ctx.db.insert("workEvents", {
      createdAt: endedAt,
      idempotencyKey: `real-run-completed:${run._id}`,
      kind: "run.completed",
      payloadJson: JSON.stringify({ classification: "Succeeded" }),
      referenceId: String(run._id),
      workId: work._id,
    });
  },
});

export const failAttempt = internalMutation({
  args: { attemptId: v.id("workAttempts"), summary: v.string() },
  handler: async (ctx, args) => {
    const attempt = await ctx.db.get(args.attemptId);
    if (!attempt || attempt.status === "terminal") {
      return;
    }
    const run = await ctx.db.get(attempt.runId);
    if (!run) {
      return;
    }
    const endedAt = Date.now();
    await ctx.db.patch(attempt._id, {
      classification: "PermanentFailure",
      endedAt,
      status: "terminal",
      summary: args.summary,
    });
    await ctx.db.patch(run._id, {
      endedAt,
      status: "terminal",
      terminalClassification: "PermanentFailure",
      terminalSummary: args.summary,
    });
    await settleSliceAndWork(ctx, run, false);
  },
});

export const cancelExecution = 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");
    }
    await requireProjectMember(ctx, (await ctx.db.get(run.workId))!.projectId);
    if (run.status !== "running") {
      return { cancelled: false };
    }
    if (run.workflowId) {
      await workflow.cancel(ctx, run.workflowId as WorkflowId);
    }
    const attempts = await ctx.db
      .query("workAttempts")
      .withIndex("by_runId_and_number", (q) => q.eq("runId", run._id))
      .collect();
    const attempt = attempts.find(
      (candidate) => candidate.status !== "terminal"
    );
    if (attempt?.workspaceKey) {
      await ctx.scheduler.runAfter(
        0,
        internal.workExecutionAgent.cancelAttempt,
        {
          workspaceKey: attempt.workspaceKey,
        }
      );
      await ctx.db.patch(attempt._id, {
        classification: "Cancelled",
        endedAt: Date.now(),
        status: "terminal",
        summary: "Execution cancelled",
      });
    }
    await ctx.db.patch(run._id, {
      endedAt: Date.now(),
      status: "cancelled",
      terminalClassification: "Cancelled",
      terminalSummary: "Execution cancelled",
    });
    if (run.sliceRowId) {
      await ctx.db.patch(run.sliceRowId, { status: "ready" });
    }
    const work = await ctx.db.get(run.workId);
    if (work) {
      await ctx.db.patch(work._id, { status: "ready", updatedAt: Date.now() });
    }
    return { cancelled: true };
  },
});