@desplega.ai/agent-swarm 2.7.5 → 2.7.6

package/src/tests/list-total-cache.test.ts

@@ -1,5 +1,5 @@
 import { describe, expect, test } from "bun:test";
-import { createListTotalCache } from "../http/list-total-cache";
+import { createListTotalCache, LIST_TOTAL_TTL_MS } from "../http/list-total-cache";
 
 function counter(values: number[]) {
   let calls = 0;
@@ -10,7 +10,7 @@ function counter(values: number[]) {
 }
 
 describe("list total cache", () => {
-  test("page 1 always counts; later pages reuse its count", async () => {
+  test("page 1 always counts by default; later pages reuse its count", async () => {
     const cache = createListTotalCache();
     const count = counter([10, 11, 12]);
 
@@ -24,6 +24,31 @@ describe("list total cache", () => {
     expect(count.calls()).toBe(2);
   });
 
+  test("page 1 polls within the first-page TTL share one count", async () => {
+    const cache = createListTotalCache(LIST_TOTAL_TTL_MS, 50);
+    const count = counter([10, 11]);
+    expect(await cache.get({ q: "a" }, 0, count.compute)).toBe(10);
+    expect(await cache.get({ q: "a" }, 0, count.compute)).toBe(10);
+    expect(count.calls()).toBe(1);
+    await Bun.sleep(60);
+    expect(await cache.get({ q: "a" }, 0, count.compute)).toBe(11);
+    expect(count.calls()).toBe(2);
+  });
+
+  test("with a first-page TTL, concurrent page 1 requests share an in-flight count", async () => {
+    const cache = createListTotalCache(LIST_TOTAL_TTL_MS, 1000);
+    let calls = 0;
+    const compute = async () => {
+      calls++;
+      await Bun.sleep(5);
+      return 4;
+    };
+    expect(await Promise.all([cache.get({}, 0, compute), cache.get({}, 0, compute)])).toEqual([
+      4, 4,
+    ]);
+    expect(calls).toBe(1);
+  });
+
   test("a cold later page counts, so total is always set", async () => {
     const cache = createListTotalCache();
     const count = counter([7]);

package/src/tests/memory-fts-populate.test.ts

@@ -2,8 +2,8 @@ import { afterAll, beforeAll, describe, expect, test } from "bun:test";
 import { unlink } from "node:fs/promises";
 import { closeDb, createAgent, getDb, getDbClient, initDb } from "../be/db";
 import {
-  FTS_DELETE_EXTRA_SQL,
-  FTS_MISSING_IDS_SQL,
+  FTS_SCAN_BATCH_SQL,
+  MEMORY_ID_BATCH_SQL,
   SqliteMemoryStore,
 } from "../be/memory/providers/sqlite-store";
 
@@ -62,16 +62,18 @@ describe("memory_fts boot backfill", () => {
     // Positive control: the old anti-join scans memory_fts once per memory.
     expect(planOf(QUADRATIC_BACKFILL_SQL).join(" | ")).toContain("CORRELATED");
 
-    for (const sql of [FTS_MISSING_IDS_SQL, FTS_DELETE_EXTRA_SQL]) {
-      const plan = planOf(sql).join(" | ");
-      expect(plan).not.toContain("CORRELATED");
-      // Exactly one pass over memory_fts per statement.
-      expect(plan.match(/SCAN memory_fts/g)?.length).toBe(1);
-    }
+    // Each batch seeks a rowid range of one table. The orphan check is a
+    // lookup on the agent_memory primary key, never a scan of memory_fts.
+    const scan = planOf(FTS_SCAN_BATCH_SQL.replace(/\?/g, "0")).join(" | ");
+    expect(scan).toMatch(/SCAN memory_fts VIRTUAL TABLE INDEX \d+:>/);
+    expect(scan.match(/memory_fts/g)?.length).toBe(1);
+    expect(scan).toContain("SEARCH m USING COVERING INDEX sqlite_autoindex_agent_memory_1 (id=?)");
+    const ids = planOf(MEMORY_ID_BATCH_SQL.replace(/\?/g, "0")).join(" | ");
+    expect(ids).toContain("SEARCH agent_memory USING INTEGER PRIMARY KEY (rowid>?)");
   });
 
   test("a new store backfills missing rows across batches and drops orphans", async () => {
-    const total = 1050; // crosses two 500-row batch boundaries
+    const total = 1050; // crosses two 500-row insert batches and one 1000-row scan batch
     const seed = await store.store({
       agentId,
       scope: "agent",

package/src/tests/memory-retrieval-purge.test.ts

@@ -0,0 +1,80 @@
+import { afterAll, beforeAll, describe, expect, spyOn, test } from "bun:test";
+import { unlink } from "node:fs/promises";
+import { closeDb, getDb, getDbClient, initDb } from "../be/db";
+import { purgeStaleSearchRetrievals, startMemoryGc, stopMemoryGc } from "../http/memory";
+
+const TEST_DB_PATH = "./test-memory-retrieval-purge.sqlite";
+const DAY = 24 * 60 * 60 * 1000;
+
+async function removeDbFiles(): Promise<void> {
+  for (const suffix of ["", "-wal", "-shm"])
+    await unlink(`${TEST_DB_PATH}${suffix}`).catch(() => undefined);
+}
+
+async function insertRetrieval(id: string, eventType: string, ageDays: number): Promise<void> {
+  await getDbClient().run(
+    `INSERT INTO memory_retrieval (id, agentId, memoryId, retrievedAt, eventType)
+     VALUES (?, 'agent-1', 'memory-1', ?, ?)`,
+    [id, new Date(Date.now() - ageDays * DAY).toISOString(), eventType],
+  );
+}
+
+describe("purgeStaleSearchRetrievals", () => {
+  beforeAll(async () => {
+    await removeDbFiles();
+    initDb(TEST_DB_PATH);
+  });
+
+  afterAll(async () => {
+    closeDb();
+    await removeDbFiles();
+  });
+
+  test("deletes stale search rows in batches and keeps fresh and get rows", async () => {
+    for (let i = 0; i < 5; i++) await insertRetrieval(`stale-search-${i}`, "search", 120);
+    await insertRetrieval("fresh-search", "search", 10);
+    await insertRetrieval("stale-get", "get", 120);
+
+    const run = spyOn(getDbClient(), "run");
+    let purged: number;
+    let calls: (typeof run.mock.calls)[number][];
+    try {
+      purged = await purgeStaleSearchRetrievals(2);
+      calls = [...run.mock.calls];
+    } finally {
+      run.mockRestore();
+    }
+    expect(purged).toBe(5);
+    // Batches of 2: 2 + 2 + 1, and the short batch ends the loop.
+    expect(calls).toHaveLength(3);
+
+    const left = await getDbClient().query<{ id: string }>(
+      "SELECT id FROM memory_retrieval ORDER BY id",
+    );
+    expect(left.map((row) => row.id)).toEqual(["fresh-search", "stale-get"]);
+
+    // Without sqlite_stat1 the plan must stay on the retrievedAt range, not
+    // read every search row through the eventType index.
+    const [sql, params] = calls[0]!;
+    const plan = getDb()
+      .query<{ detail: string }, any[]>(`EXPLAIN QUERY PLAN ${sql}`)
+      .all(...((params ?? []) as any[]))
+      .map((row) => row.detail)
+      .join("\n");
+    expect(plan).toContain("USING INDEX idx_memret_retrieved_at (retrievedAt<?)");
+    expect(plan).not.toContain("idx_memret_event_type");
+  });
+});
+
+describe("startMemoryGc", () => {
+  test("defers the first purge off the boot path", async () => {
+    const timeout = spyOn(globalThis, "setTimeout");
+    try {
+      await startMemoryGc(60 * 60 * 1000, 120_000);
+      expect(timeout.mock.calls.some(([, ms]) => ms === 120_000)).toBe(true);
+    } finally {
+      stopMemoryGc();
+      timeout.mockRestore();
+    }
+  });
+});

package/src/tests/migration-260-runtime-pool-cache.test.ts

@@ -0,0 +1,118 @@
+import { Database } from "bun:sqlite";
+import { expect, test } from "bun:test";
+import { runMigrations } from "../be/migrations/runner";
+import {
+  ControllerClientMessageSchema,
+  ControllerCommandSchema,
+} from "../runtime-providers/controller-protocol";
+
+test("migration 260 adds an optional JSON cache and the release table to a fresh database", () => {
+  const db = new Database(":memory:");
+  try {
+    runMigrations(db);
+    expect(db.query("PRAGMA table_info(runtime_pools)").all()).toContainEqual(
+      expect.objectContaining({ name: "cache", type: "TEXT", notnull: 0 }),
+    );
+    expect(db.query("PRAGMA foreign_key_check").all()).toEqual([]);
+    for (const name of ["runtime_ids", "state", "created_by", "updated_by"])
+      expect(db.query("PRAGMA table_info(runtime_cache_releases)").all()).toContainEqual(
+        expect.objectContaining({ name, type: "TEXT" }),
+      );
+  } finally {
+    db.close();
+  }
+});
+
+test("migration 260 preserves existing pools with caching disabled", async () => {
+  const db = new Database(":memory:");
+  try {
+    db.exec(
+      "CREATE TABLE _migrations (version INTEGER PRIMARY KEY, name TEXT NOT NULL, applied_at TEXT NOT NULL, checksum TEXT NOT NULL)",
+    );
+    for (const file of [...new Bun.Glob("*.sql").scanSync("src/be/migrations")]
+      .sort()
+      .filter((file) => Number(file.split("_")[0]) <= 259)) {
+      const sql = await Bun.file(`src/be/migrations/${file}`).text();
+      db.transaction(() => {
+        db.exec(sql);
+        db.run("INSERT INTO _migrations VALUES (?, ?, ?, ?)", [
+          Number(file.split("_")[0]),
+          file.replace(".sql", ""),
+          new Date().toISOString(),
+          new Bun.CryptoHasher("sha256").update(sql).digest("hex"),
+        ]);
+      }).immediate();
+    }
+    db.run(
+      "INSERT INTO runtime_pools (id, name, kind, placements, image, resources, env, lifecycle) VALUES (?, ?, ?, ?, ?, ?, ?, 'idle')",
+      [
+        "existing-pool",
+        "existing",
+        "docker",
+        '["controller"]',
+        "worker:test",
+        '{"cpu":2}',
+        '{"CUSTOM":"kept"}',
+      ],
+    );
+    runMigrations(db);
+    expect(db.query("SELECT name, image, resources, env, cache FROM runtime_pools").get()).toEqual({
+      name: "existing",
+      image: "worker:test",
+      resources: '{"cpu":2}',
+      env: '{"CUSTOM":"kept"}',
+      cache: null,
+    });
+    expect(() => db.run("UPDATE runtime_pools SET cache = 'invalid json'")).toThrow();
+    expect(db.query("SELECT COUNT(*) AS n FROM runtime_cache_releases").get()).toEqual({ n: 0 });
+  } finally {
+    db.close();
+  }
+});
+
+test("cache release runtime IDs accept arrays and reject other JSON values", () => {
+  const db = new Database(":memory:");
+  try {
+    runMigrations(db);
+    db.run(
+      "INSERT INTO runtime_controllers (id, name, kind, placement) VALUES ('controller', 'controller', 'docker', 'builtin')",
+    );
+    db.run(`INSERT INTO runtime_cache_releases (id, pool_id, controller_id, owner_id, cache_key, runtime_ids)
+      VALUES ('release', 'pool', 'controller', 'owner', 'cache', '[]')`);
+    for (const invalid of ["null", "{}", "42", "true", '"runtime"', "invalid json"])
+      expect(() =>
+        db.run("UPDATE runtime_cache_releases SET runtime_ids = ?", [invalid]),
+      ).toThrow();
+    db.run("UPDATE runtime_cache_releases SET runtime_ids = ?", ['["runtime-a","runtime-b"]']);
+    expect(db.query("SELECT runtime_ids FROM runtime_cache_releases").get()).toEqual({
+      runtime_ids: '["runtime-a","runtime-b"]',
+    });
+  } finally {
+    db.close();
+  }
+});
+
+test("controller protocol accepts legacy hello messages and retains optional cache metadata", () => {
+  const legacy = ControllerClientMessageSchema.parse({
+    type: "hello",
+    controllerId: "controller",
+    kind: "docker",
+    version: 1,
+    traits: { canProvision: true, ephemeralDisk: true, supportsPause: false },
+  });
+  if (legacy.type !== "hello") throw new Error("Expected hello");
+  expect(legacy.traits.persistentCache).toBeUndefined();
+  const payload = {
+    runtimeId: "runtime",
+    agentId: "agent",
+    poolId: "pool",
+    ownerId: "owner",
+    image: "worker:test",
+    env: {},
+    labels: {},
+    cache: { key: "swarm-cache-pool", mountPath: "/cache" },
+  };
+  expect(
+    ControllerCommandSchema.parse({ type: "cmd", id: "command", op: "provision", payload }),
+  ).toMatchObject({ payload: { cache: payload.cache } });
+});

package/src/tests/query-audit-plans.test.ts

@@ -156,6 +156,22 @@ const CASES: Array<{
     uses: ["COVERING INDEX idx_session_summaries_source (source=?)"],
     revert: { dropIndex: "idx_session_summaries_source" },
   },
+  {
+    // A user's task count hides other users' sessions. Both steps of that
+    // walk stay in an index (migration 262) instead of reading task rows.
+    name: "getTasksCount readableBy, root step",
+    match: /COUNT/,
+    run: () => getTasksCount({ status: "in_progress", readableBy: { userId: "viewer-1" } }),
+    uses: ["SEARCH r USING COVERING INDEX idx_agent_tasks_session_owner (scheduleId=?)"],
+    revert: { dropIndex: "idx_agent_tasks_session_owner" },
+  },
+  {
+    name: "getTasksCount readableBy, recursive step",
+    match: /COUNT/,
+    run: () => getTasksCount({ status: "in_progress", readableBy: { userId: "viewer-1" } }),
+    uses: ["SEARCH c USING COVERING INDEX idx_agent_tasks_parentTaskId_id (parentTaskId=?)"],
+    revert: { dropIndex: "idx_agent_tasks_parentTaskId_id" },
+  },
   {
     // Few hits: the page reads the agent_tasks_fts candidates by rowid. The
     // control is the old `task LIKE '%q%'`, which scans every task body.
@@ -221,6 +237,17 @@ describe("query audit plans without sqlite_stat1", () => {
     });
   }
 
+  test("the default task count scans the filter index, not the task rows", async () => {
+    // The heartbeat filter is a NOT IN and a NOT LIKE, so the count scans.
+    // Migration 263 lets it scan the narrow index instead of every task body.
+    const count = await capturedSql(/COUNT/, () =>
+      getTasksCount({ readableBy: { userId: "viewer-1" } }),
+    );
+    const uses = ["SCAN agent_tasks USING COVERING INDEX idx_agent_tasks_list_filter"];
+    expect(missing(plan(count), uses)).toEqual([]);
+    expect(missing(planWithoutIndex(count, "idx_agent_tasks_list_filter"), uses)).not.toEqual([]);
+  });
+
   test("task search with many hits pages by lastUpdatedAt and counts from the index", async () => {
     await getDbClient().transaction(async (tx) => {
       for (let i = 0; i < 501; i++) {