@desplega.ai/agent-swarm 1.0.2 → 1.2.0

package/src/tools/join-swarm.ts

@@ -11,6 +11,10 @@ export const registerJoinSwarmTool = (server: McpServer) => {
       title: "Join the agent swarm",
       description: "Tool for an agent to join the swarm of agents.",
       inputSchema: z.object({
+        requestedId: z
+          .string()
+          .optional()
+          .describe("Requested ID for the agent (overridden by X-Agent-ID header)."),
         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."),
       }),
@@ -20,37 +24,45 @@ export const registerJoinSwarmTool = (server: McpServer) => {
         agent: AgentSchema.optional(),
       }),
     },
-    async ({ lead, name }, requestInfo, _meta) => {
+    async ({ lead, name, requestedId }, requestInfo, _meta) => {
       // Check if agent ID is set
-      if (!requestInfo.agentId) {
+      if (!requestInfo.agentId && !requestedId) {
         return {
           content: [
             {
               type: "text",
-              text: 'Agent ID not found. The MCP client should define the "X-Agent-ID" header.',
+              text: 'Agent ID not found. The MCP client should define the "X-Agent-ID" header, or provide a requestedId.',
             },
           ],
           structuredContent: {
-            yourAgentId: requestInfo.agentId,
+            yourAgentId: requestInfo.agentId ?? requestedId,
             success: false,
-            message: 'Agent ID not found. The MCP client should define the "X-Agent-ID" header.',
+            message:
+              'Agent ID not found. The MCP client should define the "X-Agent-ID" header, or provide a requestedId.',
           },
         };
       }
 
-      const agentId = requestInfo.agentId;
+      const agentId = requestInfo.agentId ?? requestedId ?? "";
 
       try {
         const agentTx = getDb().transaction(() => {
           const agents = getAllAgents();
 
+          const existingIdAgent = agents.find((agent) => agent.id === agentId);
+
+          if (existingIdAgent) {
+            throw new Error(`Agent with ID "${agentId}" already exists.`);
+          }
+
           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.`);
           }
 
+          const existingLead = agents.find((agent) => agent.isLead);
+
           // If lead is true, demote e
           if (lead && existingLead) {
             throw new Error(
@@ -72,13 +84,13 @@ export const registerJoinSwarmTool = (server: McpServer) => {
           content: [
             {
               type: "text",
-              text: `Successfully joined swarm as agent "${agent.name}" (ID: ${agent.id}).`,
+              text: `Successfully joined swarm as ${agent.isLead ? "Lead" : "Worker"} agent "${agent.name}" (ID: ${agent.id}).`,
             },
           ],
           structuredContent: {
-            yourAgentId: requestInfo.agentId,
+            yourAgentId: agent.id,
             success: true,
-            message: `Successfully joined swarm as agent "${agent.name}" (ID: ${agent.id}).`,
+            message: `Successfully joined swarm as ${agent.isLead ? "Lead" : "Worker"} agent "${agent.name}" (ID: ${agent.id}).`,
             agent,
           },
         };
@@ -91,7 +103,7 @@ export const registerJoinSwarmTool = (server: McpServer) => {
             },
           ],
           structuredContent: {
-            yourAgentId: requestInfo.agentId,
+            yourAgentId: requestInfo.agentId ?? requestedId,
             success: false,
             message: `Failed to join swarm: ${(error as Error).message}`,
           },

package/src/tools/poll-task.ts

@@ -1,7 +1,7 @@
 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 { getAgentById, getDb, getPendingTaskForAgent, startTask, updateAgentStatus } from "@/be/db";
 import { createToolRegistrar } from "@/tools/utils";
 import { AgentTaskSchema } from "@/types";
 
@@ -46,13 +46,45 @@ export const registerPollTaskTool = (server: McpServer) => {
       const now = new Date();
       const maxTime = addMinutes(now, MAX_POLL_DURATION_MS / 60000);
 
+      const agent = getAgentById(agentId);
+      if (!agent) {
+        return {
+          content: [
+            {
+              type: "text",
+              text: `Agent with ID "${agentId}" not found in the swarm.`,
+            },
+          ],
+          structuredContent: {
+            yourAgentId: requestInfo.agentId,
+            success: false,
+            message: `Agent with ID "${agentId}" not found in the swarm.`,
+            waitedForSeconds: 0,
+          },
+        };
+      }
+
       // Poll for pending tasks
       while (new Date() < maxTime) {
         // Fetch and update in a single transaction to avoid race conditions
         const startedTask = getDb().transaction(() => {
+          const agentNow = getAgentById(agentId)!;
+
+          if (agentNow.status !== "busy") {
+            updateAgentStatus(agentId, "idle");
+          }
+
           const pendingTask = getPendingTaskForAgent(agentId);
           if (!pendingTask) return null;
-          return startTask(pendingTask.id);
+
+          const maybeTask = startTask(pendingTask.id);
+
+          if (maybeTask) {
+            // Update automatically in case the agent forgets xd
+            updateAgentStatus(agentId, "busy");
+          }
+
+          return maybeTask;
         })();
 
         if (startedTask) {

package/src/tools/send-task.ts

@@ -1,6 +1,6 @@
 import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import * as z from "zod";
-import { createTask, getAgentById, getDb, updateAgentStatus } from "@/be/db";
+import { createTask, getAgentById, getDb } from "@/be/db";
 import { createToolRegistrar } from "@/tools/utils";
 import { AgentTaskSchema } from "@/types";
 
@@ -21,6 +21,38 @@ export const registerSendTaskTool = (server: McpServer) => {
       }),
     },
     async ({ agentId, task }, 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.',
+          },
+        };
+      }
+
+      if (agentId === requestInfo.agentId) {
+        return {
+          content: [
+            {
+              type: "text",
+              text: "Cannot send a task to yourself, are you drunk?",
+            },
+          ],
+          structuredContent: {
+            yourAgentId: requestInfo.agentId,
+            success: false,
+            message: "Cannot send a task to yourself, are you drunk?",
+          },
+        };
+      }
+
       const txn = getDb().transaction(() => {
         const agent = getAgentById(agentId);
 
@@ -31,6 +63,13 @@ export const registerSendTaskTool = (server: McpServer) => {
           };
         }
 
+        if (agent.isLead) {
+          return {
+            success: false,
+            message: `Cannot assign tasks to the lead agent "${agent.name}", wtf?`,
+          };
+        }
+
         if (agent.status !== "idle") {
           return {
             success: false,
@@ -39,7 +78,6 @@ export const registerSendTaskTool = (server: McpServer) => {
         }
 
         const newTask = createTask(agentId, task);
-        updateAgentStatus(agentId, "busy");
 
         return {
           success: true,

package/src/tools/store-progress.ts

@@ -3,6 +3,7 @@ import * as z from "zod";
 import {
   completeTask,
   failTask,
+  getAgentById,
   getDb,
   getTaskById,
   updateAgentStatus,
@@ -38,7 +39,32 @@ export const registerStoreProgressTool = (server: McpServer) => {
       }),
     },
     async ({ taskId, progress, status, output, failureReason }, 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 txn = getDb().transaction(() => {
+        const agent = getAgentById(requestInfo.agentId ?? "");
+
+        if (!agent) {
+          return {
+            success: false,
+            message: `Agent with ID "${requestInfo.agentId}" not found in the swarm, register before storing task progress.`,
+          };
+        }
+
         const existingTask = getTaskById(taskId);
 
         if (!existingTask) {
@@ -69,6 +95,9 @@ export const registerStoreProgressTool = (server: McpServer) => {
             updatedTask = result;
             updateAgentStatus(existingTask.agentId, "idle");
           }
+        } else {
+          // Keep it busy if just updating progress
+          updateAgentStatus(existingTask.agentId, "busy");
         }
 
         return {