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

import { env } from "@code/env/convex";
import {
  makeIssueWorkspacePlan,
  transitionWorkspaceRun,
  WorkspaceStateError,
  WorkspaceValidationError,
} from "@code/primitives/project-workspace";
import { ConvexError, v } from "convex/values";
import { Effect } from "effect";

import { mutation, query } from "./_generated/server";
import { artifactPath } from "./artifactModel";

const requireAgent = (token: string) => {
  if (token !== env.FLUE_DB_TOKEN) {
    throw new ConvexError("Invalid agent control token");
  }
};

const agentStatus = v.union(
  v.literal("working"),
  v.literal("needs-input"),
  v.literal("completed"),
  v.literal("failed")
);

const workspacePlanFor = (input: {
  readonly issueBody: string;
  readonly issueId: string;
  readonly issueNumber: number;
  readonly issueTitle: string;
  readonly sourceUrl: string | undefined;
  readonly defaultBranch: string | undefined;
}) => {
  try {
    return Effect.runSync(
      makeIssueWorkspacePlan({
        artifacts: [],
        branchName: undefined,
        checkoutPath: undefined,
        contextFiles: [],
        defaultBranch: input.defaultBranch,
        issueBody: input.issueBody,
        issueId: input.issueId,
        issueNumber: input.issueNumber,
        issueTitle: input.issueTitle,
        sourceUrl: input.sourceUrl,
      })
    );
  } catch (error) {
    if (error instanceof WorkspaceValidationError) {
      throw new ConvexError(error.message);
    }
    throw error;
  }
};

export const get = query({
  args: { issueId: v.id("projectIssues"), token: v.string() },
  handler: async (ctx, args) => {
    requireAgent(args.token);
    const issue = await ctx.db.get("projectIssues", args.issueId);
    if (!issue) {
      throw new ConvexError("Issue not found");
    }
    const project = await ctx.db.get("projects", issue.projectId);
    if (!project) {
      throw new ConvexError("Project not found");
    }
    const artifacts = await ctx.db
      .query("projectArtifacts")
      .withIndex("by_project", (q) => q.eq("projectId", project._id))
      .take(25);
    const [source] = await ctx.db
      .query("projectSources")
      .withIndex("by_project", (q) => q.eq("projectId", project._id))
      .take(1);
    const contextDocuments = await ctx.db
      .query("projectContextDocuments")
      .withIndex("by_project", (q) => q.eq("projectId", project._id))
      .take(25);
    const run = await ctx.db
      .query("projectWorkRuns")
      .withIndex("by_issue", (q) => q.eq("issueId", issue._id))
      .unique();
    return {
      artifacts,
      contextDocuments,
      issue,
      project,
      run: run ?? null,
      source: source ?? null,
    };
  },
});

