feat(mcp): complete_task tool — close agent_runs row + finalize status
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
parent
72aa2a5f0c
commit
93b6398c76
1 changed files with 140 additions and 0 deletions
140
apps/mcp-server/src/tools/complete-task.ts
Normal file
140
apps/mcp-server/src/tools/complete-task.ts
Normal file
|
|
@ -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<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();
|
||||
|
||||
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);
|
||||
}
|
||||
},
|
||||
);
|
||||
}
|
||||
Loading…
Reference in a new issue