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 { resolveWorkspaceHandle } from "../lib/resolve-workspace.js";
import { resolveWorkflowPrompt } from "../../../../packages/database/src/markdown-backlog/resolve-prompt.ts";
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()
    .describe("UUID of the markdown_backlog_items row being claimed."),
  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.",
    ),
});

/**
 * 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.",
      inputSchema: claimTaskInputSchema,
    },
    async (args) => {
      try {
        const input = claimTaskInputSchema.parse(args);
        const ws = await resolveWorkspaceHandle(input.workspace);

        const [item] = await db
          .select({
            id: markdownBacklogItems.id,
            workspaceId: markdownBacklogItems.workspaceId,
            title: markdownBacklogItems.title,
            bodyMarkdown: markdownBacklogItems.bodyMarkdown,
            status: markdownBacklogItems.status,
          })
          .from(markdownBacklogItems)
          .where(eq(markdownBacklogItems.id, input.backlogItemId))
          .limit(1);

        if (!item) {
          return toolErr(`Backlog item ${input.backlogItemId} not found`);
        }
        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 ${input.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, input.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: input.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,
            },
            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: input.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, input.backlogItemId));
          }

          await tx.insert(auditLog).values({
            workspaceId: ws.id,
            actorUserId,
            action: "task.claimed",
            targetType: "agent_run",
            targetId: created.id,
            metadata: {
              backlog_item_id: input.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: input.backlogItemId,
        });

        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,
          },
          workspace: { id: ws.id, slug: ws.slug, name: ws.name },
        });
      } catch (e) {
        return toolCatch(e);
      }
    },
  );
}