Submatter/apps/worker/src/db/task-lifecycle.ts

import { and, eq, isNull } from "drizzle-orm";

import type { db as dbInstance } from "@tasks/database";
import { agentRuns, auditLog, markdownBacklogItems } from "@tasks/database/schema";
import { exportBacklogItemToMarkdown, resolveWorkflowPrompt } from "@tasks/database/markdown-backlog";

type Db = typeof dbInstance;

export async function findBacklogItemByRepoPath(
  db: Db,
  workspaceId: string,
  repoPath: string,
) {
  const [row] = await db
    .select({
      id: markdownBacklogItems.id,
      title: markdownBacklogItems.title,
      status: markdownBacklogItems.status,
      repoPath: markdownBacklogItems.repoPath,
    })
    .from(markdownBacklogItems)
    .where(
      and(
        eq(markdownBacklogItems.workspaceId, workspaceId),
        eq(markdownBacklogItems.repoPath, repoPath),
      ),
    )
    .limit(1);

  return row ?? null;
}

export function pickTaskRepoPath(files: string[]): string | null {
  const taskFile = files.find((file) => /Task-[^/]+\.md$/i.test(file));
  return taskFile ?? files[0] ?? null;
}

export async function claimTask(
  db: Db,
  input: {
    workspaceId: string;
    backlogItemId: string;
    client?: string;
    model?: string;
  },
): Promise<{ runId: string; workflowPrompt: string }> {
  const [item] = await db
    .select({
      id: markdownBacklogItems.id,
      workspaceId: markdownBacklogItems.workspaceId,
      status: markdownBacklogItems.status,
    })
    .from(markdownBacklogItems)
    .where(eq(markdownBacklogItems.id, input.backlogItemId))
    .limit(1);

  if (!item || item.workspaceId !== input.workspaceId) {
    throw new Error(`Backlog item ${input.backlogItemId} not found in workspace`);
  }

  const [existing] = await db
    .select({ id: agentRuns.id })
    .from(agentRuns)
    .where(
      and(eq(agentRuns.backlogItemId, input.backlogItemId), isNull(agentRuns.finishedAt)),
    )
    .limit(1);

  if (existing) {
    const resolved = await resolveWorkflowPrompt(db, {
      workspaceId: input.workspaceId,
      backlogItemId: input.backlogItemId,
    });
    return { runId: existing.id, workflowPrompt: resolved.prompt };
  }

  const priorStatus = item.status;
  const shouldTransition = priorStatus === "ready" || priorStatus === "draft";
  const metadata =
    input.client || input.model
      ? {
          ...(input.client ? { client: input.client } : {}),
          ...(input.model ? { model: input.model } : {}),
        }
      : null;

  const runId = await db.transaction(async (tx) => {
    const [created] = await tx
      .insert(agentRuns)
      .values({
        workspaceId: input.workspaceId,
        backlogItemId: input.backlogItemId,
        actorUserId: null,
        metadata,
      })
      .returning({ id: agentRuns.id });

    if (shouldTransition) {
      await tx
        .update(markdownBacklogItems)
        .set({ status: "in_progress", updatedAt: new Date() })
        .where(eq(markdownBacklogItems.id, input.backlogItemId));
    }

    await tx.insert(auditLog).values({
      workspaceId: input.workspaceId,
      actorUserId: null,
      action: "task.claimed",
      targetType: "agent_run",
      targetId: created.id,
      metadata: {
        backlog_item_id: input.backlogItemId,
        prior_status: priorStatus,
        transitioned: shouldTransition,
        system_actor: "echodo-worker",
        ...(metadata ?? {}),
      },
    });

    return created.id;
  });

  const resolved = await resolveWorkflowPrompt(db, {
    workspaceId: input.workspaceId,
    backlogItemId: input.backlogItemId,
  });

  return { runId, workflowPrompt: resolved.prompt };
}

export async function completeTask(
  db: Db,
  input: {
    runId: string;
    outcome: "succeeded" | "failed" | "cancelled" | "stalled";
    tokensTotal?: number;
    tokensInput?: number;
    tokensOutput?: number;
    notes?: string;
    error?: string;
    finalStatus?: "done" | "blocked" | "ready" | "in_progress" | "cancelled";
    modelAlias?: string;
    modelTier?: string;
    providerType?: string;
    modelId?: string;
    harnessAlias?: string | null;
    repoRoot?: string | null;
  },
) {
  const [run] = await db.select().from(agentRuns).where(eq(agentRuns.id, input.runId)).limit(1);
  if (!run) throw new Error("Run not found");
  if (run.finishedAt !== null) throw new Error("Run already finished");

  const finalStatus =
    input.finalStatus ??
    (input.outcome === "succeeded" ? "done" : input.outcome === "failed" ? "blocked" : null);

  const tokensTotal =
    input.tokensTotal ??
    (typeof input.tokensInput === "number" && typeof input.tokensOutput === "number"
      ? input.tokensInput + input.tokensOutput
      : null);

  await db.transaction(async (tx) => {
    await tx
      .update(agentRuns)
      .set({
        finishedAt: new Date(),
        outcome: input.outcome,
        tokensInput: input.tokensInput ?? null,
        tokensOutput: input.tokensOutput ?? null,
        tokensTotal,
        notes: input.notes ?? null,
        error: input.error ?? null,
        modelAlias: input.modelAlias ?? null,
        modelTier: input.modelTier ?? null,
        providerType: input.providerType ?? null,
        modelId: input.modelId ?? null,
        harnessAlias: input.harnessAlias ?? null,
        metadata: {
          ...(run.metadata ?? {}),
          client: "echodo-worker",
          harnessAlias: input.harnessAlias ?? null,
        },
      })
      .where(eq(agentRuns.id, input.runId));

    if (finalStatus) {
      await tx
        .update(markdownBacklogItems)
        .set({ status: finalStatus, updatedAt: new Date() })
        .where(eq(markdownBacklogItems.id, run.backlogItemId));
    }

    await tx.insert(auditLog).values({
      workspaceId: run.workspaceId,
      actorUserId: null,
      action: "task.completed",
      targetType: "agent_run",
      targetId: input.runId,
      metadata: {
        outcome: input.outcome,
        finalStatus,
        backlogItemId: run.backlogItemId,
        tokensTotal,
        system_actor: "echodo-worker",
      },
    });
  });

  if (input.repoRoot && finalStatus) {
    await exportBacklogItemToMarkdown(db, {
      workspaceId: run.workspaceId,
      backlogItemId: run.backlogItemId,
      repoRootAbs: input.repoRoot,
    });
  }
}