export const ensureRun = mutation({
  args: {
    daemonId: v.string(),
    issueId: v.id("projectIssues"),
    token: v.string(),
  },
  handler: async (ctx, args) => {
    requireAgent(args.token);
    const issue = await ctx.db.get("projectIssues", args.issueId);
    if (!issue) {
      throw new ConvexError("Issue not found");
    }
    const [source] = await ctx.db
      .query("projectSources")
      .withIndex("by_project", (q) => q.eq("projectId", issue.projectId))
      .take(1);
    const plan = workspacePlanFor({
      defaultBranch: source?.defaultBranch,
      issueBody: issue.body,
      issueId: String(issue._id),
      issueNumber: issue.number,
      issueTitle: issue.title,
      sourceUrl: source?.url,
    });
    const existing = await ctx.db
      .query("projectWorkRuns")
      .withIndex("by_issue", (q) => q.eq("issueId", issue._id))
      .unique();
    if (existing) {
      if (
        existing.baseBranch !== plan.baseBranch ||
        existing.branchName !== plan.branchName ||
        existing.checkoutPath !== plan.checkoutPath ||
        existing.sourceUrl !== plan.sourceUrl
      ) {
        await ctx.db.patch("projectWorkRuns", existing._id, {
          baseBranch: plan.baseBranch,
          branchName: plan.branchName,
          checkoutPath: plan.checkoutPath,
          sourceUrl: plan.sourceUrl,
          updatedAt: Date.now(),
        });
      }
      return await ctx.db.get("projectWorkRuns", existing._id);
    }
    const timestamp = Date.now();
    const actorKey = [
      "project",
      String(issue.projectId),
      "issue",
      String(issue._id),
    ];
    const runId = await ctx.db.insert("projectWorkRuns", {
      actorKey,
      agentId: String(issue._id),
      baseBranch: plan.baseBranch,
      branchName: plan.branchName,
      checkoutPath: plan.checkoutPath,
      createdAt: timestamp,
      daemonId: args.daemonId,
      issueId: issue._id,
      projectId: issue.projectId,
      sourceUrl: plan.sourceUrl,
      status: "queued",
      updatedAt: timestamp,
    });
    await ctx.db.insert("projectEvents", {
      createdAt: timestamp,
      data: {
        actorKey,
        branchName: plan.branchName,
        checkoutPath: plan.checkoutPath,
        daemonId: args.daemonId,
        sourceUrl: plan.sourceUrl,
      },
      issueId: issue._id,
      kind: "agent.workspace.created",
      projectId: issue.projectId,
      runId,
    });
    return await ctx.db.get("projectWorkRuns", runId);
  },
});

export const updateArtifact = mutation({
  args: {
    content: v.string(),
    issueId: v.id("projectIssues"),
    path: artifactPath,
    token: v.string(),
  },
  handler: async (ctx, args) => {
    requireAgent(args.token);
    if (args.content.length > 200_000) {
      throw new ConvexError("Artifact content exceeds 200000 characters");
    }
    const issue = await ctx.db.get("projectIssues", args.issueId);
    if (!issue) {
      throw new ConvexError("Issue not found");
    }
    const artifact = await ctx.db
      .query("projectArtifacts")
      .withIndex("by_project_and_path", (q) =>
        q.eq("projectId", issue.projectId).eq("path", args.path)
      )
      .unique();
    if (!artifact) {
      throw new ConvexError(`Artifact ${args.path} not found`);
    }
    const timestamp = Date.now();
    await ctx.db.patch("projectArtifacts", artifact._id, {
      content: args.content,
      revision: artifact.revision + 1,
      updatedAt: timestamp,
    });
    await ctx.db.insert("projectEvents", {
      createdAt: timestamp,
      data: { path: args.path, revision: artifact.revision + 1 },
      issueId: issue._id,
      kind: "artifact.updated",
      projectId: issue.projectId,
    });
    return artifact.revision + 1;
  },
});

export const setStatus = mutation({
  args: {
    issueId: v.id("projectIssues"),
    status: agentStatus,
    summary: v.optional(v.string()),
    token: v.string(),
  },
  handler: async (ctx, args) => {
    requireAgent(args.token);
    const issue = await ctx.db.get("projectIssues", args.issueId);
    if (!issue) {
      throw new ConvexError("Issue not found");
    }
    const run = await ctx.db
      .query("projectWorkRuns")
      .withIndex("by_issue", (q) => q.eq("issueId", issue._id))
      .unique();
    if (!run) {
      throw new ConvexError("Agent work run not found");
    }
    try {
      Effect.runSync(
        transitionWorkspaceRun({
          from: run.status,
          to: args.status,
        })
      );
    } catch (error) {
      if (error instanceof WorkspaceStateError) {
        throw new ConvexError(error.message);
      }
      throw error;
    }
    const timestamp = Date.now();
    const terminal = args.status === "completed" || args.status === "failed";
    await ctx.db.patch("projectIssues", issue._id, {
      status: args.status,
      updatedAt: timestamp,
    });
    await ctx.db.patch("projectWorkRuns", run._id, {
      status: args.status,
      summary: args.summary?.slice(0, 4000),
      updatedAt: timestamp,
      ...(args.status === "working" && run.startedAt === undefined
        ? { startedAt: timestamp }
        : {}),
      ...(terminal ? { completedAt: timestamp } : {}),
    });
    await ctx.db.insert("projectEvents", {
      createdAt: timestamp,
      data: { summary: args.summary?.slice(0, 1000) },
      issueId: issue._id,
      kind: `agent.${args.status}`,
      projectId: issue.projectId,
      runId: run._id,
    });
  },
});

