Submatter/apps/orchestrator/src/db/convoy.ts

import { and, desc, eq, inArray, isNull, lt, sql } from "drizzle-orm";

import type { db as dbInstance } from "@tasks/database";
import { agentRuns, convoyEvents, objects, workerCheckpoints } from "@tasks/database/schema";
import {
  briefObjectContentSchema,
  convoyObjectContentSchema,
  validateConvoyTransition,
  type ConvoyStatus,
} from "@tasks/shared/convoy-lifecycle";
import type { BriefSummary } from "@tasks/shared/orchestrator";

type Db = typeof dbInstance;

export type LoadedConvoy = {
  id: string;
  workspaceId: string;
  title: string;
  status: ConvoyStatus;
  skipFlags: string[];
  briefs: BriefSummary[];
  completedRoles: Set<string>;
  completedBriefIds: Set<string>;
};

function parseStatus(raw: string | null): ConvoyStatus {
  const status = raw ?? "draft";
  return status as ConvoyStatus;
}

export async function listActiveConvoys(db: Db): Promise<LoadedConvoy[]> {
  const rows = await db
    .select()
    .from(objects)
    .where(
      and(
        eq(objects.type, "project"),
        sql`${objects.content} ->> 'convoySlug' IS NOT NULL`,
        inArray(objects.status, ["draft", "ready", "in_progress"]),
      ),
    );

  const loaded: LoadedConvoy[] = [];
  for (const row of rows) {
    const content = convoyObjectContentSchema.safeParse(row.content);
    if (!content.success) continue;

    const briefRows = await db
      .select()
      .from(objects)
      .where(and(eq(objects.parentId, row.id), eq(objects.workspaceId, row.workspaceId)));

    const events = await db
      .select({ role: convoyEvents.role, eventType: convoyEvents.eventType, briefId: convoyEvents.briefId })
      .from(convoyEvents)
      .where(eq(convoyEvents.convoyId, row.id))
      .orderBy(desc(convoyEvents.occurredAt))
      .limit(200);

    const completedRoles = new Set<string>();
    const completedBriefIds = new Set<string>();
    for (const event of events) {
      if (event.eventType === "role_complete") {
        completedRoles.add(event.role);
        if (event.briefId) completedBriefIds.add(event.briefId);
      }
    }

    const briefs: BriefSummary[] = briefRows.map((brief) => {
      const parsed = briefObjectContentSchema.safeParse(brief.content);
      return {
        id: brief.id,
        briefNumber: parsed.success ? parsed.data.briefNumber : brief.sortOrder,
        status: brief.status,
        sliceDependencies: parsed.success ? parsed.data.sliceDependencies : [],
      };
    });

    loaded.push({
      id: row.id,
      workspaceId: row.workspaceId,
      title: row.title,
      status: parseStatus(row.status),
      skipFlags: content.data.skipFlags,
      briefs,
      completedRoles,
      completedBriefIds,
    });
  }

  return loaded;
}

export async function transitionConvoyInDb(
  db: Db,
  input: {
    workspaceId: string;
    convoyId: string;
    toStatus: ConvoyStatus;
    actor: string;
    reason?: string;
  },
) {
  const [convoy] = await db
    .select()
    .from(objects)
    .where(
      and(
        eq(objects.id, input.convoyId),
        eq(objects.workspaceId, input.workspaceId),
        eq(objects.type, "project"),
      ),
    )
    .limit(1);

  if (!convoy) {
    throw new Error(`Convoy ${input.convoyId} not found`);
  }

  const fromStatus = parseStatus(convoy.status);
  validateConvoyTransition({
    from: fromStatus,
    to: input.toStatus,
    actor: input.actor,
    reason: input.reason,
  });

  await db
    .update(objects)
    .set({ status: input.toStatus, updatedAt: new Date() })
    .where(eq(objects.id, convoy.id));

  await db.insert(convoyEvents).values({
    workspaceId: input.workspaceId,
    convoyId: convoy.id,
    role: input.actor,
    actor: input.actor,
    eventType: "status_change",
    metadata: {
      from: fromStatus,
      to: input.toStatus,
      reason: input.reason ?? null,
      source: "orchestrator",
    },
  });
}

export type StalledRunDispatch = {
  workspaceId: string;
  agentRunId: string;
  backlogItemId: string;
  convoyId: string;
  roleId: string;
  briefId?: string;
};

/**
 * Open agent runs past the stall threshold that have a workspace-scoped
 * checkpoint (enough context to re-dispatch with resume_from_checkpoint).
 */
export async function detectStalledRunsForResume(
  db: Db,
  stallMinutes: number,
): Promise<StalledRunDispatch[]> {
  const cutoff = new Date(Date.now() - stallMinutes * 60_000);
  const openRuns = await db
    .select({
      id: agentRuns.id,
      workspaceId: agentRuns.workspaceId,
      backlogItemId: agentRuns.backlogItemId,
      startedAt: agentRuns.startedAt,
    })
    .from(agentRuns)
    .where(and(isNull(agentRuns.finishedAt), lt(agentRuns.startedAt, cutoff)))
    .limit(50);

  if (openRuns.length === 0) return [];

  const dispatches: StalledRunDispatch[] = [];
  for (const run of openRuns) {
    const [checkpoint] = await db
      .select({
        convoyId: workerCheckpoints.convoyId,
        briefId: workerCheckpoints.briefId,
        roleId: workerCheckpoints.roleId,
      })
      .from(workerCheckpoints)
      .where(
        and(
          eq(workerCheckpoints.workspaceId, run.workspaceId),
          eq(workerCheckpoints.agentRunId, run.id),
        ),
      )
      .orderBy(desc(workerCheckpoints.checkpointSeq))
      .limit(1);

    if (!checkpoint) continue;

    dispatches.push({
      workspaceId: run.workspaceId,
      agentRunId: run.id,
      backlogItemId: run.backlogItemId,
      convoyId: checkpoint.convoyId,
      roleId: checkpoint.roleId,
      briefId: checkpoint.briefId ?? undefined,
    });
  }

  return dispatches;
}

/** @deprecated Prefer detectStalledRunsForResume — kept for log-only callers. */
export async function detectStalledConvoys(db: Db, stallMinutes: number): Promise<string[]> {
  const stalled = await detectStalledRunsForResume(db, stallMinutes);
  return stalled.map((run) => run.backlogItemId);
}