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

import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { z } from "zod";

import { and, eq, isNull } from "../drizzle.js";
import { db } from "../db.js";
import {
  agentRuns,
  auditLog,
  markdownBacklogItems,
} from "../schema.js";
import {
  BACKLOG_ITEM_NOT_FOUND_MESSAGE,
  backlogItemLookupFieldsSchema,
  resolveBacklogItemById,
  resolveBacklogItemsByLookup,
} from "../lib/resolve-backlog-item.js";
import { resolveWorkspaceHandle } from "../lib/resolve-workspace.js";
import { resolveWorkflowPrompt } from "../../../../packages/database/src/markdown-backlog/resolve-prompt.ts";
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";

/**
 * `claim_task` — first message of an agent session.
 *
 * Flow:
 *   1. Resolve workspace handle (slug or UUID).
 *   2. Confirm the backlog item belongs to that workspace.
 *   3. Check for an existing open run on this backlog item. Same actor =>
 *      idempotent re-claim (return the existing run). Different actor =>
 *      `ALREADY_CLAIMED` error so the agent knows to back off.
 *   4. Insert a new `agent_runs` row; flip the backlog item to
 *      `in_progress` only if it's currently `ready` or `draft` (we don't
 *      want to clobber a `blocked` / `done` / `in_progress` status that
 *      was set deliberately).
 *   5. Resolve the effective workflow prompt (task → epic → plan → default).
 *   6. Audit the claim.
 *
 * What this tool intentionally does NOT do:
 *   - No external fetches. The MCP server speaks only to the local DB.
 *   - No timeout-based "stale claim" eviction. If a run goes stale, the
 *     operator (or `complete_task`) closes it explicitly. Symphony-style
 *     stall detection is the deferred orchestrator's job.
 *   - No claim with a stale row written into `metadata` for "see, the row
 *     was already there." If the agent needs to know why, it can read
 *     `agent_runs` separately.
 */

const claimTaskInputSchema = z
  .object({
    workspace: z
      .string()
      .min(1)
      .describe("Workspace UUID or slug (e.g. 'acme' or '550e8400-...')"),
    backlogItemId: z
      .string()
      .uuid()
      .optional()
      .describe(
        "UUID of the markdown_backlog_items row being claimed. Optional when repoPath or slug keys are provided.",
      ),
    actorUserId: z
      .string()
      .uuid()
      .nullable()
      .optional()
      .describe(
        "Optional UUID of the human user the agent is acting on behalf of. Null = anonymous dev session.",
      ),
    notes: z
      .string()
      .max(2_000)
      .optional()
      .describe("Short opener text attached to the run (e.g. session intent)."),
    client: z
      .string()
      .min(1)
      .max(50)
      .optional()
      .describe(
        "Caller identifier — e.g. 'cursor-ide', 'cursor-cloud', 'codex', 'echodo-orchestrator'. Persisted in agent_runs.metadata so the runs UI can slice by client.",
      ),
    model: z
      .string()
      .min(1)
      .max(100)
      .optional()
      .describe(
        "Model handle driving the session — e.g. 'claude-4.6-sonnet', 'gpt-5-codex'. Free-form because new models ship faster than we can rev a zod enum.",
      ),
  })
  .merge(backlogItemLookupFieldsSchema)
  .superRefine((input, ctx) => {
    const hasId = Boolean(input.backlogItemId);
    const hasRepoPath = Boolean(input.repoPath?.trim());
    const hasSlugKeys = Boolean(input.slug?.trim() && input.planSlug?.trim());
    if (!hasId && !hasRepoPath && !hasSlugKeys) {
      ctx.addIssue({
        code: z.ZodIssueCode.custom,
        message:
          "Provide backlogItemId, repoPath, or slug + planSlug (add epicSlug for tasks/epics).",
      });
    }
  });

/**
 * Build the metadata object stored on agent_runs at claim time. Returns
 * `null` (not an empty object) when nothing identifying was passed so
 * JSONB `metadata IS NULL` queries stay clean.
 */
function buildClaimMetadata(input: {
  client?: string;
  model?: string;
}): Record<string, unknown> | null {
  const meta: Record<string, unknown> = {};
  if (input.client) meta.client = input.client;
  if (input.model) meta.model = input.model;
  return Object.keys(meta).length === 0 ? null : meta;
}

