@desplega.ai/agent-swarm 1.0.1

package/src/tools/join-swarm.ts

@@ -0,0 +1,102 @@
+import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import * as z from "zod";
+import { createAgent, getAllAgents, getDb } from "@/be/db";
+import { createToolRegistrar } from "@/tools/utils";
+import { AgentSchema } from "@/types";
+
+export const registerJoinSwarmTool = (server: McpServer) => {
+  createToolRegistrar(server)(
+    "join-swarm",
+    {
+      title: "Join the agent swarm",
+      description: "Tool for an agent to join the swarm of agents.",
+      inputSchema: z.object({
+        lead: z.boolean().default(false).describe("Whether this agent should be the lead."),
+        name: z.string().min(1).describe("The name of the agent joining the swarm."),
+      }),
+      outputSchema: z.object({
+        success: z.boolean(),
+        message: z.string(),
+        agent: AgentSchema.optional(),
+      }),
+    },
+    async ({ lead, name }, requestInfo, _meta) => {
+      // Check if agent ID is set
+      if (!requestInfo.agentId) {
+        return {
+          content: [
+            {
+              type: "text",
+              text: 'Agent ID not found. The MCP client should define the "X-Agent-ID" header.',
+            },
+          ],
+          structuredContent: {
+            yourAgentId: requestInfo.agentId,
+            success: false,
+            message: 'Agent ID not found. The MCP client should define the "X-Agent-ID" header.',
+          },
+        };
+      }
+
+      const agentId = requestInfo.agentId;
+
+      try {
+        const agentTx = getDb().transaction(() => {
+          const agents = getAllAgents();
+
+          const existingAgent = agents.find((agent) => agent.name === name);
+          const existingLead = agents.find((agent) => agent.isLead);
+
+          if (existingAgent) {
+            throw new Error(`Agent with name "${name}" already exists.`);
+          }
+
+          // If lead is true, demote e
+          if (lead && existingLead) {
+            throw new Error(
+              `Lead agent "${existingLead.name}" already exists. Only one lead agent is allowed.`,
+            );
+          }
+
+          return createAgent({
+            id: agentId,
+            name,
+            isLead: lead,
+            status: "idle",
+          });
+        });
+
+        const agent = agentTx();
+
+        return {
+          content: [
+            {
+              type: "text",
+              text: `Successfully joined swarm as agent "${agent.name}" (ID: ${agent.id}).`,
+            },
+          ],
+          structuredContent: {
+            yourAgentId: requestInfo.agentId,
+            success: true,
+            message: `Successfully joined swarm as agent "${agent.name}" (ID: ${agent.id}).`,
+            agent,
+          },
+        };
+      } catch (error) {
+        return {
+          content: [
+            {
+              type: "text",
+              text: `Failed to join swarm: ${(error as Error).message}`,
+            },
+          ],
+          structuredContent: {
+            yourAgentId: requestInfo.agentId,
+            success: false,
+            message: `Failed to join swarm: ${(error as Error).message}`,
+          },
+        };
+      }
+    },
+  );
+};

package/src/tools/my-agent-info.ts

