@okrlinkhub/agent-factory 1.0.0 → 1.0.2

package/src/component/queue.ts

@@ -48,6 +48,13 @@ const claimedJobValidator = v.object({
   payload: queuePayloadValidator,
 });
 
+const workerAssignmentValidator = v.object({
+  conversationId: v.string(),
+  agentKey: v.string(),
+  leaseId: v.string(),
+  assignedAt: v.number(),
+});
+
 const secretStatusValidator = v.object({
   secretRef: v.string(),
   hasActive: v.boolean(),
@@ -166,6 +173,10 @@ export const enqueueMessage = mutation({
         contextHistory: [],
         pendingToolCalls: [],
       });
+    } else if (existingConversation.agentKey !== args.agentKey) {
+      throw new Error(
+        `Conversation '${args.conversationId}' is already bound to agent '${existingConversation.agentKey}', cannot enqueue for '${args.agentKey}'.`,
+      );
     }
 
     const priority = Math.min(
@@ -1024,11 +1035,13 @@ export const attachMessageMetadata = mutation({
 export const claimNextJob = mutation({
   args: {
     workerId: v.string(),
+    conversationId: v.optional(v.string()),
     nowMs: v.optional(v.number()),
   },
   returns: v.union(v.null(), claimedJobValidator),
   handler: async (ctx, args) => {
     const nowMs = args.nowMs ?? Date.now();
+    const staleHeartbeatCutoff = nowMs - DEFAULT_CONFIG.lease.staleAfterMs;
     const worker = await ctx.db
       .query("workers")
       .withIndex("by_workerId", (q) => q.eq("workerId", args.workerId))
@@ -1036,6 +1049,14 @@ export const claimNextJob = mutation({
     if (worker && !isWorkerClaimable(worker.status)) {
       return null;
     }
+    if (
+      worker?.assignment &&
+      args.conversationId &&
+      worker.assignment.conversationId !== args.conversationId
+    ) {
+      return null;
+    }
+    const workers = await ctx.db.query("workers").collect();
     const candidates = await ctx.db
       .query("messageQueue")
       .withIndex("by_status_and_scheduledFor", (q) =>
@@ -1050,6 +1071,16 @@ export const claimNextJob = mutation({
     });
 
     for (const candidate of candidates) {
+      if (args.conversationId && candidate.conversationId !== args.conversationId) {
+        continue;
+      }
+      if (
+        worker?.assignment &&
+        candidate.conversationId !== worker.assignment.conversationId
+      ) {
+        continue;
+      }
+
       const conversation = await ctx.db
         .query("conversations")
         .withIndex("by_conversationId", (q) =>
@@ -1057,12 +1088,34 @@ export const claimNextJob = mutation({
         )
         .unique();
       if (!conversation) continue;
+      if (conversation.agentKey !== candidate.agentKey) continue;
+      if (
+        worker?.assignment &&
+        conversation.agentKey !== worker.assignment.agentKey
+      ) {
+        continue;
+      }
+      const existingOwner = findActiveAssignmentOwner(workers, {
+        conversationId: candidate.conversationId,
+        agentKey: candidate.agentKey,
+        excludeWorkerId: args.workerId,
+        staleHeartbeatCutoff,
+      });
+      if (existingOwner) {
+        continue;
+      }
 
       const lock = conversation.processingLock;
       if (lock && lock.leaseExpiresAt > nowMs) continue;
 
       const leaseId = `${nowMs}-${Math.random().toString(36).slice(2, 10)}`;
       const leaseExpiresAt = nowMs + DEFAULT_CONFIG.lease.leaseMs;
+      const nextAssignment = {
+        conversationId: candidate.conversationId,
+        agentKey: candidate.agentKey,
+        leaseId,
+        assignedAt: worker?.assignment?.assignedAt ?? nowMs,
+      };
 
       await ctx.db.patch(candidate._id, {
         status: "processing",
@@ -1091,6 +1144,7 @@ export const claimNextJob = mutation({
           lastClaimAt: nowMs,
           scheduledShutdownAt: undefined,
           stoppedAt: undefined,
+          assignment: nextAssignment,
           capabilities: [],
         });
       } else {
@@ -1101,6 +1155,7 @@ export const claimNextJob = mutation({
           lastClaimAt: nowMs,
           scheduledShutdownAt: undefined,
           stoppedAt: undefined,
+          assignment: nextAssignment,
         });
       }
 
@@ -1168,7 +1223,29 @@ export const heartbeatJob = mutation({
       .withIndex("by_workerId", (q) => q.eq("workerId", args.workerId))
       .unique();
     if (worker && isWorkerRunning(worker.status)) {
-      await ctx.db.patch(worker._id, { heartbeatAt: nowMs });
+      const nextPatch: {
+        heartbeatAt: number;
+        assignment?: {
+          conversationId: string;
+          agentKey: string;
+          leaseId: string;
+          assignedAt: number;
+        };
+      } = { heartbeatAt: nowMs };
+      if (
+        !worker.assignment ||
+        worker.assignment.conversationId !== message.conversationId ||
+        worker.assignment.agentKey !== message.agentKey ||
+        worker.assignment.leaseId !== args.leaseId
+      ) {
+        nextPatch.assignment = {
+          conversationId: message.conversationId,
+          agentKey: message.agentKey,
+          leaseId: args.leaseId,
+          assignedAt: worker.assignment?.assignedAt ?? nowMs,
+        };
+      }
+      await ctx.db.patch(worker._id, nextPatch);
     }
 
     return true;
@@ -1229,6 +1306,7 @@ export const completeJob = mutation({
         load: nextLoad,
         heartbeatAt: nowMs,
         scheduledShutdownAt: nextScheduledShutdownAt,
+        assignment: getAssignmentForCompletedConversation(worker, message),
       });
       if (nextScheduledShutdownAt !== undefined) {
         await scheduleIdleShutdownWatchdog(ctx, nextScheduledShutdownAt, nowMs, args.providerConfig);
@@ -1317,6 +1395,7 @@ export const failJob = mutation({
         load: nextLoad,
         heartbeatAt: nowMs,
         scheduledShutdownAt: nextScheduledShutdownAt,
+        assignment: getAssignmentForCompletedConversation(worker, message),
       });
       if (nextScheduledShutdownAt !== undefined) {
         await scheduleIdleShutdownWatchdog(ctx, nextScheduledShutdownAt, nowMs, args.providerConfig);
@@ -1402,6 +1481,7 @@ export const releaseExpiredLeases = internalMutation({
             load: nextLoad,
             heartbeatAt: nowMs,
             scheduledShutdownAt: nextScheduledShutdownAt,
+            assignment: clearAssignmentForMessage(worker, message, nextLoad),
           });
           if (nextScheduledShutdownAt !== undefined) {
             await scheduleIdleShutdownWatchdog(ctx, nextScheduledShutdownAt, nowMs);
@@ -1485,6 +1565,7 @@ export const releaseStuckJobs = mutation({
             load: nextLoad,
             heartbeatAt: nowMs,
             scheduledShutdownAt: nextScheduledShutdownAt,
+            assignment: clearAssignmentForMessage(worker, message, nextLoad),
           });
           if (nextScheduledShutdownAt !== undefined) {
             await scheduleIdleShutdownWatchdog(ctx, nextScheduledShutdownAt, nowMs);
@@ -1714,6 +1795,33 @@ export const getActiveConversationCountForScheduler = internalQuery({
   },
 });
 
+export const getActiveConversationIdsForScheduler = internalQuery({
+  args: {
+    nowMs: v.optional(v.number()),
+    limit: v.optional(v.number()),
+  },
+  returns: v.array(v.string()),
+  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);
+    return Array.from(
+      new Set([...queuedJobs, ...processingJobs].map((job) => job.conversationId)),
+    ).sort();
+  },
+});
+
 export const listJobsByStatus = query({
   args: {
     status: queueStatusValidator,
@@ -1788,6 +1896,7 @@ export const upsertWorkerState = internalMutation({
           args.status === "stopped" || args.status === "stopping"
             ? (args.stoppedAt ?? nowMs)
             : undefined,
+        assignment: undefined,
         machineRef:
           args.machineId && args.appName
             ? {
@@ -1817,6 +1926,7 @@ export const upsertWorkerState = internalMutation({
           ? (args.stoppedAt ?? worker.stoppedAt ?? nowMs)
           : undefined,
       lastSnapshotId: args.clearLastSnapshotId ? undefined : worker.lastSnapshotId,
+      assignment: worker.assignment,
       machineRef:
         args.clearMachineRef
           ? undefined
@@ -1979,10 +2089,9 @@ export const getLatestDataSnapshotForRestore = query({
         snapshot.archiveFileId !== undefined &&
         snapshot.expiresAt > nowMs,
     );
-    const preferred =
-      (args.conversationId
-        ? ready.find((snapshot) => snapshot.conversationId === args.conversationId)
-        : undefined) ?? ready[0];
+    const preferred = args.conversationId
+      ? ready.find((snapshot) => snapshot.conversationId === args.conversationId)
+      : ready[0];
     if (!preferred || !preferred.archiveFileId) return null;
     const downloadUrl = await ctx.storage.getUrl(preferred.archiveFileId);
     if (!downloadUrl) return null;
@@ -2008,6 +2117,7 @@ export const listWorkersForScheduler = internalQuery({
       scheduledShutdownAt: v.union(v.null(), v.number()),
       stoppedAt: v.union(v.null(), v.number()),
       lastSnapshotId: v.union(v.null(), v.id("dataSnapshots")),
+      assignment: v.union(v.null(), workerAssignmentValidator),
       machineId: v.union(v.null(), v.string()),
       appName: v.union(v.null(), v.string()),
       region: v.union(v.null(), v.string()),
@@ -2024,6 +2134,7 @@ export const listWorkersForScheduler = internalQuery({
       scheduledShutdownAt: worker.scheduledShutdownAt ?? null,
       stoppedAt: worker.stoppedAt ?? null,
       lastSnapshotId: worker.lastSnapshotId ?? null,
+      assignment: worker.assignment ?? null,
       machineId: worker.machineRef?.machineId ?? null,
       appName: worker.machineRef?.appName ?? null,
       region: worker.machineRef?.region ?? null,
@@ -2192,7 +2303,7 @@ async function scheduleIdleShutdownWatchdog(
   }
 }
 
-async function scheduleLeaseRecoveryWatchdog(ctx: any, nowMs: number) {
+async function scheduleLeaseRecoveryWatchdog(ctx: any, _nowMs: number) {
   const delayMs = DEFAULT_CONFIG.lease.leaseMs + 1_000;
   try {
     await ctx.scheduler.runAfter(delayMs, (internal.scheduler as any).reconcileWorkerPoolInternal, {
@@ -2222,6 +2333,89 @@ function computeNextScheduledShutdownAt(
   return worker.scheduledShutdownAt ?? shutdownBaseMs + DEFAULT_CONFIG.scaling.idleTimeoutMs;
 }
 
+function getAssignmentForCompletedConversation(
+  worker: {
+    assignment?: {
+      conversationId: string;
+      agentKey: string;
+      leaseId: string;
+      assignedAt: number;
+    };
+  },
+  message: {
+    conversationId: string;
+    agentKey: string;
+    leaseId?: string;
+  },
+) {
+  if (
+    worker.assignment &&
+    worker.assignment.conversationId === message.conversationId &&
+    worker.assignment.agentKey === message.agentKey
+  ) {
+    return {
+      ...worker.assignment,
+      leaseId: message.leaseId ?? worker.assignment.leaseId,
+    };
+  }
+  return worker.assignment;
+}
+
+function clearAssignmentForMessage(
+  worker: {
+    assignment?: {
+      conversationId: string;
+      agentKey: string;
+      leaseId: string;
+      assignedAt: number;
+    };
+  },
+  message: {
+    conversationId: string;
+    agentKey: string;
+  },
+  nextLoad: number,
+) {
+  if (
+    nextLoad === 0 &&
+    worker.assignment &&
+    worker.assignment.conversationId === message.conversationId &&
+    worker.assignment.agentKey === message.agentKey
+  ) {
+    return undefined;
+  }
+  return worker.assignment;
+}
+
+function findActiveAssignmentOwner(
+  workers: Array<{
+    workerId: string;
+    status: "active" | "draining" | "stopping" | "stopped";
+    heartbeatAt: number;
+    assignment?: {
+      conversationId: string;
+      agentKey: string;
+      leaseId: string;
+      assignedAt: number;
+    };
+  }>,
+  args: {
+    conversationId: string;
+    agentKey: string;
+    excludeWorkerId: string;
+    staleHeartbeatCutoff: number;
+  },
+) {
+  return workers.find(
+    (candidate) =>
+      candidate.workerId !== args.excludeWorkerId &&
+      isWorkerClaimable(candidate.status) &&
+      candidate.heartbeatAt > args.staleHeartbeatCutoff &&
+      candidate.assignment?.conversationId === args.conversationId &&
+      candidate.assignment.agentKey === args.agentKey,
+  );
+}
+
 function dedupeMessagesById<T extends { _id: string }>(messages: Array<T>): Array<T> {
   const seen = new Set<string>();
   const deduped: Array<T> = [];

package/src/component/scheduler.ts

@@ -64,6 +64,12 @@ type SchedulerWorkerRow = {
   scheduledShutdownAt: number | null;
   stoppedAt: number | null;
   lastSnapshotId: string | null;
+  assignment: {
+    conversationId: string;
+    agentKey: string;
+    leaseId: string;
+    assignedAt: number;
+  } | null;
   machineId: string | null;
   appName: string | null;
   region: string | null;
@@ -174,10 +180,11 @@ async function runReconcileWorkerPool(
   }
   const workspaceId = args.workspaceId ?? "default";
   const provider = resolveProvider(providerConfig.kind, flyApiToken);
-  const activeConversationCount: number = await ctx.runQuery(
-    (internal.queue as any).getActiveConversationCountForScheduler,
+  const activeConversationIds: Array<string> = await ctx.runQuery(
+    (internal.queue as any).getActiveConversationIdsForScheduler,
     { nowMs, limit: 1000 },
   );
+  const activeConversationCount = activeConversationIds.length;
   const cycle = await runWorkerLifecycleCycle(ctx, {
     nowMs,
     provider,
@@ -186,6 +193,7 @@ async function runReconcileWorkerPool(
     allowSpawn: true,
     convexUrl,
     workspaceId,
+    activeConversationIds,
     desiredActiveWorkers: clamp(activeConversationCount, 0, scaling.maxWorkers),
   });
   if (activeConversationCount > 0 || cycle.pending > 0) {
@@ -244,6 +252,7 @@ async function runEnforceIdleShutdowns(
     scaling: DEFAULT_CONFIG.scaling,
     allowSpawn: false,
     desiredActiveWorkers: 0,
+    activeConversationIds: [],
   });
 
   if (cycle.pending > 0) {
@@ -267,6 +276,7 @@ async function runWorkerLifecycleCycle(
     scaling: typeof DEFAULT_CONFIG.scaling;
     allowSpawn: boolean;
     desiredActiveWorkers: number;
+    activeConversationIds: Array<string>;
     convexUrl?: string;
     workspaceId?: string;
   },
@@ -320,9 +330,11 @@ async function runWorkerLifecycleCycle(
 
   let spawned = 0;
   if (input.allowSpawn && input.desiredActiveWorkers > 0) {
-    const claimableWorkers = filterScopedWorkers(workerRows, input.providerConfig.appName).filter(
-      (worker) => isWorkerClaimable(worker.status) && worker.heartbeatAt > staleHeartbeatCutoff,
-    ).length;
+    const claimableWorkers = countWorkersAvailableForActiveConversations(
+      filterScopedWorkers(workerRows, input.providerConfig.appName),
+      input.activeConversationIds,
+      staleHeartbeatCutoff,
+    );
     if (input.desiredActiveWorkers > claimableWorkers) {
       const toSpawn = Math.min(
         input.scaling.spawnStep,
@@ -747,6 +759,31 @@ function filterScopedWorkers(workerRows: Array<SchedulerWorkerRow>, appName: str
   return workerRows.filter((worker) => worker.appName === null || worker.appName === appName);
 }
 
+function countWorkersAvailableForActiveConversations(
+  workerRows: Array<SchedulerWorkerRow>,
+  activeConversationIds: Array<string>,
+  staleHeartbeatCutoff: number,
+) {
+  const activeConversationSet = new Set(activeConversationIds);
+  const assignedConversationKeys = new Set<string>();
+  let unassignedWorkers = 0;
+  for (const worker of workerRows) {
+    if (!isWorkerClaimable(worker.status) || worker.heartbeatAt <= staleHeartbeatCutoff) {
+      continue;
+    }
+    if (!worker.assignment) {
+      unassignedWorkers += 1;
+      continue;
+    }
+    if (activeConversationSet.has(worker.assignment.conversationId)) {
+      assignedConversationKeys.add(
+        `${worker.assignment.agentKey}::${worker.assignment.conversationId}`,
+      );
+    }
+  }
+  return unassignedWorkers + assignedConversationKeys.size;
+}
+
 function deriveScheduledShutdownAt(
   worker: SchedulerWorkerRow,
   nowMs: number,

package/src/component/schema.ts

@@ -119,6 +119,14 @@ export default defineSchema({
     scheduledShutdownAt: v.optional(v.number()),
     stoppedAt: v.optional(v.number()),
     lastSnapshotId: v.optional(v.id("dataSnapshots")),
+    assignment: v.optional(
+      v.object({
+        conversationId: v.string(),
+        agentKey: v.string(),
+        leaseId: v.string(),
+        assignedAt: v.number(),
+      }),
+    ),
     capabilities: v.array(v.string()),
   })
     .index("by_workerId", ["workerId"])