@desplega.ai/agent-swarm 2.4.4 → 2.4.5

package/src/tests/mcp-input-caps.test.ts

@@ -0,0 +1,294 @@
+import { afterAll, beforeAll, beforeEach, describe, expect, test } from "bun:test";
+import { unlink } from "node:fs/promises";
+import {
+  createServer as createHttpServer,
+  request as httpRequest,
+  type IncomingMessage,
+  type Server,
+  type ServerResponse,
+} from "node:http";
+import type { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
+import { closeDb, createAgent, createUser, getDbClient, initDb } from "../be/db";
+import { type IdentityActor, mintToken } from "../be/users";
+import { handleCore } from "../http/core";
+import { handleMcp } from "../http/mcp";
+import { handleMcpUser } from "../http/mcp-user";
+import { MCP_MAX_REQUEST_BODY_BYTES, MCP_MAX_TOOL_INPUT_ELEMENTS } from "../mcp-limits";
+import { listenOnFreePort } from "./test-net";
+
+// GHSA-46jr-4wx8-v453: both MCP surfaces must bound tools/call argument size
+// and the raw HTTP body they read themselves.
+
+const TEST_DB_PATH = "./test-mcp-input-caps.sqlite";
+const API_KEY = "example-test-mcp-caps-key";
+const ACTOR: IdentityActor = { kind: "operator", id: "test" };
+const TOO_LARGE_MESSAGE = `Payload Too Large: Request body must not exceed ${MCP_MAX_REQUEST_BODY_BYTES} bytes`;
+
+async function removeDbFiles(path: string): Promise<void> {
+  for (const suffix of ["", "-wal", "-shm"]) {
+    try {
+      await unlink(path + suffix);
+    } catch (error) {
+      if ((error as NodeJS.ErrnoException).code !== "ENOENT") throw error;
+    }
+  }
+}
+
+let server: Server;
+let port: number;
+
+beforeAll(async () => {
+  await removeDbFiles(TEST_DB_PATH);
+  initDb(TEST_DB_PATH);
+  const transports: Record<string, StreamableHTTPServerTransport> = {};
+  const transportsUser: Record<string, StreamableHTTPServerTransport> = {};
+  const sessionAgents: Record<string, string> = {};
+  const sessionUsers: Record<string, string> = {};
+  server = createHttpServer(async (req: IncomingMessage, res: ServerResponse) => {
+    const myAgentId = req.headers["x-agent-id"] as string | undefined;
+    if (await handleCore(req, res, myAgentId, API_KEY)) return;
+    if (await handleMcp(req, res, transports, {}, sessionAgents)) return;
+    if (await handleMcpUser(req, res, transportsUser, sessionUsers)) return;
+    res.writeHead(404);
+    res.end("Not Found");
+  });
+  port = await listenOnFreePort(server, "127.0.0.1");
+});
+
+afterAll(async () => {
+  await new Promise<void>((resolve) => server.close(() => resolve()));
+  closeDb();
+  await removeDbFiles(TEST_DB_PATH);
+});
+
+beforeEach(async () => {
+  const client = getDbClient();
+  await client.run("DELETE FROM user_tokens");
+  await client.run("DELETE FROM agent_tasks");
+  await client.run("DELETE FROM users");
+  await client.run("DELETE FROM agents");
+});
+
+type Surface = {
+  name: string;
+  path: "/mcp" | "/mcp-user";
+  /** Credentials headers for a fresh caller on this surface. */
+  headers: () => Promise<Record<string, string>>;
+  /** A tools/call that should succeed with the given extra argument members. */
+  call: (extra: Record<string, unknown>) => { name: string; arguments: Record<string, unknown> };
+};
+
+const surfaces: Surface[] = [
+  {
+    name: "agent MCP",
+    path: "/mcp",
+    headers: async () => {
+      const agent = await createAgent({ name: "Caps Agent", isLead: false, status: "idle" });
+      return { Authorization: `Bearer ${API_KEY}`, "X-Agent-ID": agent.id };
+    },
+    call: (extra) => ({ name: "get-swarm", arguments: { ...extra } }),
+  },
+  {
+    name: "user MCP",
+    path: "/mcp-user",
+    headers: async () => {
+      await createAgent({ name: "Caps Lead", isLead: true, status: "idle" });
+      const user = await createUser({ name: "Caps User" });
+      const token = (await mintToken(user.id, "caps", ACTOR)).plaintext;
+      return { Authorization: `Bearer ${token}` };
+    },
+    call: (extra) => ({ name: "send-task", arguments: { task: "caps check", ...extra } }),
+  },
+];
+
+function parsePayload(text: string): unknown {
+  const trimmed = text.trim();
+  if (trimmed.startsWith("event:") || trimmed.startsWith("data:")) {
+    return JSON.parse(
+      trimmed
+        .split("\n")
+        .filter((line) => line.startsWith("data:"))
+        .map((line) => line.slice("data:".length).trim())
+        .join("\n"),
+    );
+  }
+  return JSON.parse(trimmed);
+}
+
+async function post(
+  path: string,
+  headers: Record<string, string>,
+  body: string,
+  sessionId?: string,
+): Promise<{ status: number; payload: unknown; sessionId: string | null }> {
+  const response = await fetch(`http://127.0.0.1:${port}${path}`, {
+    method: "POST",
+    headers: {
+      Accept: "application/json, text/event-stream",
+      "Content-Type": "application/json",
+      ...headers,
+      ...(sessionId ? { "mcp-session-id": sessionId } : {}),
+    },
+    body,
+  });
+  const text = await response.text();
+  return {
+    status: response.status,
+    payload: text ? parsePayload(text) : null,
+    sessionId: response.headers.get("mcp-session-id"),
+  };
+}
+
+async function openSession(path: string, headers: Record<string, string>): Promise<string> {
+  const init = await post(
+    path,
+    headers,
+    JSON.stringify({
+      jsonrpc: "2.0",
+      id: 1,
+      method: "initialize",
+      params: {
+        protocolVersion: "2024-11-05",
+        clientInfo: { name: "caps-test", version: "1" },
+        capabilities: {},
+      },
+    }),
+  );
+  expect(init.status).toBe(200);
+  if (!init.sessionId) throw new Error("initialize returned no mcp-session-id");
+  const ack = await post(
+    path,
+    headers,
+    JSON.stringify({ jsonrpc: "2.0", method: "notifications/initialized" }),
+    init.sessionId,
+  );
+  expect([200, 202]).toContain(ack.status);
+  return init.sessionId;
+}
+
+/** POST a body with no Content-Length (chunked), so only the streaming cap can stop it. */
+function postChunked(
+  path: string,
+  headers: Record<string, string>,
+  totalBytes: number,
+): Promise<{ status: number; body: string }> {
+  return new Promise((resolve, reject) => {
+    const req = httpRequest(
+      {
+        host: "127.0.0.1",
+        port,
+        path,
+        method: "POST",
+        headers: { "Content-Type": "application/json", ...headers },
+      },
+      (res) => {
+        let body = "";
+        res.setEncoding("utf8");
+        res.on("data", (chunk: string) => {
+          body += chunk;
+        });
+        res.on("end", () => resolve({ status: res.statusCode ?? 0, body }));
+      },
+    );
+    // The server may close the socket mid-upload after answering 413.
+    req.on("error", (error) => {
+      if (!req.writableFinished) return;
+      reject(error);
+    });
+    const chunk = Buffer.alloc(64 * 1024, 0x20);
+    let sent = 0;
+    const pump = () => {
+      while (sent < totalBytes) {
+        sent += chunk.byteLength;
+        if (!req.write(chunk)) {
+          req.once("drain", pump);
+          return;
+        }
+      }
+      req.end();
+    };
+    pump();
+  });
+}
+
+for (const surface of surfaces) {
+  describe(`${surface.name} input caps`, () => {
+    test("a normal tools/call passes", async () => {
+      const headers = await surface.headers();
+      const sessionId = await openSession(surface.path, headers);
+      const { status, payload } = await post(
+        surface.path,
+        headers,
+        JSON.stringify({
+          jsonrpc: "2.0",
+          id: 2,
+          method: "tools/call",
+          params: surface.call({}),
+        }),
+        sessionId,
+      );
+      expect(status).toBe(200);
+      const result = payload as { result?: { isError?: boolean } };
+      expect(result.result).toBeDefined();
+      expect(result.result?.isError).not.toBe(true);
+    });
+
+    test("a tools/call over the element limit is rejected before the handler runs", async () => {
+      const headers = await surface.headers();
+      const sessionId = await openSession(surface.path, headers);
+      const { status, payload } = await post(
+        surface.path,
+        headers,
+        JSON.stringify({
+          jsonrpc: "2.0",
+          id: 3,
+          method: "tools/call",
+          params: surface.call({ tags: new Array(MCP_MAX_TOOL_INPUT_ELEMENTS + 1).fill("t") }),
+        }),
+        sessionId,
+      );
+      expect(status).toBe(200);
+      const result = payload as {
+        result: { isError: boolean; content: Array<{ type: string; text: string }> };
+      };
+      expect(result.result.isError).toBe(true);
+      expect(result.result.content[0]?.text).toContain(
+        `more than the maximum of ${MCP_MAX_TOOL_INPUT_ELEMENTS} elements`,
+      );
+      const tasks = await getDbClient().query<{ id: string }>("SELECT id FROM agent_tasks");
+      expect(tasks).toHaveLength(0);
+    });
+
+    test("a body over the cap with Content-Length gets 413 and a JSON-RPC error", async () => {
+      const headers = await surface.headers();
+      const oversized = JSON.stringify({
+        jsonrpc: "2.0",
+        id: 4,
+        method: "tools/call",
+        params: surface.call({ blob: "x".repeat(MCP_MAX_REQUEST_BODY_BYTES) }),
+      });
+      const { status, payload } = await post(surface.path, headers, oversized);
+      expect(status).toBe(413);
+      expect(payload).toEqual({
+        jsonrpc: "2.0",
+        error: { code: -32000, message: TOO_LARGE_MESSAGE },
+        id: null,
+      });
+    });
+
+    test("a chunked body that crosses the cap gets 413", async () => {
+      const headers = await surface.headers();
+      const { status, body } = await postChunked(
+        surface.path,
+        headers,
+        MCP_MAX_REQUEST_BODY_BYTES + 128 * 1024,
+      );
+      expect(status).toBe(413);
+      expect(JSON.parse(body)).toEqual({
+        jsonrpc: "2.0",
+        error: { code: -32000, message: TOO_LARGE_MESSAGE },
+        id: null,
+      });
+    });
+  });
+}

package/src/tests/memory-chunked-search-starvation.test.ts

@@ -0,0 +1,200 @@
+import { afterAll, afterEach, beforeAll, describe, expect, mock, test } from "bun:test";
+import { randomUUID } from "node:crypto";
+import { unlink } from "node:fs/promises";
+import type { IncomingMessage, ServerResponse } from "node:http";
+import { Readable } from "node:stream";
+import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import { closeDb, createAgent, initDb } from "../be/db";
+import { taskCompletionMemoryKey } from "../be/memory/index-content";
+import { SqliteMemoryStore } from "../be/memory/providers/sqlite-store";
+import { countCandidateDocuments, searchMemoryCandidates } from "../be/memory/search-candidates";
+
+// Candidate starvation (PR #2302 review): rerank() collapses the chunks of one
+// task_completion document, but the store's candidate window (limit x 3) is
+// cut before the collapse. 32 high-similarity chunks of one task filled the
+// whole window, so limit 3 returned one result while five relevant manual
+// memories sat past the cutoff. Everything below the query embedding is real:
+// the SQLite store (vec and hybrid paths), the HTTP routes and the MCP tool.
+//
+// The 128-chunk document is the round-2 case: the hybrid and FTS arms capped
+// their fetch at 100 rows, so widening past 100 returned the same rows and
+// read as exhaustion while the manual memories sat behind chunk 100.
+
+/** Unit vector: `cos` on axis 2 (the query axis), the remainder on `axis`. */
+function mix(cos: number, axis: number): Float32Array {
+  const embedding = new Float32Array(512);
+  embedding[2] = cos;
+  embedding[axis] = Math.sqrt(1 - cos * cos);
+  return embedding;
+}
+
+const queryEmbedding = mix(1, 3);
+
+mock.module("../be/memory", () => ({
+  getEmbeddingProvider: () => ({
+    name: "test",
+    dimensions: 512,
+    embed: async () => queryEmbedding,
+    embedBatch: async (texts: string[]) => texts.map(() => queryEmbedding),
+    isConfigured: () => true,
+  }),
+  getMemoryStore: () => new SqliteMemoryStore(),
+}));
+
+const { handleMemory } = await import("../http/memory");
+const { registerMemorySearchTool } = await import("../tools/memory-search");
+
+const TEST_DB_PATH = "./test-memory-chunked-search-starvation.sqlite";
+const QUERY = "starvation probe";
+const MANUAL_DOCS = 5;
+const LIMIT = 3;
+let prevHybridFlag: string | undefined;
+
+/** One agent per chunk count: agent-scope rows keep the scenarios apart. */
+const scenarios = [32, 128].map((chunks) => ({
+  chunks,
+  agentId: randomUUID(),
+  manualIds: new Set<string>(),
+}));
+
+async function postJson(agentId: string, url: string, segments: string[]): Promise<string[]> {
+  const body = { query: QUERY, limit: LIMIT };
+  const req = Readable.from([Buffer.from(JSON.stringify(body))]) as IncomingMessage;
+  req.method = "POST";
+  req.url = url;
+  req.headers = {};
+  let payload: { results: { id: string }[] } = { results: [] };
+  const res = {
+    writeHead() {
+      return this;
+    },
+    end(chunk?: unknown) {
+      payload = typeof chunk === "string" ? JSON.parse(chunk) : (chunk as typeof payload);
+      return this;
+    },
+  } as unknown as ServerResponse;
+  expect(await handleMemory(req, res, segments, agentId)).toBe(true);
+  return payload.results.map((r) => r.id);
+}
+
+const searchHttp = (agentId: string) =>
+  postJson(agentId, "/api/memory/search", ["api", "memory", "search"]);
+const listHttp = (agentId: string) =>
+  postJson(agentId, "/api/memory/list", ["api", "memory", "list"]);
+
+async function searchTool(agentId: string): Promise<string[]> {
+  const server = new McpServer({ name: "chunked-search-starvation", version: "1.0.0" });
+  registerMemorySearchTool(server);
+  const tool = (
+    server as unknown as {
+      _registeredTools: Record<
+        string,
+        {
+          handler: (
+            args: unknown,
+            extra: unknown,
+          ) => Promise<{ structuredContent: { results: { id: string }[] } }>;
+        }
+      >;
+    }
+  )._registeredTools["memory-search"]!;
+  const result = await tool.handler(
+    { query: QUERY, limit: LIMIT, scope: "all" },
+    { sessionId: "starvation", requestInfo: { headers: { "x-agent-id": agentId } } },
+  );
+  return result.structuredContent.results.map((r) => r.id);
+}
+
+beforeAll(async () => {
+  prevHybridFlag = process.env.MEMORY_HYBRID_SEARCH;
+  for (const suffix of ["", "-wal", "-shm"]) {
+    try {
+      await unlink(TEST_DB_PATH + suffix);
+    } catch {}
+  }
+  initDb(TEST_DB_PATH);
+  const store = new SqliteMemoryStore();
+  for (const { chunks, agentId, manualIds } of scenarios) {
+    await createAgent({
+      id: agentId,
+      name: `Starvation Agent ${chunks}`,
+      isLead: false,
+      status: "idle",
+    });
+    const taskKey = taskCompletionMemoryKey(randomUUID());
+    // One long task output, every chunk closer to the query than any manual
+    // memory.
+    for (let i = 0; i < chunks; i++) {
+      const chunk = await store.store({
+        agentId,
+        scope: "agent",
+        name: "Task: long dreaming run",
+        content: `${QUERY} chunk ${i}: ${"lane starvation probe finding ".repeat(4)}`,
+        source: "task_completion",
+        key: taskKey,
+        chunkIndex: i,
+        totalChunks: chunks,
+      });
+      await store.updateEmbedding(chunk.id, mix(0.99 - i * 0.0001, 10 + i), "test");
+    }
+    for (let j = 0; j < MANUAL_DOCS; j++) {
+      const doc = await store.store({
+        agentId,
+        scope: "agent",
+        name: `manual note ${j}`,
+        content: `${QUERY} manual note ${j}`,
+        source: "manual",
+      });
+      manualIds.add(doc.id);
+      await store.updateEmbedding(doc.id, mix(0.8 - j * 0.01, 300 + j), "test");
+    }
+  }
+});
+
+afterEach(() => {
+  if (prevHybridFlag === undefined) delete process.env.MEMORY_HYBRID_SEARCH;
+  else process.env.MEMORY_HYBRID_SEARCH = prevHybridFlag;
+});
+
+afterAll(async () => {
+  closeDb();
+  for (const suffix of ["", "-wal", "-shm"]) {
+    try {
+      await unlink(TEST_DB_PATH + suffix);
+    } catch {}
+  }
+});
+
+describe("chunks of one task do not starve the result set", () => {
+  for (const { chunks, agentId, manualIds } of scenarios) {
+    for (const hybrid of ["0", "1"]) {
+      for (const [surface, search] of [
+        ["POST /api/memory/search", searchHttp],
+        ["POST /api/memory/list", listHttp],
+        ["memory-search tool", searchTool],
+      ] as const) {
+        test(`${chunks} chunks, ${surface}, hybrid=${hybrid}: limit ${LIMIT} returns ${LIMIT} documents`, async () => {
+          process.env.MEMORY_HYBRID_SEARCH = hybrid;
+          const ids = await search(agentId);
+          expect(ids).toHaveLength(LIMIT);
+          expect(new Set(ids).size).toBe(LIMIT);
+          // At most one slot for the task document; the rest are manual memories.
+          expect(ids.filter((id) => !manualIds.has(id)).length).toBeLessThanOrEqual(1);
+        });
+      }
+    }
+  }
+
+  test("FTS-only path (no query embedding) widens past 100 chunks of one task", async () => {
+    const { agentId, manualIds } = scenarios[1]!;
+    // A missing query embedding routes the store to searchFts.
+    const candidates = await searchMemoryCandidates(
+      new SqliteMemoryStore(),
+      new Float32Array(0),
+      agentId,
+      { limit: LIMIT, queryText: QUERY, scope: "all" },
+    );
+    expect(countCandidateDocuments(candidates)).toBeGreaterThanOrEqual(LIMIT);
+    expect(candidates.filter((c) => manualIds.has(c.id)).length).toBeGreaterThan(0);
+  });
+});

package/src/tests/memory-embedding-input-cap.test.ts

@@ -0,0 +1,114 @@
+import { afterAll, beforeAll, beforeEach, describe, expect, test } from "bun:test";
+import {
+  MAX_EMBEDDING_INPUT_CHARS,
+  OpenAIEmbeddingProvider,
+} from "../be/memory/providers/openai-embedding";
+
+// A gateway that rejects any input over REJECT_OVER chars the way the OpenAI
+// API rejects an input over 8,192 tokens: a 400 for the whole request.
+const REJECT_OVER = 10_000;
+let server: ReturnType<typeof Bun.serve>;
+const requests: string[][] = [];
+const previousBaseUrl = process.env.EMBEDDING_API_BASE_URL;
+
+beforeAll(() => {
+  server = Bun.serve({
+    hostname: "127.0.0.1",
+    port: 0,
+    async fetch(req) {
+      const body = (await req.json()) as { input: string | string[] };
+      const inputs = Array.isArray(body.input) ? body.input : [body.input];
+      requests.push(inputs);
+      const tooLong = inputs.findIndex((input) => input.length > REJECT_OVER);
+      if (tooLong !== -1) {
+        const param = Array.isArray(body.input) ? `input[${tooLong}]` : "input";
+        return Response.json(
+          {
+            error: {
+              message: `Invalid '${param}': maximum input length is 8192 tokens.`,
+              type: "invalid_request_error",
+              param: null,
+              code: null,
+            },
+          },
+          { status: 400 },
+        );
+      }
+      return Response.json({
+        object: "list",
+        data: inputs.map((input, index) => ({
+          object: "embedding",
+          index,
+          embedding: [input.length, 0, 0, 1],
+        })),
+        model: "test-embedding-model",
+        usage: { prompt_tokens: 1, total_tokens: 1 },
+      });
+    },
+  });
+  process.env.EMBEDDING_API_BASE_URL = `http://127.0.0.1:${server.port}/v1`;
+});
+
+afterAll(() => {
+  server.stop(true);
+  if (previousBaseUrl === undefined) delete process.env.EMBEDDING_API_BASE_URL;
+  else process.env.EMBEDDING_API_BASE_URL = previousBaseUrl;
+});
+
+beforeEach(() => {
+  requests.length = 0;
+});
+
+function provider(): OpenAIEmbeddingProvider {
+  return new OpenAIEmbeddingProvider({
+    apiKey: "test-key",
+    model: "test-embedding-model",
+    dimensions: 4,
+  });
+}
+
+/** The embedded length, which the fake gateway echoes as the first component. */
+function embeddedLength(vector: Float32Array | null): number | null {
+  return vector ? vector[0]! : null;
+}
+
+describe("embedding input cap", () => {
+  test("an over-cap input is cut down until the provider accepts it", async () => {
+    const vector = await provider().embed("x".repeat(50_000));
+
+    expect(embeddedLength(vector)).toBe(MAX_EMBEDDING_INPUT_CHARS / 4);
+    expect(requests.map((r) => r[0]!.length)).toEqual([
+      MAX_EMBEDDING_INPUT_CHARS,
+      MAX_EMBEDDING_INPUT_CHARS / 2,
+      MAX_EMBEDDING_INPUT_CHARS / 4,
+    ]);
+  });
+
+  test("no input is ever sent over the char cap", async () => {
+    await provider().embedBatch(["a".repeat(100_000), "b".repeat(30_000)]);
+    for (const request of requests) {
+      for (const input of request)
+        expect(input.length).toBeLessThanOrEqual(MAX_EMBEDDING_INPUT_CHARS);
+    }
+  });
+
+  test("one oversized input in a batch does not fail the other inputs", async () => {
+    const vectors = await provider().embedBatch(["small one", "y".repeat(40_000), "small two"]);
+
+    expect(vectors.map(embeddedLength)).toEqual([
+      "small one".length,
+      MAX_EMBEDDING_INPUT_CHARS / 4,
+      "small two".length,
+    ]);
+    // The whole batch first, then each input alone after the rejection.
+    expect(requests[0]).toHaveLength(3);
+    expect(requests.slice(1).every((r) => r.length === 1)).toBe(true);
+  });
+
+  test("a batch with no oversized input goes out as one request", async () => {
+    const vectors = await provider().embedBatch(["one", "", "three"]);
+
+    expect(vectors.map(embeddedLength)).toEqual([3, null, 5]);
+    expect(requests).toEqual([["one", "three"]]);
+  });
+});

package/src/tests/memory.test.ts

@@ -214,6 +214,20 @@ describe("Memory System", () => {
       expect(xChars).toBeGreaterThanOrEqual(2500);
     });
 
+    test("keeps sections under MIN_CHUNK_SIZE, packed into shared chunks", () => {
+      const preamble = "Report preamble with enough words to stand as its own chunk here.";
+      const sections = Array.from({ length: 100 }, (_, i) => `## Lane ${i}\n\nfact-${i} ok`);
+      const chunks = chunkContent([preamble, ...sections].join("\n\n"));
+
+      const stored = chunks.map((c) => c.content).join("\n");
+      for (let i = 0; i < 100; i++) {
+        expect(stored).toContain(`## Lane ${i}\n\nfact-${i} ok`);
+      }
+      expect(chunks[0]!.content).toBe(preamble);
+      expect(chunks.length).toBeLessThan(6);
+      for (const chunk of chunks) expect(chunk.content.length).toBeLessThanOrEqual(2000);
+    });
+
     test("chunk indices are sequential", () => {
       const longText = "A ".repeat(3000);
       const chunks = chunkContent(longText);