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

import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { z } from "zod";
import { modelTierSchema, providerTypeSchema } from "@tasks/shared/model-routing";
import { db } from "../db.js";
import { eq } from "../drizzle.js";
import { agentRuns, auditLog, markdownBacklogItems } from "../schema.js";
import { exportBacklogItemToMarkdown } from "../../../../packages/database/src/markdown-backlog/export.ts";
import { resolveRepoRoot } from "../lib/repo-root.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.",
    ),
  modelAlias: z
    .string()
    .min(1)
    .max(120)
    .optional()
    .describe("Registry alias used for this session (e.g. qwen-coder)."),
  modelTier: modelTierSchema.optional(),
  providerType: providerTypeSchema.optional(),
  modelId: z.string().min(1).max(200).optional(),
});

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;
        if (input.modelAlias) incomingMeta.modelAlias = input.modelAlias;
        if (input.modelTier) incomingMeta.modelTier = input.modelTier;
        if (input.providerType) incomingMeta.providerType = input.providerType;
        if (input.modelId) incomingMeta.modelId = input.modelId;
        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,
              ...(input.modelAlias !== undefined
                ? { modelAlias: input.modelAlias }
                : {}),
              ...(input.modelTier !== undefined ? { modelTier: input.modelTier } : {}),
              ...(input.providerType !== undefined
                ? { providerType: input.providerType }
                : {}),
              ...(input.modelId !== undefined ? { modelId: input.modelId } : {}),
              ...(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 ?? {}),
            },
          });
        });

        // Project the new state back to the markdown file. This closes
        // the sync loop so a subsequent importer pass is a no-op. We do
        // this OUTSIDE the txn because writing to disk inside a DB
        // transaction would hold locks across an I/O call; the worst
        // case here is "DB updated but file write failed," which the
        // operator can reconcile by re-running the importer in either
        // direction.
        let exportSummary:
          | { changed: boolean; repoPath: string }
          | { error: string }
          | null = null;
        const repoRoot = resolveRepoRoot();
        if (repoRoot && finalStatus !== null) {
          try {
            const result = await exportBacklogItemToMarkdown(db, {
              workspaceId: run.workspaceId,
              backlogItemId: run.backlogItemId,
              repoRootAbs: repoRoot,
            });
            if (result.ok) {
              exportSummary = {
                changed: result.changed,
                repoPath: result.repoPath,
              };
            } else {
              exportSummary = { error: `${result.reason}: ${result.detail}` };
              console.error(
                `[complete_task] markdown export failed (${result.reason}): ${result.detail}`,
              );
            }
          } catch (e) {
            const msg = e instanceof Error ? e.message : String(e);
            exportSummary = { error: msg };
            console.error("[complete_task] markdown export threw:", e);
          }
        }

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