diff --git a/apps/mcp-server/src/tools/complete-task.ts b/apps/mcp-server/src/tools/complete-task.ts new file mode 100644 index 0000000..9233ef1 --- /dev/null +++ b/apps/mcp-server/src/tools/complete-task.ts @@ -0,0 +1,140 @@ +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(), +}); + +type CompleteTaskInput = z.infer; + +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 | 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(); + + 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, + }) + .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, + }, + }); + }); + + return toolOk({ + runId: input.runId, + finishedAt: finishedAt.toISOString(), + finalStatus: finalStatus ?? null, + tokensTotal, + outcome: input.outcome, + }); + } catch (e) { + return toolCatch(e); + } + }, + ); +}