const lifecycleStatus = v.union(
  v.literal("no_changes"),
  v.literal("committed"),
  v.literal("pushed"),
  v.literal("pull_request_open"),
  v.literal("failed")
);

const pullRequestMetadata = v.object({
  baseBranch: v.string(),
  branch: v.string(),
  number: v.number(),
  status: v.union(v.literal("open"), v.literal("closed"), v.literal("merged")),
  url: v.string(),
});

export const recordGiteaLifecycle = mutation({
  args: {
    baseBranch: v.string(),
    branch: v.string(),
    commitSha: v.optional(v.string()),
    error: v.optional(v.string()),
    issueId: v.id("projectIssues"),
    pullRequest: v.optional(pullRequestMetadata),
    status: lifecycleStatus,
    token: v.string(),
  },
  handler: async (ctx, args) => {
    requireAgent(args.token);
    const issue = await ctx.db.get("projectIssues", args.issueId);
    if (!issue) {
      throw new ConvexError("Issue not found");
    }
    const run = await ctx.db
      .query("projectWorkRuns")
      .withIndex("by_issue", (q) => q.eq("issueId", issue._id))
      .unique();
    if (!run) {
      throw new ConvexError("Agent work run not found");
    }

    const events = await ctx.db
      .query("projectEvents")
      .withIndex("by_issue_and_createdAt", (q) => q.eq("issueId", issue._id))
      .order("desc")
      .take(100);
    const existing = events.find((event) => {
      if (!event.kind.startsWith("gitea.")) {
        return false;
      }
      const data = event.data as {
        branch?: unknown;
        commitSha?: unknown;
        status?: unknown;
      };
      return (
        data.branch === args.branch &&
        data.commitSha === args.commitSha &&
        data.status === args.status
      );
    });
    if (existing) {
      return existing.data;
    }

    const timestamp = Date.now();
    const data = {
      baseBranch: args.baseBranch,
      branch: args.branch,
      ...(args.commitSha === undefined ? {} : { commitSha: args.commitSha }),
      ...(args.error === undefined ? {} : { error: args.error.slice(0, 2000) }),
      ...(args.pullRequest === undefined
        ? {}
        : { pullRequest: args.pullRequest }),
      status: args.status,
    };
    const artifact = await ctx.db
      .query("projectArtifacts")
      .withIndex("by_project_and_path", (q) =>
        q.eq("projectId", issue.projectId).eq("path", "artifacts.md")
      )
      .unique();
    if (artifact) {
      const pullRequestLine = args.pullRequest
        ? `- Pull request: [#${args.pullRequest.number}](${args.pullRequest.url}) (${args.pullRequest.status})\n`
        : "";
      const commitLine = args.commitSha ? `- Commit: ${args.commitSha}\n` : "";
      await ctx.db.patch("projectArtifacts", artifact._id, {
        content: `${artifact.content}\n## Git lifecycle\n\n- Status: ${args.status}\n- Branch: ${args.branch}\n- Base branch: ${args.baseBranch}\n${commitLine}${pullRequestLine}${args.error ? `- Error: ${args.error.slice(0, 1000)}\n` : ""}`,
        revision: artifact.revision + 1,
        updatedAt: timestamp,
      });
    }
    await ctx.db.insert("projectEvents", {
      createdAt: timestamp,
      data,
      issueId: issue._id,
      kind:
        args.pullRequest === undefined
          ? "gitea.lifecycle.updated"
          : "gitea.pull_request.created",
      projectId: issue.projectId,
      runId: run._id,
    });
    return data;
  },
});