@@ -0,0 +1,61 @@
+import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import * as z from "zod";
+import { getAgentById } from "@/be/db";
+import { createToolRegistrar } from "@/tools/utils";
+
+export const registerMyAgentInfoTool = (server: McpServer) => {
+  createToolRegistrar(server)(
+    "my-agent-info",
+    {
+      title: "Get your agent info",
+      description: "Returns your agent ID based on the X-Agent-ID header.",
+      inputSchema: z.object({}),
+      outputSchema: z.object({
+        success: z.boolean(),
+        message: z.string(),
+        agentId: z.string().optional(),
+      }),
+    },
+    async (_input, requestInfo, _meta) => {
+      if (!requestInfo.agentId) {
+        return {
+          content: [
+            {
+              type: "text",
+              text: 'Agent ID not found. The MCP client should define the "X-Agent-ID" header.',
+            },
+          ],
+          structuredContent: {
+            yourAgentId: requestInfo.agentId,
+            success: false,
+            message: 'Agent ID not found. The MCP client should define the "X-Agent-ID" header.',
+          },
+        };
+      }
+
+      const maybeAgent = getAgentById(requestInfo.agentId);
+
+      let registeredMessage =
+        " You are not registered as an agent, use the 'join-swarm' tool to register, use a nice name related to the project you are working on if not provided by the user.";
+
+      if (maybeAgent) {
+        registeredMessage = ` You are registered as agent "${maybeAgent.name}".`;
+      }
+
+      return {
+        content: [
+          {
+            type: "text",
+            text: `Your agent ID is: ${requestInfo.agentId}.${registeredMessage}`,
+          },
+        ],
+        structuredContent: {
+          yourAgentId: requestInfo.agentId,
+          yourAgentInfo: maybeAgent,
+          success: true,
+          message: `Your agent ID is: ${requestInfo.agentId}.${registeredMessage}`,
+        },
+      };
+    },
+  );
+};

package/src/tools/poll-task.ts

@@ -0,0 +1,109 @@
+import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import { addMinutes } from "date-fns";
+import * as z from "zod";
+import { getDb, getPendingTaskForAgent, startTask } from "@/be/db";
+import { createToolRegistrar } from "@/tools/utils";
+import { AgentTaskSchema } from "@/types";
+
+const DEFAULT_POLL_INTERVAL_MS = 2000;
+const MAX_POLL_DURATION_MS = 1 * 60 * 1000;
+
+export const registerPollTaskTool = (server: McpServer) => {
+  createToolRegistrar(server)(
+    "poll-task",
+    {
+      title: "Poll for a task",
+      description:
+        "Tool for an agent to poll for a new task assignment, to be used recursively until a task is assigned.",
+      inputSchema: z.object({}),
+      outputSchema: z.object({
+        success: z.boolean(),
+        message: z.string(),
+        task: AgentTaskSchema.optional(),
+        waitedForSeconds: z.number().describe("Seconds waited before receiving the task."),
+      }),
+    },
+    async (_input, requestInfo, meta) => {
+      // Check if agent ID is set
+      if (!requestInfo.agentId) {
+        return {
+          content: [
+            {
+              type: "text",
+              text: 'Agent ID not found. The MCP client should define the "X-Agent-ID" header.',
+            },
+          ],
+          structuredContent: {
+            yourAgentId: requestInfo.agentId,
+            success: false,
+            message: 'Agent ID not found. The MCP client should define the "X-Agent-ID" header.',
+            waitedForSeconds: 0,
+          },
+        };
+      }
+
+      const agentId = requestInfo.agentId;
+      const now = new Date();
+      const maxTime = addMinutes(now, MAX_POLL_DURATION_MS / 60000);
+
+      // Poll for pending tasks
+      while (new Date() < maxTime) {
+        // Fetch and update in a single transaction to avoid race conditions
+        const startedTask = getDb().transaction(() => {
+          const pendingTask = getPendingTaskForAgent(agentId);
+          if (!pendingTask) return null;
+          return startTask(pendingTask.id);
+        })();
+
+        if (startedTask) {
+          const waitedFor = Math.round((Date.now() - now.getTime()) / 1000);
+
+          return {
+            content: [
+              {
+                type: "text",
+                text: `Task "${startedTask.id}" assigned and started.`,
+              },
+            ],
+            structuredContent: {
+              yourAgentId: requestInfo.agentId,
+              success: true,
+              message: `Task "${startedTask.id}" assigned and started.`,
+              task: startedTask,
+              waitedForSeconds: waitedFor,
+            },
+          };
+        }
+
+        await meta.sendNotification({
+          method: "notifications/message",
+          params: {
+            level: "info",
+            data: `Polling for task assignment...`,
+          },
+        });
+
+        // Wait for a short period before polling again
+        await new Promise((resolve) => setTimeout(resolve, DEFAULT_POLL_INTERVAL_MS));
+      }
+
+      const waitedForSeconds = Math.round((Date.now() - now.getTime()) / 1000);
+
+      // If no task was found within the time limit
+      return {
+        content: [
+          {
+            type: "text",
+            text: `No task assigned within the polling duration.`,
+          },
+        ],
+        structuredContent: {
+          yourAgentId: requestInfo.agentId,
+          success: false,
+          message: `No task assigned within the polling duration, please keep polling until a task is assigned.`,
+          waitedForSeconds,
+        },
+      };
+    },
+  );
+};

