@okrlinkhub/agent-factory 3.0.1 → 3.0.3

package/src/component/lib.test.ts

@@ -833,6 +833,7 @@ describe("component lib", () => {
       machineId,
       appName: TEST_PROVIDER_CONFIG.appName,
       region: TEST_PROVIDER_CONFIG.region,
+      volumeId,
     });
 
     const completionTime = claimTime + 60_000;
@@ -926,9 +927,6 @@ describe("component lib", () => {
       if (url.endsWith(`/apps/${TEST_PROVIDER_CONFIG.appName}/machines`) && method === "GET") {
         return jsonResponse([]);
       }
-      if (url.endsWith(`/machines/machine-orphan-1`) && method === "GET") {
-        return new Response("not found", { status: 404 });
-      }
       if (url.endsWith(`/apps/${TEST_PROVIDER_CONFIG.appName}/volumes`) && method === "GET") {
         return jsonResponse([
           {
@@ -955,6 +953,7 @@ describe("component lib", () => {
       machineId: "machine-orphan-1",
       appName: TEST_PROVIDER_CONFIG.appName,
       region: TEST_PROVIDER_CONFIG.region,
+      volumeId,
     });
 
     const result = await t.action(api.scheduler.checkIdleShutdowns, {
@@ -971,12 +970,195 @@ describe("component lib", () => {
         ((call[1] as RequestInit | undefined)?.method ?? "GET") === "DELETE",
     );
     expect(deleteVolumeCalls).toHaveLength(1);
+    const machineDetailCalls = fetchMock.mock.calls.filter(
+      (call) =>
+        String(call[0]).endsWith(`/machines/machine-orphan-1`) &&
+        ((call[1] as RequestInit | undefined)?.method ?? "GET") === "GET",
+    );
+    expect(machineDetailCalls).toHaveLength(0);
 
     const workers = await t.query((internal.queue as any).listWorkersForScheduler, {});
     const worker = workers.find((row: { workerId: string }) => row.workerId === workerId);
     expect(worker?.status).toBe("stopped");
   });
 
+  test("reconcile should persist volumeId on the worker row after spawn", async () => {
+    const t = initConvexTest();
+    const nowMs = Date.UTC(2026, 0, 2, 9, 0, 0);
+    const volumeId = "vol-persist-1";
+    const machineId = "machine-persist-1";
+    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(
+          fetchMock.mock.calls.some(
+            (call) =>
+              String(call[0]).endsWith(`/apps/${TEST_PROVIDER_CONFIG.appName}/machines`) &&
+              ((call[1] as RequestInit | undefined)?.method ?? "GET") === "POST",
+          )
+            ? [
+                {
+                  id: machineId,
+                  name: `afw-${nowMs}-0`,
+                  region: TEST_PROVIDER_CONFIG.region,
+                  state: "started",
+                  config: { image: TEST_PROVIDER_CONFIG.image },
+                },
+              ]
+            : [],
+        );
+      }
+      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: volumeId,
+          name: buildDedicatedVolumeName(TEST_PROVIDER_CONFIG.volumeName, `afw-${nowMs}-0`),
+          region: TEST_PROVIDER_CONFIG.region,
+        });
+      }
+      if (url.endsWith(`/apps/${TEST_PROVIDER_CONFIG.appName}/machines`) && method === "POST") {
+        return jsonResponse({
+          id: machineId,
+          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:persist-volume",
+      agentKey: "support-agent",
+      payload: {
+        provider: "telegram",
+        providerUserId: "u-persist-volume",
+        messageText: "hello",
+      },
+      nowMs,
+    });
+
+    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,
+    });
+
+    const workers = await t.query((internal.queue as any).listWorkersForScheduler, {});
+    const worker = workers.find((row: { workerId: string }) => row.workerId === `afw-${nowMs}-0`);
+    expect(worker?.machineId).toBe(machineId);
+    expect(worker?.volumeId).toBe(volumeId);
+  });
+
+  test("runFlyCleanup should destroy machines and volumes with final verification", async () => {
+    const t = initConvexTest();
+    await t.mutation(api.queue.importPlaintextSecret, {
+      secretRef: "fly.apiToken",
+      plaintextValue: "fly-token",
+    });
+    await t.mutation(api.queue.setProviderRuntimeConfig, {
+      providerConfig: TEST_PROVIDER_CONFIG,
+    });
+
+    let machineListCalls = 0;
+    let volumeListCalls = 0;
+    const deletedMachines: Array<string> = [];
+    const deletedVolumes: Array<string> = [];
+
+    vi.stubGlobal(
+      "fetch",
+      vi.fn(async (input: RequestInfo | URL, init?: RequestInit) => {
+        const url = String(input);
+        const method = init?.method ?? "GET";
+        const appBase = "https://api.machines.dev/v1/apps/agent-factory-workers-test";
+
+        if (url === `${appBase}` && method === "GET") {
+          return jsonResponse({ name: "agent-factory-workers-test" });
+        }
+        if (url === `${appBase}/machines` && method === "GET") {
+          machineListCalls += 1;
+          return jsonResponse(machineListCalls === 1 ? [{ id: "machine-1" }, { id: "machine-2" }] : []);
+        }
+        if (url === `${appBase}/volumes` && method === "GET") {
+          volumeListCalls += 1;
+          return jsonResponse(
+            volumeListCalls === 1
+              ? [{ id: "volume-1" }, { id: "volume-2" }, { id: "volume-3" }]
+              : [],
+          );
+        }
+        if (url.includes("/machines/") && url.endsWith("/cordon") && method === "POST") {
+          return jsonResponse({});
+        }
+        if (url.includes("/machines/") && url.endsWith("/stop") && method === "POST") {
+          return jsonResponse({});
+        }
+        if (url.includes("/machines/") && method === "DELETE") {
+          const segments = url.split("/");
+          deletedMachines.push(segments[segments.length - 1]!);
+          return new Response(null, { status: 204 });
+        }
+        if (url.includes("/volumes/") && method === "DELETE") {
+          const segments = url.split("/");
+          deletedVolumes.push(segments[segments.length - 1]!);
+          return new Response(null, { status: 204 });
+        }
+        throw new Error(`Unexpected fetch ${method} ${url}`);
+      }),
+    );
+
+    const report = await t.action((api.lib as any).runFlyCleanup, {});
+
+    expect(report.appName).toBe("agent-factory-workers-test");
+    expect(report.machinesFound).toBe(2);
+    expect(report.machinesDeleted).toBe(2);
+    expect(report.machinesRemaining).toBe(0);
+    expect(report.volumesFound).toBe(3);
+    expect(report.volumesDeleted).toBe(3);
+    expect(report.volumesRemaining).toBe(0);
+    expect(report.errors).toEqual([]);
+    expect(report.warnings).toEqual([]);
+    expect(deletedMachines.sort()).toEqual(["machine-1", "machine-2"]);
+    expect(deletedVolumes.sort()).toEqual(["volume-1", "volume-2", "volume-3"]);
+  });
+
+  test("runFlyCleanup should require an active fly secret", async () => {
+    const t = initConvexTest();
+    await t.mutation(api.queue.setProviderRuntimeConfig, {
+      providerConfig: TEST_PROVIDER_CONFIG,
+    });
+
+    await expect(t.action((api.lib as any).runFlyCleanup, {})).rejects.toThrow(
+      "Missing active 'fly.apiToken' secret.",
+    );
+  });
+
   test("scheduler count includes queued and in-progress conversations", async () => {
     const t = initConvexTest();
     await t.mutation(api.queue.upsertAgentProfile, {
@@ -1248,6 +1430,14 @@ describe("component lib", () => {
     const t = initConvexTest();
     const nowMs = Date.UTC(2026, 0, 1, 17, 0, 0);
     vi.setSystemTime(nowMs);
+    let machineCreateBody:
+      | {
+          name?: string;
+          config?: {
+            env?: Record<string, string>;
+          };
+        }
+      | null = null;
     const fetchMock = vi.fn(async (input: RequestInfo | URL, init?: RequestInit) => {
       const url = String(input);
       const method = init?.method ?? "GET";
@@ -1273,7 +1463,13 @@ describe("component lib", () => {
         });
       }
       if (url.endsWith(`/apps/${TEST_PROVIDER_CONFIG.appName}/machines`) && method === "POST") {
-        const body = JSON.parse(String(init?.body ?? "{}")) as { name?: string };
+        const body = JSON.parse(String(init?.body ?? "{}")) as {
+          name?: string;
+          config?: {
+            env?: Record<string, string>;
+          };
+        };
+        machineCreateBody = body;
         return jsonResponse({
           id: "machine-new-worker",
           name: body.name,
@@ -1349,6 +1545,28 @@ describe("component lib", () => {
 
     expect(reconcile.spawned).toBe(1);
     expect(reconcile.activeWorkers).toBe(2);
+    const machineCreateEnv = (
+      machineCreateBody as
+        | {
+            config?: {
+              env?: Record<string, string>;
+            };
+          }
+        | null
+    )?.config?.env;
+    expect(machineCreateEnv?.OPENCLAW_CONVERSATION_ID).toBe("telegram:chat:spawn-b");
+    expect(machineCreateEnv?.OPENCLAW_AGENT_KEY).toBe("support-agent");
+
+    const workers = await t.query((internal.queue as any).listWorkersForScheduler, {});
+    const spawnedWorker = workers.find(
+      (row: { workerId: string }) => row.workerId === `afw-${nowMs}-0`,
+    );
+    expect(spawnedWorker?.assignment).toEqual({
+      conversationId: "telegram:chat:spawn-b",
+      agentKey: "support-agent",
+      leaseId: `spawn:afw-${nowMs}-0`,
+      assignedAt: nowMs,
+    });
   });
 
   test("scheduler should forward OpenClaw bridge env to spawned machines", async () => {
@@ -1623,6 +1841,60 @@ describe("component lib", () => {
     });
   });
 
+  test("preassigned workers should claim only their assigned conversation on first claim", async () => {
+    const t = initConvexTest();
+    const nowMs = Date.UTC(2026, 0, 1, 17, 35, 0);
+    vi.setSystemTime(nowMs);
+    await t.mutation(api.queue.upsertAgentProfile, {
+      agentKey: "support-agent",
+      version: "1.0.0",
+      secretsRef: [],
+      enabled: true,
+    });
+
+    await t.mutation(api.queue.enqueueMessage, {
+      conversationId: "telegram:chat:preassign-a",
+      agentKey: "support-agent",
+      payload: {
+        provider: "telegram",
+        providerUserId: "u-preassign-a",
+        messageText: "first",
+      },
+      nowMs,
+    });
+    const messageB = await t.mutation(api.queue.enqueueMessage, {
+      conversationId: "telegram:chat:preassign-b",
+      agentKey: "support-agent",
+      payload: {
+        provider: "telegram",
+        providerUserId: "u-preassign-b",
+        messageText: "second",
+      },
+      nowMs: nowMs + 1,
+    });
+
+    await t.mutation(internal.queue.upsertWorkerState, {
+      workerId: "worker-preassigned-1",
+      provider: "fly",
+      status: "active",
+      load: 0,
+      nowMs,
+      assignment: {
+        conversationId: "telegram:chat:preassign-b",
+        agentKey: "support-agent",
+        leaseId: "spawn:worker-preassigned-1",
+        assignedAt: nowMs,
+      },
+    });
+
+    const claim = await t.mutation(api.lib.claim, {
+      workerId: "worker-preassigned-1",
+      nowMs: nowMs + 10,
+    });
+    expect(claim?.messageId).toBe(messageB);
+    expect(claim?.conversationId).toBe("telegram:chat:preassign-b");
+  });
+
   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);
@@ -2465,6 +2737,9 @@ describe("component lib", () => {
       if (url.endsWith(`/apps/${TEST_PROVIDER_CONFIG.appName}/machines`) && method === "POST") {
         throw new Error("simulated spawn failure");
       }
+      if (url.endsWith(`/volumes/vol-spawn-failure-1`) && method === "DELETE") {
+        return emptyResponse();
+      }
       throw new Error(`Unexpected fetch ${method} ${url}`);
     });
     vi.stubGlobal("fetch", fetchMock);
@@ -2514,6 +2789,7 @@ describe("component lib", () => {
     const worker = workers.find((row: { workerId: string }) => row.workerId === `afw-${nowMs}-0`);
     expect(worker?.status).toBe("stopped");
     expect(worker?.machineId).toBeNull();
+    expect(worker?.volumeId).toBeNull();
   });
 
   test("push jobs should dispatch scheduled messages with a stable user-agent conversation id", async () => {

package/src/component/lib.ts

@@ -34,6 +34,7 @@ export {
   checkIdleShutdowns,
 } from "./scheduler.js";
 export { deleteFlyVolumeManual as deleteFlyVolume } from "./providers/fly.js";
+export { runFlyCleanup } from "./flyCleanup.js";
 
 export {
   bindUserAgent,

package/src/component/providers/fly.ts

@@ -10,9 +10,8 @@ export type SpawnWorkerInput = {
   appName: string;
   image: string;
   region: string;
-  volumeName: string;
+  volumeId: string;
   volumePath: string;
-  volumeSizeGb: number;
   cpuKind?: string;
   cpus?: number;
   memoryMb?: number;
@@ -22,13 +21,27 @@ export type SpawnWorkerInput = {
 export type ProviderWorker = {
   workerId: string;
   machineId: string;
+  volumeId?: string;
   region?: string;
   image?: string;
   status: WorkerProviderStatus;
   rawState?: string;
 };
 
+export type WorkerVolume = {
+  volumeId: string;
+  volumeName: string;
+  region?: string;
+};
+
 export interface WorkerProvider {
+  ensureWorkerVolume(input: {
+    appName: string;
+    workerId: string;
+    region: string;
+    volumeName: string;
+    volumeSizeGb: number;
+  }): Promise<WorkerVolume>;
   spawnWorker(input: SpawnWorkerInput): Promise<ProviderWorker>;
   listWorkers(appName: string): Promise<Array<ProviderWorker>>;
   terminateWorker(appName: string, machineId: string): Promise<void>;
@@ -40,6 +53,7 @@ export interface WorkerProvider {
     machineId?: string | null;
     region?: string;
     volumeName: string;
+    volumeId?: string | null;
   }): Promise<void>;
 }
 
@@ -72,7 +86,14 @@ export class FlyMachinesProvider implements WorkerProvider {
     private readonly baseUrl: string = "https://api.machines.dev/v1",
   ) {}
 
-  async spawnWorker(input: SpawnWorkerInput): Promise<ProviderWorker> {
+  async ensureWorkerVolume(input: {
+    appName: string;
+    workerId: string;
+    region: string;
+    volumeName: string;
+    volumeSizeGb: number;
+  }): Promise<WorkerVolume> {
+    const volumeName = buildDedicatedVolumeName(input.volumeName, input.workerId);
     const volumeId = await this.resolveOrCreateVolumeId(
       input.appName,
       input.volumeName,
@@ -80,6 +101,14 @@ export class FlyMachinesProvider implements WorkerProvider {
       input.region,
       input.volumeSizeGb,
     );
+    return {
+      volumeId,
+      volumeName,
+      region: input.region,
+    };
+  }
+
+  async spawnWorker(input: SpawnWorkerInput): Promise<ProviderWorker> {
     const payload = {
       name: input.workerId,
       region: input.region,
@@ -92,7 +121,7 @@ export class FlyMachinesProvider implements WorkerProvider {
         },
         mounts: [
           {
-            volume: volumeId,
+            volume: input.volumeId,
             path: input.volumePath,
           } satisfies FlyMachineMount,
         ],
@@ -110,6 +139,7 @@ export class FlyMachinesProvider implements WorkerProvider {
     return {
       workerId: input.workerId,
       machineId: machine.id,
+      volumeId: input.volumeId,
       region: machine.region ?? input.region,
       image: machine.config?.image_ref ?? machine.config?.image ?? input.image,
       status: mapFlyStateToProviderStatus(machine.state),
@@ -182,9 +212,13 @@ export class FlyMachinesProvider implements WorkerProvider {
     machineId?: string | null;
     region?: string;
     volumeName: string;
+    volumeId?: string | null;
   }): Promise<void> {
     const volumeIds = new Set<string>();
-    if (input.machineId) {
+    if (input.volumeId) {
+      volumeIds.add(input.volumeId);
+    }
+    if (!input.volumeId && input.machineId) {
       const machineVolumeIds = await this.getMachineVolumeIds(input.appName, input.machineId);
       for (const volumeId of machineVolumeIds) {
         volumeIds.add(volumeId);

package/src/component/queue.ts

@@ -188,6 +188,11 @@ const workerSpawnOpenClawEnvValidator = v.object({
   OPENCLAW_LINKING_SHARED_SECRET: v.optional(v.string()),
 });
 
+const schedulerConversationTargetValidator = v.object({
+  conversationId: v.string(),
+  agentKey: v.string(),
+});
+
 const messageRuntimeConfigValidator = v.object({
   systemPrompt: v.optional(v.string()),
   telegramAttachmentRetentionMs: v.optional(v.number()),
@@ -489,6 +494,47 @@ export const getWorkerSpawnOpenClawEnv = internalQuery({
   },
 });
 
+export const getActiveConversationsForScheduler = internalQuery({
+  args: {
+    nowMs: v.optional(v.number()),
+    limit: v.optional(v.number()),
+  },
+  returns: v.array(schedulerConversationTargetValidator),
+  handler: async (ctx, args) => {
+    const nowMs = args.nowMs ?? Date.now();
+    const limit = Math.max(1, args.limit ?? 1000);
+    const queuedJobs = await ctx.db
+      .query("messageQueue")
+      .withIndex("by_status_and_scheduledFor", (q) =>
+        q.eq("status", "queued").lte("scheduledFor", nowMs),
+      )
+      .take(limit);
+    const processingJobs = await ctx.db
+      .query("messageQueue")
+      .withIndex("by_status_and_leaseExpiresAt", (q) =>
+        q.eq("status", "processing").gt("leaseExpiresAt", nowMs),
+      )
+      .take(limit);
+
+    const conversations = new Map<string, { conversationId: string; agentKey: string }>();
+    for (const job of [...queuedJobs, ...processingJobs]) {
+      const key = `${job.agentKey}::${job.conversationId}`;
+      if (!conversations.has(key)) {
+        conversations.set(key, {
+          conversationId: job.conversationId,
+          agentKey: job.agentKey,
+        });
+      }
+    }
+
+    return Array.from(conversations.values()).sort(
+      (left, right) =>
+        left.agentKey.localeCompare(right.agentKey) ||
+        left.conversationId.localeCompare(right.conversationId),
+    );
+  },
+});
+
 export const getProviderRuntimeConfig = internalQuery({
   args: {},
   returns: v.union(v.null(), providerConfigValidator),
@@ -2277,8 +2323,11 @@ export const upsertWorkerState = internalMutation({
     machineId: v.optional(v.string()),
     appName: v.optional(v.string()),
     region: v.optional(v.string()),
+    volumeId: v.optional(v.string()),
+    assignment: v.optional(v.union(v.null(), workerAssignmentValidator)),
     clearLastSnapshotId: v.optional(v.boolean()),
     clearMachineRef: v.optional(v.boolean()),
+    clearVolumeId: v.optional(v.boolean()),
   },
   returns: v.null(),
   handler: async (ctx, args) => {
@@ -2299,7 +2348,8 @@ export const upsertWorkerState = internalMutation({
           args.status === "stopped" || args.status === "stopping"
             ? (args.stoppedAt ?? nowMs)
             : undefined,
-        assignment: undefined,
+        volumeId: args.volumeId,
+        assignment: args.assignment ?? undefined,
         machineRef:
           args.machineId && args.appName
             ? {
@@ -2329,7 +2379,8 @@ export const upsertWorkerState = internalMutation({
           ? (args.stoppedAt ?? worker.stoppedAt ?? nowMs)
           : undefined,
       lastSnapshotId: args.clearLastSnapshotId ? undefined : worker.lastSnapshotId,
-      assignment: worker.assignment,
+      volumeId: args.clearVolumeId ? undefined : (args.volumeId ?? worker.volumeId),
+      assignment: args.assignment === undefined ? worker.assignment : (args.assignment ?? undefined),
       machineRef:
         args.clearMachineRef
           ? undefined
@@ -2548,6 +2599,7 @@ export const listWorkersForScheduler = internalQuery({
       machineId: v.union(v.null(), v.string()),
       appName: v.union(v.null(), v.string()),
       region: v.union(v.null(), v.string()),
+      volumeId: v.union(v.null(), v.string()),
     }),
   ),
   handler: async (ctx) => {
@@ -2565,6 +2617,7 @@ export const listWorkersForScheduler = internalQuery({
       machineId: worker.machineRef?.machineId ?? null,
       appName: worker.machineRef?.appName ?? null,
       region: worker.machineRef?.region ?? null,
+      volumeId: worker.volumeId ?? null,
     }));
   },
 });