Submatter/apps/mcp-server/src/tools/complete-task.ts

import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { z } from "zod";
import { db } from "../db.js";
import { eq } from "../drizzle.js";
import { agentRuns, auditLog, markdownBacklogItems } from "../schema.js";
import { toolCatch, toolErr, toolOk } from "./tool-result.js";

const outcomeSchema = z.enum(["succeeded", "failed", "cancelled", "stalled"]);

const finalStatusSchema = z.enum([
  "done",
  "blocked",
  "ready",
  "in_progress",
  "cancelled",
]);

const tokenCount = z.number().int().nonnegative();

const completeTaskInputSchema = z.object({
  runId: z.string().uuid(),
  outcome: outcomeSchema,
  tokensInput: tokenCount.optional(),
  tokensOutput: tokenCount.optional(),
  tokensTotal: tokenCount.optional(),
  notes: z.string().max(2000).optional(),
  error: z.string().max(2000).optional(),
  finalStatus: finalStatusSchema.optional(),
  client: z
    .string()
    .min(1)
    .max(50)
    .optional()
    .describe(
      "Late-bound client identifier — only needed if you didn't supply one at claim time, or if it changed mid-session. Merged into existing run metadata; existing keys take precedence unless overridden.",
    ),
  model: z
    .string()
    .min(1)
    .max(100)
    .optional()
    .describe(
      "Late-bound model handle. Same merge semantics as `client` — pass this when the model that closed the session differs from the one that opened it.",
    ),
});

type CompleteTaskInput = z.infer<typeof completeTaskInputSchema>;

function resolveTokensTotal(input: CompleteTaskInput): number | null {
  if (typeof input.tokensTotal === "number") return input.tokensTotal;
  if (
    typeof input.tokensInput === "number" &&
    typeof input.tokensOutput === "number"
  ) {
    return input.tokensInput + input.tokensOutput;
  }
  return null;
}

function resolveFinalStatus(
  input: CompleteTaskInput,
): z.infer<typeof finalStatusSchema> | null {
  if (input.finalStatus) return input.finalStatus;
  if (input.outcome === "succeeded") return "done";
  if (input.outcome === "failed") return "blocked";
  return null;
}

export function registerCompleteTaskTool(mcp: McpServer): void {
  mcp.registerTool(
    "complete_task",
    {
      description:
        "Close an agent_runs row at session end with an outcome and token totals. Optionally flips the backlog item to a terminal status. Errors if the run is already closed.",
      inputSchema: completeTaskInputSchema,
    },
    async (args) => {
      try {
        const input = completeTaskInputSchema.parse(args);

        const [run] = await db
          .select()
          .from(agentRuns)
          .where(eq(agentRuns.id, input.runId))
          .limit(1);

        if (!run) return toolErr("Run not found");
        if (run.finishedAt !== null) {
          return toolErr(
            `Run already finished (outcome: ${run.outcome ?? "unknown"})`,
          );
        }

        const [backlogItem] = await db
          .select()
          .from(markdownBacklogItems)
          .where(eq(markdownBacklogItems.id, run.backlogItemId))
          .limit(1);

        if (!backlogItem || backlogItem.workspaceId !== run.workspaceId) {
          return toolErr("Backlog item missing");
        }

        const tokensTotal = resolveTokensTotal(input);
        const finalStatus = resolveFinalStatus(input);
        const finishedAt = new Date();

        // Build the merged metadata object once, outside the txn, so we
        // can also stamp the same shape into the audit log without
        // re-reading the row. Late-bound `client`/`model` from this call
        // override whatever claim_task wrote; that's the right policy
        // because the latest caller is the one we can verify.
        const incomingMeta: Record<string, unknown> = {};
        if (input.client) incomingMeta.client = input.client;
        if (input.model) incomingMeta.model = input.model;
        const hasIncomingMeta = Object.keys(incomingMeta).length > 0;
        const mergedMeta = hasIncomingMeta
          ? { ...(run.metadata ?? {}), ...incomingMeta }
          : run.metadata;

        await db.transaction(async (tx) => {
          await tx
            .update(agentRuns)
            .set({
              finishedAt,
              outcome: input.outcome,
              tokensInput: input.tokensInput ?? null,
              tokensOutput: input.tokensOutput ?? null,
              tokensTotal,
              notes: input.notes ?? null,
              error: input.error ?? null,
              ...(hasIncomingMeta ? { metadata: mergedMeta } : {}),
            })
            .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: run.actorUserId ?? null,
            action: "task.completed",
            targetType: "agent_run",
            targetId: input.runId,
            metadata: {
              outcome: input.outcome,
              finalStatus,
              backlogItemId: run.backlogItemId,
              tokensTotal,
              ...(mergedMeta ?? {}),
            },
          });
        });

        return toolOk({
          runId: input.runId,
          finishedAt: finishedAt.toISOString(),
          finalStatus: finalStatus ?? null,
          tokensTotal,
          outcome: input.outcome,
        });
      } catch (e) {
        return toolCatch(e);
      }
    },
  );
}