package/src/tools/send-task.ts

@@ -0,0 +1,62 @@
+import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import * as z from "zod";
+import { createTask, getAgentById, getDb, updateAgentStatus } from "@/be/db";
+import { createToolRegistrar } from "@/tools/utils";
+import { AgentTaskSchema } from "@/types";
+
+export const registerSendTaskTool = (server: McpServer) => {
+  createToolRegistrar(server)(
+    "send-task",
+    {
+      title: "Send a task",
+      description: "Sends a task to a specific agent in the swarm.",
+      inputSchema: z.object({
+        agentId: z.uuid().describe("The ID of the agent to send the task to."),
+        task: z.string().min(1).describe("The task description to send."),
+      }),
+      outputSchema: z.object({
+        success: z.boolean(),
+        message: z.string(),
+        task: AgentTaskSchema.optional(),
+      }),
+    },
+    async ({ agentId, task }, requestInfo, _meta) => {
+      const txn = getDb().transaction(() => {
+        const agent = getAgentById(agentId);
+
+        if (!agent) {
+          return {
+            success: false,
+            message: `Agent with ID "${agentId}" not found.`,
+          };
+        }
+
+        if (agent.status !== "idle") {
+          return {
+            success: false,
+            message: `Agent "${agent.name}" is not idle (status: ${agent.status}). Cannot assign task.`,
+          };
+        }
+
+        const newTask = createTask(agentId, task);
+        updateAgentStatus(agentId, "busy");
+
+        return {
+          success: true,
+          message: `Task "${newTask.id}" sent to agent "${agent.name}".`,
+          task: newTask,
+        };
+      });
+
+      const result = txn();
+
+      return {
+        content: [{ type: "text", text: result.message }],
+        structuredContent: {
+          yourAgentId: requestInfo.agentId,
+          ...result,
+        },
+      };
+    },
+  );
+};

package/src/tools/store-progress.ts

@@ -0,0 +1,94 @@
+import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import * as z from "zod";
+import {
+  completeTask,
+  failTask,
+  getDb,
+  getTaskById,
+  updateAgentStatus,
+  updateTaskProgress,
+} from "@/be/db";
+import { createToolRegistrar } from "@/tools/utils";
+import { AgentTaskSchema } from "@/types";
+
+export const registerStoreProgressTool = (server: McpServer) => {
+  createToolRegistrar(server)(
+    "store-progress",
+    {
+      title: "Store task progress",
+      description:
+        "Stores the progress of a specific task. Can also mark task as completed or failed, which will set the agent back to idle.",
+      inputSchema: z.object({
+        taskId: z.uuid().describe("The ID of the task to update progress for."),
+        progress: z.string().optional().describe("The progress update to store."),
+        status: z
+          .enum(["completed", "failed"])
+          .optional()
+          .describe("Set to 'completed' or 'failed' to finish the task."),
+        output: z.string().optional().describe("The output of the task (used when completing)."),
+        failureReason: z
+          .string()
+          .optional()
+          .describe("The reason for failure (used when failing)."),
+      }),
+      outputSchema: z.object({
+        success: z.boolean(),
+        message: z.string(),
+        task: AgentTaskSchema.optional(),
+      }),
+    },
+    async ({ taskId, progress, status, output, failureReason }, requestInfo, _meta) => {
+      const txn = getDb().transaction(() => {
+        const existingTask = getTaskById(taskId);
+
+        if (!existingTask) {
+          return {
+            success: false,
+            message: `Task with ID "${taskId}" not found.`,
+          };
+        }
+
+        let updatedTask = existingTask;
+
+        // Update progress if provided
+        if (progress) {
+          const result = updateTaskProgress(taskId, progress);
+          if (result) updatedTask = result;
+        }
+
+        // Handle status change
+        if (status === "completed") {
+          const result = completeTask(taskId, output);
+          if (result) {
+            updatedTask = result;
+            updateAgentStatus(existingTask.agentId, "idle");
+          }
+        } else if (status === "failed") {
+          const result = failTask(taskId, failureReason ?? "Unknown failure");
+          if (result) {
+            updatedTask = result;
+            updateAgentStatus(existingTask.agentId, "idle");
+          }
+        }
+
+        return {
+          success: true,
+          message: status
+            ? `Task "${taskId}" marked as ${status}.`
+            : `Progress stored for task "${taskId}".`,
+          task: updatedTask,
+        };
+      });
+
+      const result = txn();
+
+      return {
+        content: [{ type: "text", text: result.message }],
+        structuredContent: {
+          yourAgentId: requestInfo.agentId,
+          ...result,
+        },
+      };
+    },
+  );
+};

