@okrlinkhub/agent-factory 1.0.0 → 1.0.2

package/dist/component/queue.d.ts

@@ -2,10 +2,10 @@ export declare const enqueueMessage: import("convex/server").RegisteredMutation<
     nowMs?: number | undefined;
     providerConfig?: {
         appName: string;
-        region: string;
         kind: "fly" | "runpod" | "ecs";
         organizationSlug: string;
         image: string;
+        region: string;
         volumeName: string;
         volumePath: string;
         volumeSizeGb: number;
@@ -72,10 +72,10 @@ export declare const getActiveSecretPlaintext: import("convex/server").Registere
 }, Promise<string | null>>;
 export declare const getProviderRuntimeConfig: import("convex/server").RegisteredQuery<"internal", {}, Promise<{
     appName: string;
-    region: string;
     kind: "fly" | "runpod" | "ecs";
     organizationSlug: string;
     image: string;
+    region: string;
     volumeName: string;
     volumePath: string;
     volumeSizeGb: number;
@@ -84,10 +84,10 @@ export declare const upsertProviderRuntimeConfig: import("convex/server").Regist
     nowMs?: number | undefined;
     providerConfig: {
         appName: string;
-        region: string;
         kind: "fly" | "runpod" | "ecs";
         organizationSlug: string;
         image: string;
+        region: string;
         volumeName: string;
         volumePath: string;
         volumeSizeGb: number;
@@ -95,10 +95,10 @@ export declare const upsertProviderRuntimeConfig: import("convex/server").Regist
 }, Promise<null>>;
 export declare const providerRuntimeConfig: import("convex/server").RegisteredQuery<"public", {}, Promise<{
     appName: string;
-    region: string;
     kind: "fly" | "runpod" | "ecs";
     organizationSlug: string;
     image: string;
+    region: string;
     volumeName: string;
     volumePath: string;
     volumeSizeGb: number;
@@ -107,10 +107,10 @@ export declare const setProviderRuntimeConfig: import("convex/server").Registere
     nowMs?: number | undefined;
     providerConfig: {
         appName: string;
-        region: string;
         kind: "fly" | "runpod" | "ecs";
         organizationSlug: string;
         image: string;
+        region: string;
         volumeName: string;
         volumePath: string;
         volumeSizeGb: number;
@@ -155,9 +155,9 @@ export declare const deployGlobalSkill: import("convex/server").RegisteredMutati
     releaseChannel: "stable" | "canary";
 }>>;
 export declare const listGlobalSkills: import("convex/server").RegisteredQuery<"public", {
+    status?: "active" | "disabled" | undefined;
     limit?: number | undefined;
     releaseChannel?: "stable" | "canary" | undefined;
-    status?: "active" | "disabled" | undefined;
 }, Promise<{
     skillId: any;
     slug: string;
@@ -197,8 +197,8 @@ export declare const getWorkerGlobalSkillsManifest: import("convex/server").Regi
 export declare const setGlobalSkillStatus: import("convex/server").RegisteredMutation<"public", {
     nowMs?: number | undefined;
     actor?: string | undefined;
-    slug: string;
     status: "active" | "disabled";
+    slug: string;
 }, Promise<{
     updated: boolean;
     slug: string;
@@ -224,6 +224,7 @@ export declare const attachMessageMetadata: import("convex/server").RegisteredMu
 }, Promise<boolean>>;
 export declare const claimNextJob: import("convex/server").RegisteredMutation<"public", {
     nowMs?: number | undefined;
+    conversationId?: string | undefined;
     workerId: string;
 }, Promise<{
     messageId: import("convex/values").GenericId<"messageQueue">;
@@ -250,10 +251,10 @@ export declare const completeJob: import("convex/server").RegisteredMutation<"pu
     nowMs?: number | undefined;
     providerConfig?: {
         appName: string;
-        region: string;
         kind: "fly" | "runpod" | "ecs";
         organizationSlug: string;
         image: string;
+        region: string;
         volumeName: string;
         volumePath: string;
         volumeSizeGb: number;
@@ -266,10 +267,10 @@ export declare const failJob: import("convex/server").RegisteredMutation<"public
     nowMs?: number | undefined;
     providerConfig?: {
         appName: string;
-        region: string;
         kind: "fly" | "runpod" | "ecs";
         organizationSlug: string;
         image: string;
+        region: string;
         volumeName: string;
         volumePath: string;
         volumeSizeGb: number;
@@ -352,6 +353,10 @@ export declare const getActiveConversationCountForScheduler: import("convex/serv
     nowMs?: number | undefined;
     limit?: number | undefined;
 }, Promise<number>>;
+export declare const getActiveConversationIdsForScheduler: import("convex/server").RegisteredQuery<"internal", {
+    nowMs?: number | undefined;
+    limit?: number | undefined;
+}, Promise<string[]>>;
 export declare const listJobsByStatus: import("convex/server").RegisteredQuery<"public", {
     limit?: number | undefined;
     status: "done" | "failed" | "queued" | "processing" | "dead_letter";
@@ -370,14 +375,14 @@ export declare const listJobsByStatus: import("convex/server").RegisteredQuery<"
 export declare const upsertWorkerState: import("convex/server").RegisteredMutation<"internal", {
     nowMs?: number | undefined;
     appName?: string | undefined;
-    machineId?: string | undefined;
     region?: string | undefined;
+    machineId?: string | undefined;
     scheduledShutdownAt?: number | undefined;
     stoppedAt?: number | undefined;
     clearLastSnapshotId?: boolean | undefined;
     clearMachineRef?: boolean | undefined;
-    workerId: string;
     status: "active" | "draining" | "stopping" | "stopped";
+    workerId: string;
     provider: string;
     load: number;
 }, Promise<null>>;
@@ -408,8 +413,8 @@ export declare const finalizeDataSnapshotUpload: import("convex/server").Registe
 }, Promise<boolean>>;
 export declare const failDataSnapshotUpload: import("convex/server").RegisteredMutation<"public", {
     nowMs?: number | undefined;
-    workerId: string;
     error: string;
+    workerId: string;
     snapshotId: import("convex/values").GenericId<"dataSnapshots">;
 }, Promise<boolean>>;
 export declare const getLatestDataSnapshotForRestore: import("convex/server").RegisteredQuery<"public", {
@@ -433,6 +438,12 @@ export declare const listWorkersForScheduler: import("convex/server").Registered
     scheduledShutdownAt: number | null;
     stoppedAt: number | null;
     lastSnapshotId: import("convex/values").GenericId<"dataSnapshots"> | null;
+    assignment: {
+        agentKey: string;
+        conversationId: string;
+        leaseId: string;
+        assignedAt: number;
+    } | null;
     machineId: string | null;
     appName: string | null;
     region: string | null;

package/dist/component/queue.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"queue.d.ts","sourceRoot":"","sources":["../../src/component/queue.ts"],"names":[],"mappings":"AAwGA,eAAO,MAAM,cAAc;;;;;;;;;;;;;;;;;;;;;;;;;8DA+FzB,CAAC;AAEH,eAAO,MAAM,0BAA0B;;;;;;;;;;;;GAiErC,CAAC;AAEH,eAAO,MAAM,kBAAkB;;;;;;;;;;;;;+DAyB7B,CAAC;AAEH,eAAO,MAAM,qBAAqB;;;;;;;;GA6ChC,CAAC;AAEH,eAAO,MAAM,gBAAgB;;;eAOZ,MAAM;eACN,OAAO;aACT,MAAM,GAAG,IAAI;KAiB1B,CAAC;AAEH,eAAO,MAAM,wBAAwB;;0BAiBnC,CAAC;AAEH,eAAO,MAAM,wBAAwB;;;;;;;;;UAanC,CAAC;AAEH,eAAO,MAAM,2BAA2B;;;;;;;;;;;;iBA0BtC,CAAC;AAEH,eAAO,MAAM,qBAAqB;;;;;;;;;UAahC,CAAC;AAEH,eAAO,MAAM,wBAAwB;;;;;;;;;;;;iBA0BnC,CAAC;AAEH,eAAO,MAAM,uBAAuB;;UAalC,CAAC;AAEH,eAAO,MAAM,0BAA0B;;;;;iBAiCrC,CAAC;AAEH,eAAO,MAAM,oBAAoB;;UAa/B,CAAC;AAEH,eAAO,MAAM,uBAAuB;;;;;iBAiClC,CAAC;AAEH,eAAO,MAAM,iBAAiB;;;;;;;;;;;;;;;;;;;GAgI5B,CAAC;AAEH,eAAO,MAAM,gBAAgB;;;;;aA0Cd,GAAG;UACN,MAAM;iBACC,MAAM;kBACL,MAAM;YACZ,QAAQ,GAAG,UAAU;eAClB,MAAM;mBACF;QACb,SAAS,EAAE,GAAG,CAAC;QACf,SAAS,EAAE,GAAG,CAAC;QACf,OAAO,EAAE,MAAM,CAAC;QAChB,MAAM,EAAE,MAAM,CAAC;QACf,YAAY,EAAE,KAAK,GAAG,KAAK,CAAC;QAC5B,UAAU,EAAE,MAAM,CAAC;QACnB,cAAc,EAAE,QAAQ,GAAG,QAAQ,CAAC;QACpC,WAAW,EAAE,MAAM,CAAC;KACrB,GAAG,IAAI;KAwCZ,CAAC;AAEH,eAAO,MAAM,6BAA6B;;;;;;;;;;cAuB9B,MAAM;iBACH,MAAM;sBACD,KAAK,GAAG,KAAK;oBACf,MAAM;kBACR,MAAM;gBACR,MAAM;;GA0ClB,CAAC;AAEH,eAAO,MAAM,oBAAoB;;;;;;;;;GA8B/B,CAAC;AAEH,eAAO,MAAM,iBAAiB;;;;;;;GAsE5B,CAAC;AAEH,eAAO,MAAM,sBAAsB;;GASjC,CAAC;AAEH,eAAO,MAAM,iBAAiB;;0BAQ5B,CAAC;AAEH,eAAO,MAAM,qBAAqB;;;oBAoBhC,CAAC;AAEH,eAAO,MAAM,YAAY;;;;;;;;;;;;;;;;;UAgGvB,CAAC;AAEH,eAAO,MAAM,YAAY;;;;;oBAsDvB,CAAC;AAEH,eAAO,MAAM,WAAW;;;;;;;;;;;;;;;oBA6DtB,CAAC;AAEH,eAAO,MAAM,OAAO;;;;;;;;;;;;;;;;;;;;GA2FlB,CAAC;AAEH,eAAO,MAAM,oBAAoB;;;;;;GAiF/B,CAAC;AAEH,eAAO,MAAM,gBAAgB;;;;;;GAiF3B,CAAC;AAEH,eAAO,MAAM,+BAA+B;;;;;;;;;;;;;;;;;;;;;;;;;;;;;iBAkmBjC,MAAM,GAAG,IAAI;2BACH,MAAM,GAAG,IAAI;mBACrB,MAAM,GAAG,IAAI;gBAChB,MAAM,GAAG,IAAI;oBACT,MAAM,GAAG,IAAI;6BACJ,MAAM,GAAG,IAAI;;UAngBlC,CAAC;AAEH,eAAO,MAAM,aAAa;;;;;;GAgCxB,CAAC;AAEH,eAAO,MAAM,4BAA4B;;oBAevC,CAAC;AAEH,eAAO,MAAM,qCAAqC;;;mBA6BhD,CAAC;AAEH,eAAO,MAAM,sCAAsC;;;mBA+BjD,CAAC;AAEH,eAAO,MAAM,gBAAgB;;;;;;;;;;;;;;KAsC3B,CAAC;AAEH,eAAO,MAAM,iBAAiB;;;;;;;;;;;;;iBA4E5B,CAAC;AAEH,eAAO,MAAM,qBAAqB;;;;GAsBhC,CAAC;AAEH,eAAO,MAAM,yBAAyB;;;;;;;;;;;GA+BpC,CAAC;AAEH,eAAO,MAAM,0BAA0B;;;;;;;oBAmCrC,CAAC;AAEH,eAAO,MAAM,sBAAsB;;;;;oBAmBjC,CAAC;AAEH,eAAO,MAAM,+BAA+B;;;;;;;;;;;UA+C1C,CAAC;AAEH,eAAO,MAAM,uBAAuB;;;;;;;;;;;;KAiClC,CAAC;AAEH,eAAO,MAAM,sBAAsB;;;mBAoBjC,CAAC;AAEH,eAAO,MAAM,cAAc;;;;;;;;;;;GAsCzB,CAAC"}
+{"version":3,"file":"queue.d.ts","sourceRoot":"","sources":["../../src/component/queue.ts"],"names":[],"mappings":"AA+GA,eAAO,MAAM,cAAc;;;;;;;;;;;;;;;;;;;;;;;;;8DAmGzB,CAAC;AAEH,eAAO,MAAM,0BAA0B;;;;;;;;;;;;GAiErC,CAAC;AAEH,eAAO,MAAM,kBAAkB;;;;;;;;;;;;;+DAyB7B,CAAC;AAEH,eAAO,MAAM,qBAAqB;;;;;;;;GA6ChC,CAAC;AAEH,eAAO,MAAM,gBAAgB;;;eAOZ,MAAM;eACN,OAAO;aACT,MAAM,GAAG,IAAI;KAiB1B,CAAC;AAEH,eAAO,MAAM,wBAAwB;;0BAiBnC,CAAC;AAEH,eAAO,MAAM,wBAAwB;;;;;;;;;UAanC,CAAC;AAEH,eAAO,MAAM,2BAA2B;;;;;;;;;;;;iBA0BtC,CAAC;AAEH,eAAO,MAAM,qBAAqB;;;;;;;;;UAahC,CAAC;AAEH,eAAO,MAAM,wBAAwB;;;;;;;;;;;;iBA0BnC,CAAC;AAEH,eAAO,MAAM,uBAAuB;;UAalC,CAAC;AAEH,eAAO,MAAM,0BAA0B;;;;;iBAiCrC,CAAC;AAEH,eAAO,MAAM,oBAAoB;;UAa/B,CAAC;AAEH,eAAO,MAAM,uBAAuB;;;;;iBAiClC,CAAC;AAEH,eAAO,MAAM,iBAAiB;;;;;;;;;;;;;;;;;;;GAgI5B,CAAC;AAEH,eAAO,MAAM,gBAAgB;;;;;aA0Cd,GAAG;UACN,MAAM;iBACC,MAAM;kBACL,MAAM;YACZ,QAAQ,GAAG,UAAU;eAClB,MAAM;mBACF;QACb,SAAS,EAAE,GAAG,CAAC;QACf,SAAS,EAAE,GAAG,CAAC;QACf,OAAO,EAAE,MAAM,CAAC;QAChB,MAAM,EAAE,MAAM,CAAC;QACf,YAAY,EAAE,KAAK,GAAG,KAAK,CAAC;QAC5B,UAAU,EAAE,MAAM,CAAC;QACnB,cAAc,EAAE,QAAQ,GAAG,QAAQ,CAAC;QACpC,WAAW,EAAE,MAAM,CAAC;KACrB,GAAG,IAAI;KAwCZ,CAAC;AAEH,eAAO,MAAM,6BAA6B;;;;;;;;;;cAuB9B,MAAM;iBACH,MAAM;sBACD,KAAK,GAAG,KAAK;oBACf,MAAM;kBACR,MAAM;gBACR,MAAM;;GA0ClB,CAAC;AAEH,eAAO,MAAM,oBAAoB;;;;;;;;;GA8B/B,CAAC;AAEH,eAAO,MAAM,iBAAiB;;;;;;;GAsE5B,CAAC;AAEH,eAAO,MAAM,sBAAsB;;GASjC,CAAC;AAEH,eAAO,MAAM,iBAAiB;;0BAQ5B,CAAC;AAEH,eAAO,MAAM,qBAAqB;;;oBAoBhC,CAAC;AAEH,eAAO,MAAM,YAAY;;;;;;;;;;;;;;;;;;UA4IvB,CAAC;AAEH,eAAO,MAAM,YAAY;;;;;oBA4EvB,CAAC;AAEH,eAAO,MAAM,WAAW;;;;;;;;;;;;;;;oBA8DtB,CAAC;AAEH,eAAO,MAAM,OAAO;;;;;;;;;;;;;;;;;;;;GA4FlB,CAAC;AAEH,eAAO,MAAM,oBAAoB;;;;;;GAkF/B,CAAC;AAEH,eAAO,MAAM,gBAAgB;;;;;;GAkF3B,CAAC;AAEH,eAAO,MAAM,+BAA+B;;;;;;;;;;;;;;;;;;;;;;;;;;;;;iBAgoBjC,MAAM,GAAG,IAAI;2BACH,MAAM,GAAG,IAAI;mBACrB,MAAM,GAAG,IAAI;gBAChB,MAAM,GAAG,IAAI;oBACT,MAAM,GAAG,IAAI;6BACJ,MAAM,GAAG,IAAI;;UAjiBlC,CAAC;AAEH,eAAO,MAAM,aAAa;;;;;;GAgCxB,CAAC;AAEH,eAAO,MAAM,4BAA4B;;oBAevC,CAAC;AAEH,eAAO,MAAM,qCAAqC;;;mBA6BhD,CAAC;AAEH,eAAO,MAAM,sCAAsC;;;mBA+BjD,CAAC;AAEH,eAAO,MAAM,oCAAoC;;;qBAyB/C,CAAC;AAEH,eAAO,MAAM,gBAAgB;;;;;;;;;;;;;;KAsC3B,CAAC;AAEH,eAAO,MAAM,iBAAiB;;;;;;;;;;;;;iBA8E5B,CAAC;AAEH,eAAO,MAAM,qBAAqB;;;;GAsBhC,CAAC;AAEH,eAAO,MAAM,yBAAyB;;;;;;;;;;;GA+BpC,CAAC;AAEH,eAAO,MAAM,0BAA0B;;;;;;;oBAmCrC,CAAC;AAEH,eAAO,MAAM,sBAAsB;;;;;oBAmBjC,CAAC;AAEH,eAAO,MAAM,+BAA+B;;;;;;;;;;;UA8C1C,CAAC;AAEH,eAAO,MAAM,uBAAuB;;;;;;;;;;;;;;;;;;KAmClC,CAAC;AAEH,eAAO,MAAM,sBAAsB;;;mBAoBjC,CAAC;AAEH,eAAO,MAAM,cAAc;;;;;;;;;;;GAsCzB,CAAC"}

package/dist/component/queue.js

@@ -22,6 +22,12 @@ const claimedJobValidator = v.object({
     leaseExpiresAt: v.number(),
     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(),
@@ -121,6 +127,9 @@ export const enqueueMessage = mutation({
                 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(DEFAULT_CONFIG.queue.maxPriority, Math.max(0, args.priority ?? DEFAULT_CONFIG.queue.defaultPriority));
         const messageId = await ctx.db.insert("messageQueue", {
             conversationId: args.conversationId,
@@ -864,11 +873,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))
@@ -876,6 +887,12 @@ 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) => q.eq("status", "queued").lte("scheduledFor", nowMs))
@@ -888,17 +905,45 @@ export const claimNextJob = mutation({
             return a._creationTime - b._creationTime;
         });
         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) => q.eq("conversationId", candidate.conversationId))
                 .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",
                 claimedBy: args.workerId,
@@ -924,6 +969,7 @@ export const claimNextJob = mutation({
                     lastClaimAt: nowMs,
                     scheduledShutdownAt: undefined,
                     stoppedAt: undefined,
+                    assignment: nextAssignment,
                     capabilities: [],
                 });
             }
@@ -935,6 +981,7 @@ export const claimNextJob = mutation({
                     lastClaimAt: nowMs,
                     scheduledShutdownAt: undefined,
                     stoppedAt: undefined,
+                    assignment: nextAssignment,
                 });
             }
             await scheduleLeaseRecoveryWatchdog(ctx, nowMs);
@@ -990,7 +1037,19 @@ 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: 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;
     },
@@ -1040,6 +1099,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);
@@ -1118,6 +1178,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);
@@ -1187,6 +1248,7 @@ export const releaseExpiredLeases = internalMutation({
                         load: nextLoad,
                         heartbeatAt: nowMs,
                         scheduledShutdownAt: nextScheduledShutdownAt,
+                        assignment: clearAssignmentForMessage(worker, message, nextLoad),
                     });
                     if (nextScheduledShutdownAt !== undefined) {
                         await scheduleIdleShutdownWatchdog(ctx, nextScheduledShutdownAt, nowMs);
@@ -1254,6 +1316,7 @@ export const releaseStuckJobs = mutation({
                         load: nextLoad,
                         heartbeatAt: nowMs,
                         scheduledShutdownAt: nextScheduledShutdownAt,
+                        assignment: clearAssignmentForMessage(worker, message, nextLoad),
                     });
                     if (nextScheduledShutdownAt !== undefined) {
                         await scheduleIdleShutdownWatchdog(ctx, nextScheduledShutdownAt, nowMs);
@@ -1436,6 +1499,26 @@ export const getActiveConversationCountForScheduler = internalQuery({
         return conversationIds.size;
     },
 });
+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,
@@ -1506,6 +1589,7 @@ export const upsertWorkerState = internalMutation({
                 stoppedAt: args.status === "stopped" || args.status === "stopping"
                     ? (args.stoppedAt ?? nowMs)
                     : undefined,
+                assignment: undefined,
                 machineRef: args.machineId && args.appName
                     ? {
                         appName: args.appName,
@@ -1529,6 +1613,7 @@ export const upsertWorkerState = internalMutation({
                 ? (args.stoppedAt ?? worker.stoppedAt ?? nowMs)
                 : undefined,
             lastSnapshotId: args.clearLastSnapshotId ? undefined : worker.lastSnapshotId,
+            assignment: worker.assignment,
             machineRef: args.clearMachineRef
                 ? undefined
                 : args.machineId && args.appName
@@ -1677,9 +1762,9 @@ export const getLatestDataSnapshotForRestore = query({
         const ready = candidates.filter((snapshot) => snapshot.status === "ready" &&
             snapshot.archiveFileId !== undefined &&
             snapshot.expiresAt > nowMs);
-        const preferred = (args.conversationId
+        const preferred = args.conversationId
             ? ready.find((snapshot) => snapshot.conversationId === args.conversationId)
-            : undefined) ?? ready[0];
+            : ready[0];
         if (!preferred || !preferred.archiveFileId)
             return null;
         const downloadUrl = await ctx.storage.getUrl(preferred.archiveFileId);
@@ -1705,6 +1790,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()),
@@ -1720,6 +1806,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,
@@ -1822,7 +1909,7 @@ async function scheduleIdleShutdownWatchdog(ctx, scheduledShutdownAt, nowMs, pro
         console.warn(`[queue] failed to schedule idle-shutdown watchdog: ${error instanceof Error ? error.message : String(error)}`);
     }
 }
-async function scheduleLeaseRecoveryWatchdog(ctx, nowMs) {
+async function scheduleLeaseRecoveryWatchdog(ctx, _nowMs) {
     const delayMs = DEFAULT_CONFIG.lease.leaseMs + 1_000;
     try {
         await ctx.scheduler.runAfter(delayMs, internal.scheduler.reconcileWorkerPoolInternal, {
@@ -1840,6 +1927,33 @@ function computeNextScheduledShutdownAt(worker, nextLoad, nowMs) {
     const shutdownBaseMs = worker.lastClaimAt ?? nowMs;
     return worker.scheduledShutdownAt ?? shutdownBaseMs + DEFAULT_CONFIG.scaling.idleTimeoutMs;
 }
+function getAssignmentForCompletedConversation(worker, message) {
+    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, message, nextLoad) {
+    if (nextLoad === 0 &&
+        worker.assignment &&
+        worker.assignment.conversationId === message.conversationId &&
+        worker.assignment.agentKey === message.agentKey) {
+        return undefined;
+    }
+    return worker.assignment;
+}
+function findActiveAssignmentOwner(workers, args) {
+    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(messages) {
     const seen = new Set();
     const deduped = [];