@okrlinkhub/agent-factory 3.0.1 → 3.0.3

package/src/component/scheduler.ts

@@ -73,6 +73,12 @@ type SchedulerWorkerRow = {
   machineId: string | null;
   appName: string | null;
   region: string | null;
+  volumeId: string | null;
+};
+
+type SchedulerConversationTarget = {
+  conversationId: string;
+  agentKey: string;
 };
 
 const PROVIDER_RECONCILE_GRACE_MS = 90_000;
@@ -182,11 +188,11 @@ async function runReconcileWorkerPool(
   }
   const workspaceId = args.workspaceId ?? "default";
   const provider = resolveProvider(providerConfig.kind, flyApiToken);
-  const activeConversationIds: Array<string> = await ctx.runQuery(
-    (internal.queue as any).getActiveConversationIdsForScheduler,
+  const activeConversations: Array<SchedulerConversationTarget> = await ctx.runQuery(
+    (internal.queue as any).getActiveConversationsForScheduler,
     { nowMs, limit: 1000 },
   );
-  const activeConversationCount = activeConversationIds.length;
+  const activeConversationCount = activeConversations.length;
   const cycle = await runWorkerLifecycleCycle(ctx, {
     nowMs,
     provider,
@@ -195,7 +201,7 @@ async function runReconcileWorkerPool(
     allowSpawn: true,
     convexUrl,
     workspaceId,
-    activeConversationIds,
+    activeConversations,
     desiredActiveWorkers: clamp(activeConversationCount, 0, scaling.maxWorkers),
   });
   if (activeConversationCount > 0 || cycle.pending > 0) {
@@ -258,7 +264,7 @@ async function runEnforceIdleShutdowns(
     scaling: DEFAULT_CONFIG.scaling,
     allowSpawn: false,
     desiredActiveWorkers: 0,
-    activeConversationIds: [],
+    activeConversations: [],
   });
 
   if (cycle.pending > 0) {
@@ -282,7 +288,7 @@ async function runWorkerLifecycleCycle(
     scaling: typeof DEFAULT_CONFIG.scaling;
     allowSpawn: boolean;
     desiredActiveWorkers: number;
-    activeConversationIds: Array<string>;
+    activeConversations: Array<SchedulerConversationTarget>;
     convexUrl?: string;
     workspaceId?: string;
   },
@@ -338,7 +344,7 @@ async function runWorkerLifecycleCycle(
   if (input.allowSpawn && input.desiredActiveWorkers > 0) {
     const claimableWorkers = countWorkersAvailableForActiveConversations(
       filterScopedWorkers(workerRows, input.providerConfig.appName),
-      input.activeConversationIds,
+      input.activeConversations,
       staleHeartbeatCutoff,
     );
     if (input.desiredActiveWorkers > claimableWorkers) {
@@ -350,8 +356,30 @@ async function runWorkerLifecycleCycle(
         input.scaling.spawnStep,
         input.desiredActiveWorkers - claimableWorkers,
       );
+      const spawnTargets = selectSpawnTargetsForActiveConversations(
+        filterScopedWorkers(workerRows, input.providerConfig.appName),
+        input.activeConversations,
+        staleHeartbeatCutoff,
+        toSpawn,
+      );
       for (let index = 0; index < toSpawn; index += 1) {
         const workerId = `afw-${input.nowMs}-${index}`;
+        const target = spawnTargets[index];
+        const assignment = target
+          ? {
+              conversationId: target.conversationId,
+              agentKey: target.agentKey,
+              leaseId: `spawn:${workerId}`,
+              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,
@@ -359,6 +387,8 @@ async function runWorkerLifecycleCycle(
           load: 0,
           nowMs: input.nowMs,
           scheduledShutdownAt: input.nowMs + input.scaling.idleTimeoutMs,
+          volumeId: workerVolume.volumeId,
+          assignment,
         });
         let created;
         try {
@@ -367,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,
@@ -377,6 +406,8 @@ async function runWorkerLifecycleCycle(
               WORKSPACE_ID: input.workspaceId ?? "default",
               WORKER_ID: workerId,
               WORKER_IDLE_TIMEOUT_MS: String(input.scaling.idleTimeoutMs),
+              OPENCLAW_AGENT_KEY: target?.agentKey,
+              OPENCLAW_CONVERSATION_ID: target?.conversationId,
             }),
           });
         } catch (error) {
@@ -385,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,
             {
@@ -400,6 +438,7 @@ async function runWorkerLifecycleCycle(
               machineId: null,
               appName: input.providerConfig.appName,
               region: input.providerConfig.region,
+              volumeId: workerVolume.volumeId,
             },
             input.providerConfig,
             input.nowMs,
@@ -420,6 +459,7 @@ async function runWorkerLifecycleCycle(
               machineId: null,
               appName: input.providerConfig.appName,
               region: input.providerConfig.region,
+              volumeId: workerVolume.volumeId,
             },
             input.providerConfig,
             input.nowMs,
@@ -440,6 +480,7 @@ async function runWorkerLifecycleCycle(
               machineId: null,
               appName: input.providerConfig.appName,
               region: input.providerConfig.region,
+              volumeId: workerVolume.volumeId,
             },
             input.providerConfig,
             input.nowMs,
@@ -456,6 +497,8 @@ async function runWorkerLifecycleCycle(
           machineId: created.machineId,
           appName: input.providerConfig.appName,
           region: created.region,
+          volumeId: created.volumeId ?? workerVolume.volumeId,
+          assignment,
         });
         await scheduleIdleShutdownWatch(
           ctx,
@@ -655,6 +698,7 @@ async function transitionWorkerToStopped(
     scheduledShutdownAt: worker.scheduledShutdownAt ?? nowMs,
     stoppedAt: worker.stoppedAt ?? nowMs,
     clearMachineRef: true,
+    clearVolumeId: true,
   });
 }
 
@@ -682,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;
 }
@@ -862,10 +907,46 @@ function filterScopedWorkers(workerRows: Array<SchedulerWorkerRow>, appName: str
 
 function countWorkersAvailableForActiveConversations(
   workerRows: Array<SchedulerWorkerRow>,
-  activeConversationIds: Array<string>,
+  activeConversations: Array<SchedulerConversationTarget>,
+  staleHeartbeatCutoff: number,
+) {
+  const coverage = summarizeWorkerConversationCoverage(
+    workerRows,
+    activeConversations,
+    staleHeartbeatCutoff,
+  );
+  return coverage.unassignedWorkers + coverage.assignedConversationKeys.size;
+}
+
+function selectSpawnTargetsForActiveConversations(
+  workerRows: Array<SchedulerWorkerRow>,
+  activeConversations: Array<SchedulerConversationTarget>,
   staleHeartbeatCutoff: number,
+  limit: number,
 ) {
-  const activeConversationSet = new Set(activeConversationIds);
+  const coverage = summarizeWorkerConversationCoverage(
+    workerRows,
+    activeConversations,
+    staleHeartbeatCutoff,
+  );
+  const uncoveredConversations = activeConversations.filter(
+    (conversation) =>
+      !coverage.assignedConversationKeys.has(getConversationTargetKey(conversation)),
+  );
+  return uncoveredConversations.slice(
+    coverage.unassignedWorkers,
+    coverage.unassignedWorkers + Math.max(0, limit),
+  );
+}
+
+function summarizeWorkerConversationCoverage(
+  workerRows: Array<SchedulerWorkerRow>,
+  activeConversations: Array<SchedulerConversationTarget>,
+  staleHeartbeatCutoff: number,
+) {
+  const activeConversationSet = new Set(
+    activeConversations.map((conversation) => getConversationTargetKey(conversation)),
+  );
   const assignedConversationKeys = new Set<string>();
   let unassignedWorkers = 0;
   for (const worker of workerRows) {
@@ -876,13 +957,16 @@ function countWorkersAvailableForActiveConversations(
       unassignedWorkers += 1;
       continue;
     }
-    if (activeConversationSet.has(worker.assignment.conversationId)) {
-      assignedConversationKeys.add(
-        `${worker.assignment.agentKey}::${worker.assignment.conversationId}`,
-      );
+    const assignmentKey = getConversationTargetKey(worker.assignment);
+    if (activeConversationSet.has(assignmentKey)) {
+      assignedConversationKeys.add(assignmentKey);
     }
   }
-  return unassignedWorkers + assignedConversationKeys.size;
+  return { assignedConversationKeys, unassignedWorkers };
+}
+
+function getConversationTargetKey(input: { conversationId: string; agentKey: string }) {
+  return `${input.agentKey}::${input.conversationId}`;
 }
 
 function deriveScheduledShutdownAt(

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(),