@okrlinkhub/agent-factory 3.0.2 → 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, {
@@ -2555,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);
@@ -2604,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

@@ -2323,9 +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) => {
@@ -2346,6 +2348,7 @@ export const upsertWorkerState = internalMutation({
           args.status === "stopped" || args.status === "stopping"
             ? (args.stoppedAt ?? nowMs)
             : undefined,
+        volumeId: args.volumeId,
         assignment: args.assignment ?? undefined,
         machineRef:
           args.machineId && args.appName
@@ -2376,6 +2379,7 @@ export const upsertWorkerState = internalMutation({
           ? (args.stoppedAt ?? worker.stoppedAt ?? nowMs)
           : undefined,
       lastSnapshotId: args.clearLastSnapshotId ? undefined : worker.lastSnapshotId,
+      volumeId: args.clearVolumeId ? undefined : (args.volumeId ?? worker.volumeId),
       assignment: args.assignment === undefined ? worker.assignment : (args.assignment ?? undefined),
       machineRef:
         args.clearMachineRef
@@ -2595,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) => {
@@ -2612,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,
     }));
   },
 });

package/src/component/scheduler.ts

@@ -73,6 +73,7 @@ type SchedulerWorkerRow = {
   machineId: string | null;
   appName: string | null;
   region: string | null;
+  volumeId: string | null;
 };
 
 type SchedulerConversationTarget = {
@@ -372,6 +373,13 @@ async function runWorkerLifecycleCycle(
               assignedAt: input.nowMs,
             }
           : undefined;
+        const workerVolume = await input.provider.ensureWorkerVolume({
+          appName: input.providerConfig.appName,
+          workerId,
+          region: input.providerConfig.region,
+          volumeName: input.providerConfig.volumeName,
+          volumeSizeGb: input.providerConfig.volumeSizeGb,
+        });
         await ctx.runMutation(internal.queue.upsertWorkerState, {
           workerId,
           provider: input.providerConfig.kind,
@@ -379,6 +387,7 @@ async function runWorkerLifecycleCycle(
           load: 0,
           nowMs: input.nowMs,
           scheduledShutdownAt: input.nowMs + input.scaling.idleTimeoutMs,
+          volumeId: workerVolume.volumeId,
           assignment,
         });
         let created;
@@ -388,9 +397,8 @@ async function runWorkerLifecycleCycle(
             appName: input.providerConfig.appName,
             image: input.providerConfig.image,
             region: input.providerConfig.region,
-            volumeName: input.providerConfig.volumeName,
+            volumeId: workerVolume.volumeId,
             volumePath: input.providerConfig.volumePath,
-            volumeSizeGb: input.providerConfig.volumeSizeGb,
             env: compactEnv({
               ...DEFAULT_WORKER_RUNTIME_ENV,
               ...forwardedOpenClawEnv,
@@ -408,6 +416,13 @@ async function runWorkerLifecycleCycle(
               error instanceof Error ? error.message : String(error)
             }`,
           );
+          await input.provider.cleanupWorkerStorage({
+            appName: input.providerConfig.appName,
+            workerId,
+            region: input.providerConfig.region,
+            volumeName: input.providerConfig.volumeName,
+            volumeId: workerVolume.volumeId,
+          });
           await transitionWorkerToDraining(
             ctx,
             {
@@ -423,6 +438,7 @@ async function runWorkerLifecycleCycle(
               machineId: null,
               appName: input.providerConfig.appName,
               region: input.providerConfig.region,
+              volumeId: workerVolume.volumeId,
             },
             input.providerConfig,
             input.nowMs,
@@ -443,6 +459,7 @@ async function runWorkerLifecycleCycle(
               machineId: null,
               appName: input.providerConfig.appName,
               region: input.providerConfig.region,
+              volumeId: workerVolume.volumeId,
             },
             input.providerConfig,
             input.nowMs,
@@ -463,6 +480,7 @@ async function runWorkerLifecycleCycle(
               machineId: null,
               appName: input.providerConfig.appName,
               region: input.providerConfig.region,
+              volumeId: workerVolume.volumeId,
             },
             input.providerConfig,
             input.nowMs,
@@ -479,6 +497,7 @@ async function runWorkerLifecycleCycle(
           machineId: created.machineId,
           appName: input.providerConfig.appName,
           region: created.region,
+          volumeId: created.volumeId ?? workerVolume.volumeId,
           assignment,
         });
         await scheduleIdleShutdownWatch(
@@ -679,6 +698,7 @@ async function transitionWorkerToStopped(
     scheduledShutdownAt: worker.scheduledShutdownAt ?? nowMs,
     stoppedAt: worker.stoppedAt ?? nowMs,
     clearMachineRef: true,
+    clearVolumeId: true,
   });
 }
 
@@ -706,6 +726,7 @@ async function finalizeWorkerTeardown(input: {
     machineId,
     region: input.worker.region ?? input.providerConfig.region,
     volumeName: input.providerConfig.volumeName,
+    volumeId: input.worker.volumeId,
   });
   return true;
 }

package/src/component/schema.ts

@@ -127,6 +127,7 @@ export default defineSchema({
   workers: defineTable({
     workerId: v.string(),
     provider: v.string(),
+    volumeId: v.optional(v.string()),
     machineRef: v.optional(
       v.object({
         appName: v.string(),