package/src/tools/utils.ts

@@ -0,0 +1,115 @@
+import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import type {
+  AnySchema,
+  SchemaOutput,
+  ShapeOutput,
+  ZodRawShapeCompat,
+} from "@modelcontextprotocol/sdk/server/zod-compat.js";
+import type { RequestHandlerExtra } from "@modelcontextprotocol/sdk/shared/protocol.js";
+import type {
+  CallToolResult,
+  ServerNotification,
+  ServerRequest,
+  ToolAnnotations,
+} from "@modelcontextprotocol/sdk/types.js";
+
+type Meta = RequestHandlerExtra<ServerRequest, ServerNotification>;
+
+export type RequestInfo = {
+  sessionId: string | undefined;
+  agentId: string | undefined;
+};
+
+export const getRequestInfo = (req: Meta): RequestInfo => {
+  const agentIdHeader = req.requestInfo?.headers?.["x-agent-id"];
+
+  let agentId: string | undefined;
+
+  if (Array.isArray(agentIdHeader)) {
+    agentId = agentIdHeader?.[0];
+  } else if (typeof agentIdHeader === "string") {
+    agentId = agentIdHeader;
+  }
+
+  return {
+    sessionId: req.sessionId || undefined,
+    agentId,
+  };
+};
+
+// Infer the input type from the schema
+type InferInput<Args extends undefined | ZodRawShapeCompat | AnySchema> =
+  Args extends ZodRawShapeCompat
+    ? ShapeOutput<Args>
+    : Args extends AnySchema
+      ? SchemaOutput<Args>
+      : undefined;
+
+// Callback type with requestInfo injected as second parameter
+type ToolCallbackWithInfo<Args extends undefined | ZodRawShapeCompat | AnySchema = undefined> =
+  Args extends undefined
+    ? (requestInfo: RequestInfo, meta: Meta) => CallToolResult | Promise<CallToolResult>
+    : (
+        args: InferInput<Args>,
+        requestInfo: RequestInfo,
+        meta: Meta,
+      ) => CallToolResult | Promise<CallToolResult>;
+
+type ToolConfig<
+  InputArgs extends undefined | ZodRawShapeCompat | AnySchema,
+  OutputArgs extends ZodRawShapeCompat | AnySchema,
+> = {
+  title?: string;
+  description?: string;
+  inputSchema?: InputArgs;
+  outputSchema?: OutputArgs;
+  annotations?: ToolAnnotations;
+  _meta?: Record<string, unknown>;
+};
+
+/**
+ * Creates a tool registration helper that automatically extracts request info
+ * and passes it as the second parameter to the callback.
+ *
+ * @example
+ * const registerTool = createToolRegistrar(server);
+ *
+ * registerTool(
+ *   "my-tool",
+ *   { inputSchema: z.object({ name: z.string() }) },
+ *   async ({ name }, requestInfo, meta) => {
+ *     // requestInfo.sessionId and requestInfo.agentId are available
+ *     return { content: [{ type: "text", text: `Hello ${name}` }] };
+ *   }
+ * );
+ */
+export const createToolRegistrar = (server: McpServer) => {
+  return <
+    OutputArgs extends ZodRawShapeCompat | AnySchema,
+    InputArgs extends undefined | ZodRawShapeCompat | AnySchema = undefined,
+  >(
+    name: string,
+    config: ToolConfig<InputArgs, OutputArgs>,
+    cb: ToolCallbackWithInfo<InputArgs>,
+  ) => {
+    return server.registerTool(name, config, ((args: InferInput<InputArgs>, meta: Meta) => {
+      const requestInfo = getRequestInfo(meta);
+
+      // Handle zero-argument case
+      if (config.inputSchema === undefined) {
+        return (
+          cb as (requestInfo: RequestInfo, meta: Meta) => CallToolResult | Promise<CallToolResult>
+        )(requestInfo, meta);
+      }
+
+      // Handle with-arguments case
+      return (
+        cb as (
+          args: InferInput<InputArgs>,
+          requestInfo: RequestInfo,
+          meta: Meta,
+        ) => CallToolResult | Promise<CallToolResult>
+      )(args, requestInfo, meta);
+    }) as Parameters<typeof server.registerTool>[2]);
+  };
+};

