Submatter/apps/worker/src/db/checkpoints.ts

import { createHash } from "node:crypto";

import type { db as dbInstance } from "@tasks/database";
import { workerCheckpoints } from "@tasks/database/schema";
import {
  WORKER_CHECKPOINT_TOOL_HISTORY_LIMIT,
  workerCheckpointPayloadSchema,
  type WorkerCheckpointFileTouch,
  type WorkerCheckpointToolStep,
} from "@tasks/shared/worker";
import { and, desc, eq, sql } from "drizzle-orm";

type Db = typeof dbInstance;

export type SaveWorkerCheckpointInput = {
  workspaceId: string;
  agentRunId: string;
  convoyId: string;
  briefId?: string;
  roleId: string;
  checkpointSeq: number;
  toolHistory: WorkerCheckpointToolStep[];
  filesTouched: WorkerCheckpointFileTouch[];
  branch: string | null;
  worktreePath: string | null;
};

export type WorkerCheckpointRow = {
  id: string;
  workspaceId: string;
  agentRunId: string;
  convoyId: string;
  briefId: string | null;
  roleId: string;
  checkpointSeq: number;
  toolHistory: WorkerCheckpointToolStep[];
  filesTouched: WorkerCheckpointFileTouch[];
  branch: string | null;
  worktreePath: string | null;
  occurredAt: Date;
};

export function hashContent(content: string): string {
  return createHash("sha256").update(content).digest("hex");
}

export function digestResult(output: string, ok: boolean): string {
  const prefix = ok ? "ok" : "err";
  return `${prefix}:${hashContent(output).slice(0, 16)}`;
}

export async function saveWorkerCheckpoint(
  db: Db,
  input: SaveWorkerCheckpointInput,
): Promise<WorkerCheckpointRow> {
  const payload = workerCheckpointPayloadSchema.parse({
    toolHistory: input.toolHistory.slice(-WORKER_CHECKPOINT_TOOL_HISTORY_LIMIT),
    filesTouched: input.filesTouched.slice(-200),
    branch: input.branch,
    worktreePath: input.worktreePath,
  });

  const [row] = await db
    .insert(workerCheckpoints)
    .values({
      workspaceId: input.workspaceId,
      agentRunId: input.agentRunId,
      convoyId: input.convoyId,
      briefId: input.briefId ?? null,
      roleId: input.roleId,
      checkpointSeq: input.checkpointSeq,
      toolHistory: payload.toolHistory,
      filesTouched: payload.filesTouched,
      branch: payload.branch,
      worktreePath: payload.worktreePath,
    })
    .onConflictDoUpdate({
      target: [workerCheckpoints.agentRunId, workerCheckpoints.checkpointSeq],
      set: {
        toolHistory: payload.toolHistory,
        filesTouched: payload.filesTouched,
        branch: payload.branch,
        worktreePath: payload.worktreePath,
        occurredAt: sql`now()`,
      },
    })
    .returning();

  if (!row) {
    throw new Error("Failed to persist worker checkpoint");
  }

  return {
    id: row.id,
    workspaceId: row.workspaceId,
    agentRunId: row.agentRunId,
    convoyId: row.convoyId,
    briefId: row.briefId,
    roleId: row.roleId,
    checkpointSeq: row.checkpointSeq,
    toolHistory: row.toolHistory,
    filesTouched: row.filesTouched,
    branch: row.branch,
    worktreePath: row.worktreePath,
    occurredAt: row.occurredAt,
  };
}

/** Latest checkpoint for an agent run — always scoped by workspace_id. */
export async function getLatestCheckpointForRun(
  db: Db,
  workspaceId: string,
  agentRunId: string,
): Promise<WorkerCheckpointRow | null> {
  const [row] = await db
    .select()
    .from(workerCheckpoints)
    .where(
      and(
        eq(workerCheckpoints.workspaceId, workspaceId),
        eq(workerCheckpoints.agentRunId, agentRunId),
      ),
    )
    .orderBy(desc(workerCheckpoints.checkpointSeq))
    .limit(1);

  if (!row) return null;

  return {
    id: row.id,
    workspaceId: row.workspaceId,
    agentRunId: row.agentRunId,
    convoyId: row.convoyId,
    briefId: row.briefId,
    roleId: row.roleId,
    checkpointSeq: row.checkpointSeq,
    toolHistory: row.toolHistory,
    filesTouched: row.filesTouched,
    branch: row.branch,
    worktreePath: row.worktreePath,
    occurredAt: row.occurredAt,
  };
}