@okrlinkhub/agent-factory 1.0.0 → 1.0.2

package/src/component/lib.test.ts

@@ -2,6 +2,7 @@
 
 import { afterEach, beforeEach, describe, expect, test, vi } from "vitest";
 import { api, internal } from "./_generated/api.js";
+import { DEFAULT_CONFIG } from "./config.js";
 import { initConvexTest } from "./setup.test.js";
 import { canTransitionWorkerStatus } from "./workerLifecycle.js";
 
@@ -1090,6 +1091,499 @@ describe("component lib", () => {
     expect(reconcile.activeWorkers).toBe(2);
   });
 
+  test("worker assignment should prevent cross-conversation claims after completion", async () => {
+    const t = initConvexTest();
+    const nowMs = Date.UTC(2026, 0, 1, 17, 30, 0);
+    vi.setSystemTime(nowMs);
+    await t.mutation(api.queue.upsertAgentProfile, {
+      agentKey: "support-agent",
+      version: "1.0.0",
+      secretsRef: [],
+      enabled: true,
+    });
+
+    const conversationA = "telegram:chat:sticky-a";
+    const conversationB = "telegram:chat:sticky-b";
+    const messageA = await t.mutation(api.queue.enqueueMessage, {
+      conversationId: conversationA,
+      agentKey: "support-agent",
+      payload: {
+        provider: "telegram",
+        providerUserId: "u-sticky-a",
+        messageText: "first",
+      },
+      nowMs,
+    });
+
+    const firstClaim = await t.mutation(api.lib.claim, {
+      workerId: "worker-sticky-1",
+      conversationId: conversationA,
+      nowMs,
+    });
+    expect(firstClaim?.messageId).toBe(messageA);
+
+    const completed = await t.mutation(api.lib.complete, {
+      workerId: "worker-sticky-1",
+      messageId: messageA,
+      leaseId: firstClaim?.leaseId ?? "",
+      nowMs: nowMs + 1_000,
+    });
+    expect(completed).toBe(true);
+
+    await t.mutation(api.queue.enqueueMessage, {
+      conversationId: conversationB,
+      agentKey: "support-agent",
+      payload: {
+        provider: "telegram",
+        providerUserId: "u-sticky-b",
+        messageText: "second",
+      },
+      nowMs: nowMs + 2_000,
+    });
+
+    const mismatchedClaim = await t.mutation(api.lib.claim, {
+      workerId: "worker-sticky-1",
+      nowMs: nowMs + 2_000,
+    });
+    expect(mismatchedClaim).toBeNull();
+
+    const queuedJobs = await t.query(api.queue.listJobsByStatus, {
+      status: "queued",
+      limit: 10,
+    });
+    expect(queuedJobs.some((job) => job.conversationId === conversationB)).toBe(true);
+
+    const workers = await t.query((internal.queue as any).listWorkersForScheduler, {});
+    const worker = workers.find((row: { workerId: string }) => row.workerId === "worker-sticky-1");
+    expect(worker?.assignment).toEqual({
+      conversationId: conversationA,
+      agentKey: "support-agent",
+      leaseId: firstClaim?.leaseId ?? "",
+      assignedAt: nowMs,
+    });
+  });
+
+  test("exclusive ownership should block another worker and let the owner reclaim", async () => {
+    const t = initConvexTest();
+    const nowMs = Date.UTC(2026, 0, 1, 17, 40, 0);
+    vi.setSystemTime(nowMs);
+    await t.mutation(api.queue.upsertAgentProfile, {
+      agentKey: "support-agent",
+      version: "1.0.0",
+      secretsRef: [],
+      enabled: true,
+    });
+
+    const conversationId = "telegram:chat:exclusive-owner";
+    const firstMessageId = await t.mutation(api.queue.enqueueMessage, {
+      conversationId,
+      agentKey: "support-agent",
+      payload: {
+        provider: "telegram",
+        providerUserId: "u-exclusive-1",
+        messageText: "first",
+      },
+      nowMs,
+    });
+    const firstClaim = await t.mutation(api.lib.claim, {
+      workerId: "worker-exclusive-1",
+      conversationId,
+      nowMs,
+    });
+    expect(firstClaim?.messageId).toBe(firstMessageId);
+
+    await t.mutation(api.lib.complete, {
+      workerId: "worker-exclusive-1",
+      messageId: firstMessageId,
+      leaseId: firstClaim?.leaseId ?? "",
+      nowMs: nowMs + 1_000,
+    });
+
+    const secondMessageId = await t.mutation(api.queue.enqueueMessage, {
+      conversationId,
+      agentKey: "support-agent",
+      payload: {
+        provider: "telegram",
+        providerUserId: "u-exclusive-2",
+        messageText: "second",
+      },
+      nowMs: nowMs + 2_000,
+    });
+
+    const blockedClaim = await t.mutation(api.lib.claim, {
+      workerId: "worker-exclusive-2",
+      conversationId,
+      nowMs: nowMs + 2_000,
+    });
+    expect(blockedClaim).toBeNull();
+
+    const ownerReclaim = await t.mutation(api.lib.claim, {
+      workerId: "worker-exclusive-1",
+      nowMs: nowMs + 2_000,
+    });
+    expect(ownerReclaim?.messageId).toBe(secondMessageId);
+    expect(ownerReclaim?.conversationId).toBe(conversationId);
+  });
+
+  test("scheduler should spawn when only an idle worker pinned to another conversation exists", async () => {
+    const t = initConvexTest();
+    const nowMs = Date.UTC(2026, 0, 1, 17, 45, 0);
+    vi.setSystemTime(nowMs);
+    const fetchMock = vi.fn(async (input: RequestInfo | URL, init?: RequestInit) => {
+      const url = String(input);
+      const method = init?.method ?? "GET";
+      if (url.endsWith(`/apps/${TEST_PROVIDER_CONFIG.appName}/machines`) && method === "GET") {
+        return jsonResponse([]);
+      }
+      if (url.endsWith(`/apps/${TEST_PROVIDER_CONFIG.appName}/volumes`) && method === "GET") {
+        return jsonResponse([]);
+      }
+      if (url.endsWith(`/apps/${TEST_PROVIDER_CONFIG.appName}/volumes`) && method === "POST") {
+        return jsonResponse({
+          id: "vol-pinned-new-worker",
+          name: buildDedicatedVolumeName(TEST_PROVIDER_CONFIG.volumeName, "afw-177"),
+          region: TEST_PROVIDER_CONFIG.region,
+        });
+      }
+      if (url.endsWith(`/apps/${TEST_PROVIDER_CONFIG.appName}/machines`) && method === "POST") {
+        const body = JSON.parse(String(init?.body ?? "{}")) as { name?: string };
+        return jsonResponse({
+          id: "machine-pinned-new-worker",
+          name: body.name,
+          region: TEST_PROVIDER_CONFIG.region,
+          state: "started",
+          config: { image: TEST_PROVIDER_CONFIG.image },
+        });
+      }
+      throw new Error(`Unexpected fetch ${method} ${url}`);
+    });
+    vi.stubGlobal("fetch", fetchMock);
+
+    await t.mutation(api.queue.upsertAgentProfile, {
+      agentKey: "support-agent",
+      version: "1.0.0",
+      secretsRef: [],
+      enabled: true,
+    });
+    await t.mutation(api.queue.importPlaintextSecret, {
+      secretRef: "fly.apiToken",
+      plaintextValue: "fly-token",
+    });
+    await t.mutation(api.queue.importPlaintextSecret, {
+      secretRef: "convex.url",
+      plaintextValue: "https://example.convex.cloud",
+    });
+
+    const messageA = await t.mutation(api.queue.enqueueMessage, {
+      conversationId: "telegram:chat:pinned-a",
+      agentKey: "support-agent",
+      payload: {
+        provider: "telegram",
+        providerUserId: "u-pinned-a",
+        messageText: "first",
+      },
+      nowMs,
+    });
+    const claimA = await t.mutation(api.lib.claim, {
+      workerId: "worker-pinned-1",
+      conversationId: "telegram:chat:pinned-a",
+      nowMs,
+    });
+    expect(claimA?.messageId).toBe(messageA);
+
+    await t.mutation(api.lib.complete, {
+      workerId: "worker-pinned-1",
+      messageId: messageA,
+      leaseId: claimA?.leaseId ?? "",
+      nowMs: nowMs + 1_000,
+    });
+
+    await t.mutation(api.queue.enqueueMessage, {
+      conversationId: "telegram:chat:pinned-b",
+      agentKey: "support-agent",
+      payload: {
+        provider: "telegram",
+        providerUserId: "u-pinned-b",
+        messageText: "second",
+      },
+      nowMs: nowMs + 2_000,
+    });
+
+    const reconcile = await t.action(api.scheduler.reconcileWorkerPool, {
+      nowMs: nowMs + 2_000,
+      flyApiToken: "fly-token",
+      convexUrl: "https://example.convex.cloud",
+      scalingPolicy: {
+        maxWorkers: 5,
+        queuePerWorkerTarget: 1,
+        spawnStep: 1,
+        idleTimeoutMs: 300_000,
+        reconcileIntervalMs: 15_000,
+      },
+      providerConfig: TEST_PROVIDER_CONFIG,
+    });
+
+    expect(reconcile.spawned).toBe(1);
+    expect(reconcile.activeWorkers).toBe(2);
+  });
+
+  test("stale owner should allow another worker to take over the conversation", async () => {
+    const t = initConvexTest();
+    const nowMs = Date.UTC(2026, 0, 1, 17, 47, 0);
+    vi.setSystemTime(nowMs);
+    await t.mutation(api.queue.upsertAgentProfile, {
+      agentKey: "support-agent",
+      version: "1.0.0",
+      secretsRef: [],
+      enabled: true,
+    });
+
+    const conversationId = "telegram:chat:stale-owner";
+    const firstMessageId = await t.mutation(api.queue.enqueueMessage, {
+      conversationId,
+      agentKey: "support-agent",
+      payload: {
+        provider: "telegram",
+        providerUserId: "u-stale-1",
+        messageText: "first",
+      },
+      nowMs,
+    });
+    const firstClaim = await t.mutation(api.lib.claim, {
+      workerId: "worker-stale-1",
+      conversationId,
+      nowMs,
+    });
+    expect(firstClaim?.messageId).toBe(firstMessageId);
+
+    await t.mutation(api.lib.complete, {
+      workerId: "worker-stale-1",
+      messageId: firstMessageId,
+      leaseId: firstClaim?.leaseId ?? "",
+      nowMs: nowMs + 1_000,
+    });
+
+    await t.run(async (ctx) => {
+      const worker = await ctx.db
+        .query("workers")
+        .withIndex("by_workerId", (q) => q.eq("workerId", "worker-stale-1"))
+        .unique();
+      expect(worker).not.toBeNull();
+      await ctx.db.patch(worker!._id, {
+        heartbeatAt: nowMs + 2_000 - DEFAULT_CONFIG.lease.staleAfterMs - 1,
+      });
+    });
+
+    const secondMessageId = await t.mutation(api.queue.enqueueMessage, {
+      conversationId,
+      agentKey: "support-agent",
+      payload: {
+        provider: "telegram",
+        providerUserId: "u-stale-2",
+        messageText: "second",
+      },
+      nowMs: nowMs + 2_000,
+    });
+
+    const takeoverClaim = await t.mutation(api.lib.claim, {
+      workerId: "worker-stale-2",
+      conversationId,
+      nowMs: nowMs + 2_000,
+    });
+    expect(takeoverClaim?.messageId).toBe(secondMessageId);
+    expect(takeoverClaim?.conversationId).toBe(conversationId);
+  });
+
+  test("scheduler should dedupe duplicated sticky workers for the same conversation", async () => {
+    const t = initConvexTest();
+    const nowMs = Date.UTC(2026, 0, 1, 17, 48, 0);
+    vi.setSystemTime(nowMs);
+    const fetchMock = vi.fn(async (input: RequestInfo | URL, init?: RequestInit) => {
+      const url = String(input);
+      const method = init?.method ?? "GET";
+      if (url.endsWith(`/apps/${TEST_PROVIDER_CONFIG.appName}/machines`) && method === "GET") {
+        return jsonResponse([]);
+      }
+      if (url.endsWith(`/apps/${TEST_PROVIDER_CONFIG.appName}/volumes`) && method === "GET") {
+        return jsonResponse([]);
+      }
+      if (url.endsWith(`/apps/${TEST_PROVIDER_CONFIG.appName}/volumes`) && method === "POST") {
+        return jsonResponse({
+          id: "vol-deduped-sticky-worker",
+          name: buildDedicatedVolumeName(TEST_PROVIDER_CONFIG.volumeName, "afw-deduped"),
+          region: TEST_PROVIDER_CONFIG.region,
+        });
+      }
+      if (url.endsWith(`/apps/${TEST_PROVIDER_CONFIG.appName}/machines`) && method === "POST") {
+        const body = JSON.parse(String(init?.body ?? "{}")) as { name?: string };
+        return jsonResponse({
+          id: "machine-deduped-sticky-worker",
+          name: body.name,
+          region: TEST_PROVIDER_CONFIG.region,
+          state: "started",
+          config: { image: TEST_PROVIDER_CONFIG.image },
+        });
+      }
+      throw new Error(`Unexpected fetch ${method} ${url}`);
+    });
+    vi.stubGlobal("fetch", fetchMock);
+
+    await t.mutation(api.queue.upsertAgentProfile, {
+      agentKey: "support-agent",
+      version: "1.0.0",
+      secretsRef: [],
+      enabled: true,
+    });
+    await t.mutation(api.queue.importPlaintextSecret, {
+      secretRef: "fly.apiToken",
+      plaintextValue: "fly-token",
+    });
+    await t.mutation(api.queue.importPlaintextSecret, {
+      secretRef: "convex.url",
+      plaintextValue: "https://example.convex.cloud",
+    });
+    await t.mutation(api.queue.enqueueMessage, {
+      conversationId: "telegram:chat:dedupe-a",
+      agentKey: "support-agent",
+      payload: {
+        provider: "telegram",
+        providerUserId: "u-dedupe-a",
+        messageText: "first",
+      },
+      nowMs,
+    });
+    await t.mutation(api.queue.enqueueMessage, {
+      conversationId: "telegram:chat:dedupe-b",
+      agentKey: "support-agent",
+      payload: {
+        provider: "telegram",
+        providerUserId: "u-dedupe-b",
+        messageText: "second",
+      },
+      nowMs,
+    });
+
+    await t.run(async (ctx) => {
+      for (const workerId of ["worker-dedupe-1", "worker-dedupe-2", "worker-dedupe-3"]) {
+        await ctx.db.insert("workers", {
+          workerId,
+          provider: "fly",
+          status: "active",
+          load: 0,
+          heartbeatAt: nowMs,
+          lastClaimAt: nowMs,
+          scheduledShutdownAt: nowMs + 300_000,
+          assignment: {
+            conversationId: "telegram:chat:dedupe-a",
+            agentKey: "support-agent",
+            leaseId: `${workerId}-lease`,
+            assignedAt: nowMs,
+          },
+          machineRef: {
+            appName: TEST_PROVIDER_CONFIG.appName,
+            machineId: `${workerId}-machine`,
+            region: TEST_PROVIDER_CONFIG.region,
+          },
+          capabilities: [],
+        });
+      }
+    });
+
+    const reconcile = await t.action(api.scheduler.reconcileWorkerPool, {
+      nowMs,
+      flyApiToken: "fly-token",
+      convexUrl: "https://example.convex.cloud",
+      scalingPolicy: {
+        maxWorkers: 5,
+        queuePerWorkerTarget: 1,
+        spawnStep: 1,
+        idleTimeoutMs: 300_000,
+        reconcileIntervalMs: 15_000,
+      },
+      providerConfig: TEST_PROVIDER_CONFIG,
+    });
+
+    expect(reconcile.spawned).toBe(1);
+  });
+
+  test("snapshot restore should require a matching conversation when conversationId is provided", async () => {
+    const t = initConvexTest();
+    const nowMs = Date.UTC(2026, 0, 1, 17, 50, 0);
+    const snapshot = await t.mutation(api.queue.prepareDataSnapshotUpload as any, {
+      workerId: "worker-snapshot-1",
+      workspaceId: "default",
+      agentKey: "support-agent",
+      conversationId: "telegram:chat:snapshot-a",
+      reason: "manual",
+      nowMs,
+    });
+    const storageId = await t.run(async (ctx) => {
+      return await ctx.storage.store(new Blob(["snapshot-a"]));
+    });
+    const finalized = await t.mutation(api.queue.finalizeDataSnapshotUpload as any, {
+      workerId: "worker-snapshot-1",
+      snapshotId: snapshot.snapshotId,
+      storageId,
+      sha256: "beadfeed",
+      sizeBytes: 10,
+      nowMs: nowMs + 1,
+    });
+    expect(finalized).toBe(true);
+
+    const missingConversationSnapshot = await t.query(
+      api.queue.getLatestDataSnapshotForRestore as any,
+      {
+        workspaceId: "default",
+        agentKey: "support-agent",
+        conversationId: "telegram:chat:snapshot-b",
+        nowMs: nowMs + 2,
+      },
+    );
+    expect(missingConversationSnapshot).toBeNull();
+  });
+
+  test("release stuck jobs should clear worker assignment when the lease is recovered", async () => {
+    const t = initConvexTest();
+    const nowMs = Date.UTC(2026, 0, 1, 17, 55, 0);
+    vi.setSystemTime(nowMs);
+    await t.mutation(api.queue.upsertAgentProfile, {
+      agentKey: "support-agent",
+      version: "1.0.0",
+      secretsRef: [],
+      enabled: true,
+    });
+
+    const messageId = await t.mutation(api.queue.enqueueMessage, {
+      conversationId: "telegram:chat:lease-clear",
+      agentKey: "support-agent",
+      payload: {
+        provider: "telegram",
+        providerUserId: "u-lease-clear",
+        messageText: "recover me",
+      },
+      nowMs,
+    });
+    const claim = await t.mutation(api.lib.claim, {
+      workerId: "worker-lease-clear-1",
+      conversationId: "telegram:chat:lease-clear",
+      nowMs,
+    });
+    expect(claim?.messageId).toBe(messageId);
+
+    const released = await t.mutation(api.queue.releaseStuckJobs, {
+      nowMs: nowMs + DEFAULT_CONFIG.lease.leaseMs + 1,
+      limit: 10,
+    });
+    expect(released.requeued).toBe(1);
+    expect(released.unlocked).toBe(1);
+
+    const workers = await t.query((internal.queue as any).listWorkersForScheduler, {});
+    const worker = workers.find(
+      (row: { workerId: string }) => row.workerId === "worker-lease-clear-1",
+    );
+    expect(worker?.assignment).toBeNull();
+  });
+
   test("checkIdleShutdowns should backfill missing scheduledShutdownAt for idle active workers", async () => {
     const t = initConvexTest();
     const nowMs = Date.UTC(2026, 0, 1, 18, 0, 0);