package/src/types.ts

@@ -0,0 +1,42 @@
+import * as z from "zod";
+
+export const AgentTaskStatusSchema = z.enum(["pending", "in_progress", "completed", "failed"]);
+
+export const AgentTaskSchema = z.object({
+  id: z.uuid(),
+  agentId: z.uuid(),
+  task: z.string().min(1),
+  status: AgentTaskStatusSchema,
+
+  createdAt: z.iso.datetime().default(() => new Date().toISOString()),
+  lastUpdatedAt: z.iso.datetime().default(() => new Date().toISOString()),
+
+  finishedAt: z.iso.datetime().optional(),
+
+  failureReason: z.string().optional(),
+  output: z.string().optional(),
+  progress: z.string().optional(),
+});
+
+export const AgentStatusSchema = z.enum(["idle", "busy", "offline"]);
+
+export const AgentSchema = z.object({
+  id: z.uuid(),
+  name: z.string().min(1),
+  isLead: z.boolean().default(false),
+  status: AgentStatusSchema,
+
+  createdAt: z.iso.datetime().default(() => new Date().toISOString()),
+  lastUpdatedAt: z.iso.datetime().default(() => new Date().toISOString()),
+});
+
+export const AgentWithTasksSchema = AgentSchema.extend({
+  tasks: z.array(AgentTaskSchema).default([]),
+});
+
+export type AgentTaskStatus = z.infer<typeof AgentTaskStatusSchema>;
+export type AgentTask = z.infer<typeof AgentTaskSchema>;
+
+export type AgentStatus = z.infer<typeof AgentStatusSchema>;
+export type Agent = z.infer<typeof AgentSchema>;
+export type AgentWithTasks = z.infer<typeof AgentWithTasksSchema>;

package/tsconfig.json

@@ -0,0 +1,35 @@
+{
+  "compilerOptions": {
+    // Environment setup & latest features
+    "lib": ["ESNext"],
+    "target": "ESNext",
+    "module": "Preserve",
+    "moduleDetection": "force",
+    "jsx": "react-jsx",
+    "allowJs": true,
+
+    // Bundler mode
+    "moduleResolution": "bundler",
+    "allowImportingTsExtensions": true,
+
+    // Path aliases
+    "baseUrl": ".",
+    "paths": {
+      "@/*": ["src/*"]
+    },
+    "verbatimModuleSyntax": true,
+    "noEmit": true,
+
+    // Best practices
+    "strict": true,
+    "skipLibCheck": true,
+    "noFallthroughCasesInSwitch": true,
+    "noUncheckedIndexedAccess": true,
+    "noImplicitOverride": true,
+
+    // Some stricter flags (disabled by default)
+    "noUnusedLocals": false,
+    "noUnusedParameters": false,
+    "noPropertyAccessFromIndexSignature": false
+  }
+}