export function registerClaimTaskTool(mcp: McpServer): void {
  mcp.registerTool(
    "claim_task",
    {
      description:
        "Open an agent_runs row against a backlog item, transition it to in_progress (if eligible), and return the effective workflow prompt for the session. Identify the item by backlogItemId, repoPath, or slug + planSlug.",
      inputSchema: claimTaskInputSchema,
    },
    async (args) => {
      try {
        const input = claimTaskInputSchema.parse(args);
        const ws = await resolveWorkspaceHandle(input.workspace);

        let item;
        let backlogItemId: string;

        if (input.backlogItemId) {
          const found = await resolveBacklogItemById(input.backlogItemId);
          if (!found) {
            return toolErr(`Backlog item ${input.backlogItemId} not found`);
          }
          item = found;
          backlogItemId = found.id;
        } else {
          const matches = await resolveBacklogItemsByLookup(ws.id, {
            repoPath: input.repoPath,
            slug: input.slug,
            planSlug: input.planSlug,
            epicSlug: input.epicSlug,
            kind: input.kind,
          });
          if (matches.length === 0) {
            return toolErr(BACKLOG_ITEM_NOT_FOUND_MESSAGE);
          }
          if (matches.length > 1) {
            return toolErr(
              `Ambiguous backlog lookup (${matches.length} matches). Add epicSlug/kind or use repoPath.`,
            );
          }
          item = matches[0]!;
          backlogItemId = item.id;
        }

        if (item.workspaceId !== ws.id) {
          // Tenancy fence: the backlog item exists but in a different
          // workspace than the resolved handle. Refuse rather than 404 so
          // the caller can't fish for cross-tenant existence.
          return toolErr(
            `Backlog item ${backlogItemId} does not belong to workspace ${ws.slug}`,
          );
        }

        const actorUserId = input.actorUserId ?? null;

        // Look for an already-open run on this backlog item. We deliberately
        // match on (backlogItemId, finished_at IS NULL) and inspect actor
        // in code rather than baking the actor match into SQL — the
        // distinction "same actor vs different actor" produces different
        // outcomes (idempotent vs error), which is clearer at the JS layer.
        const [existing] = await db
          .select({
            id: agentRuns.id,
            actorUserId: agentRuns.actorUserId,
            startedAt: agentRuns.startedAt,
            notes: agentRuns.notes,
            metadata: agentRuns.metadata,
          })
          .from(agentRuns)
          .where(
            and(
              eq(agentRuns.backlogItemId, backlogItemId),
              isNull(agentRuns.finishedAt),
            ),
          )
          .limit(1);

        if (existing) {
          const sameActor = existing.actorUserId === actorUserId;
          if (!sameActor) {
            return toolErr(
              `ALREADY_CLAIMED: backlog item is already claimed by ${existing.actorUserId ?? "anonymous"} (run ${existing.id})`,
            );
          }
          // Idempotent re-claim. Refresh `notes` if the caller provided
          // new ones; leave timing untouched (started_at is sacred).
          // For client/model, merge into existing metadata so a fresh
          // claim from a different model (e.g. user switched models in
          // Cursor mid-session) updates the attribution without losing
          // prior keys.
          const incomingMeta = buildClaimMetadata(input);
          const mergedMeta =
            incomingMeta === null
              ? existing.metadata
              : { ...(existing.metadata ?? {}), ...incomingMeta };
          const notesChanged = Boolean(
            input.notes && input.notes !== existing.notes,
          );
          const metaChanged = incomingMeta !== null;
          if (notesChanged || metaChanged) {
            await db
              .update(agentRuns)
              .set({
                ...(notesChanged ? { notes: input.notes! } : {}),
                ...(metaChanged ? { metadata: mergedMeta } : {}),
              })
              .where(eq(agentRuns.id, existing.id));
          }

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

          return toolOk({
            runId: existing.id,
            reused: true,
            workflowPrompt: resolved.prompt,
            workflowPromptSource: resolved.source,
            backlogItem: {
              id: item.id,
              title: item.title,
              bodyMarkdown: item.bodyMarkdown,
              status: item.status,
              repoPath: item.repoPath,
            },
            workspace: { id: ws.id, slug: ws.slug, name: ws.name },
          });
        }

        // No open run — open one + transition status + audit, all in one txn.
        const priorStatus = item.status;
        const shouldTransition =
          priorStatus === "ready" || priorStatus === "draft";

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

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

          await tx.insert(auditLog).values({
            workspaceId: ws.id,
            actorUserId,
            action: "task.claimed",
            targetType: "agent_run",
            targetId: created.id,
            metadata: {
              backlog_item_id: backlogItemId,
              prior_status: priorStatus,
              transitioned: shouldTransition,
              ...(claimMeta ?? {}),
              ...(actorUserId === null ? { system_actor: "mcp:claim_task" } : {}),
            },
          });

          return created;
        });

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

        // Only export when the row's status actually changed. On a
        // no-transition claim (already in_progress / blocked / done)
        // the file already matches the DB and a write would just bump
        // updated_at for no reason.
        let exportSummary:
          | { changed: boolean; repoPath: string }
          | { error: string }
          | null = null;
        const repoRoot = resolveRepoRoot();
        if (repoRoot && shouldTransition) {
          try {
            const exportResult = await exportBacklogItemToMarkdown(db, {
              workspaceId: ws.id,
              backlogItemId,
              repoRootAbs: repoRoot,
            });
            if (exportResult.ok) {
              exportSummary = {
                changed: exportResult.changed,
                repoPath: exportResult.repoPath,
              };
            } else {
              exportSummary = {
                error: `${exportResult.reason}: ${exportResult.detail}`,
              };
              console.error(
                `[claim_task] markdown export failed (${exportResult.reason}): ${exportResult.detail}`,
              );
            }
          } catch (e) {
            const msg = e instanceof Error ? e.message : String(e);
            exportSummary = { error: msg };
            console.error("[claim_task] markdown export threw:", e);
          }
        }

        return toolOk({
          runId: result.id,
          reused: false,
          workflowPrompt: resolved.prompt,
          workflowPromptSource: resolved.source,
          backlogItem: {
            id: item.id,
            title: item.title,
            bodyMarkdown: item.bodyMarkdown,
            status: shouldTransition ? "in_progress" : priorStatus,
            repoPath: item.repoPath,
          },
          workspace: { id: ws.id, slug: ws.slug, name: ws.name },
          export: exportSummary,
        });
      } catch (e) {
        return toolCatch(e);
      }
    },
  );
}