@okrlinkhub/agent-factory 1.0.1 → 1.0.2

package/src/component/queue.ts

@@ -1041,6 +1041,7 @@ export const claimNextJob = mutation({
   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))
@@ -1055,6 +1056,7 @@ export const claimNextJob = mutation({
     ) {
       return null;
     }
+    const workers = await ctx.db.query("workers").collect();
     const candidates = await ctx.db
       .query("messageQueue")
       .withIndex("by_status_and_scheduledFor", (q) =>
@@ -1093,6 +1095,15 @@ export const claimNextJob = mutation({
       ) {
         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;
@@ -2292,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, {
@@ -2376,6 +2387,35 @@ function clearAssignmentForMessage(
   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

@@ -765,20 +765,23 @@ function countWorkersAvailableForActiveConversations(
   staleHeartbeatCutoff: number,
 ) {
   const activeConversationSet = new Set(activeConversationIds);
-  let available = 0;
+  const assignedConversationKeys = new Set<string>();
+  let unassignedWorkers = 0;
   for (const worker of workerRows) {
     if (!isWorkerClaimable(worker.status) || worker.heartbeatAt <= staleHeartbeatCutoff) {
       continue;
     }
     if (!worker.assignment) {
-      available += 1;
+      unassignedWorkers += 1;
       continue;
     }
     if (activeConversationSet.has(worker.assignment.conversationId)) {
-      available += 1;
+      assignedConversationKeys.add(
+        `${worker.assignment.agentKey}::${worker.assignment.conversationId}`,
+      );
     }
   }
-  return available;
+  return unassignedWorkers + assignedConversationKeys.size;
 }
 
 function deriveScheduledShutdownAt(