@okrlinkhub/agent-factory 3.0.0 → 3.0.2

package/src/component/queue.ts

@@ -188,6 +188,11 @@ const workerSpawnOpenClawEnvValidator = v.object({
   OPENCLAW_LINKING_SHARED_SECRET: v.optional(v.string()),
 });
 
+const schedulerConversationTargetValidator = v.object({
+  conversationId: v.string(),
+  agentKey: v.string(),
+});
+
 const messageRuntimeConfigValidator = v.object({
   systemPrompt: v.optional(v.string()),
   telegramAttachmentRetentionMs: v.optional(v.number()),
@@ -489,6 +494,47 @@ export const getWorkerSpawnOpenClawEnv = internalQuery({
   },
 });
 
+export const getActiveConversationsForScheduler = internalQuery({
+  args: {
+    nowMs: v.optional(v.number()),
+    limit: v.optional(v.number()),
+  },
+  returns: v.array(schedulerConversationTargetValidator),
+  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);
+
+    const conversations = new Map<string, { conversationId: string; agentKey: string }>();
+    for (const job of [...queuedJobs, ...processingJobs]) {
+      const key = `${job.agentKey}::${job.conversationId}`;
+      if (!conversations.has(key)) {
+        conversations.set(key, {
+          conversationId: job.conversationId,
+          agentKey: job.agentKey,
+        });
+      }
+    }
+
+    return Array.from(conversations.values()).sort(
+      (left, right) =>
+        left.agentKey.localeCompare(right.agentKey) ||
+        left.conversationId.localeCompare(right.conversationId),
+    );
+  },
+});
+
 export const getProviderRuntimeConfig = internalQuery({
   args: {},
   returns: v.union(v.null(), providerConfigValidator),
@@ -2013,6 +2059,67 @@ export const sendMessageToUserAgent = mutation({
   },
 });
 
+export const sendMessageTemplateToUserAgent = mutation({
+  args: {
+    consumerUserId: v.string(),
+    agentKey: v.string(),
+    templateId: v.id("messageTemplates"),
+    metadata: v.optional(v.record(v.string(), v.string())),
+    nowMs: v.optional(v.number()),
+    providerConfig: v.optional(providerConfigValidator),
+  },
+  returns: v.object({
+    messageId: v.id("messageQueue"),
+    usageCount: v.number(),
+  }),
+  handler: async (ctx, args) => {
+    const template = await ctx.db.get(args.templateId);
+    if (!template || !template.enabled) {
+      throw new Error("Message template not found");
+    }
+
+    const nowMs = args.nowMs ?? Date.now();
+    const target = await resolveConversationTargetForUserAgent(
+      ctx,
+      args.consumerUserId,
+      args.agentKey,
+      true,
+    );
+    const providerUserId =
+      target.telegramUserId ?? target.telegramChatId ?? args.consumerUserId;
+    const usageCount = template.usageCount + 1;
+    const messageId = await enqueueMessageRecord(ctx, {
+      conversationId: target.conversationId,
+      agentKey: args.agentKey,
+      payload: {
+        provider: target.provider,
+        providerUserId,
+        messageText: template.text,
+        metadata: {
+          ...(args.metadata ?? {}),
+          consumerUserId: args.consumerUserId,
+          source: "message_template",
+          templateId: String(template._id),
+          templateKey: template.templateKey,
+          ...(target.telegramChatId ? { telegramChatId: target.telegramChatId } : {}),
+          ...(target.telegramUserId ? { telegramUserId: target.telegramUserId } : {}),
+        },
+      },
+      scheduledFor: nowMs,
+      providerConfig: args.providerConfig,
+    });
+
+    await ctx.db.patch(template._id, {
+      usageCount,
+    });
+
+    return {
+      messageId,
+      usageCount,
+    };
+  },
+});
+
 export const listSnapshotsForConversation = query({
   args: {
     conversationId: v.string(),
@@ -2216,6 +2323,7 @@ export const upsertWorkerState = internalMutation({
     machineId: v.optional(v.string()),
     appName: v.optional(v.string()),
     region: v.optional(v.string()),
+    assignment: v.optional(v.union(v.null(), workerAssignmentValidator)),
     clearLastSnapshotId: v.optional(v.boolean()),
     clearMachineRef: v.optional(v.boolean()),
   },
@@ -2238,7 +2346,7 @@ export const upsertWorkerState = internalMutation({
           args.status === "stopped" || args.status === "stopping"
             ? (args.stoppedAt ?? nowMs)
             : undefined,
-        assignment: undefined,
+        assignment: args.assignment ?? undefined,
         machineRef:
           args.machineId && args.appName
             ? {
@@ -2268,7 +2376,7 @@ export const upsertWorkerState = internalMutation({
           ? (args.stoppedAt ?? worker.stoppedAt ?? nowMs)
           : undefined,
       lastSnapshotId: args.clearLastSnapshotId ? undefined : worker.lastSnapshotId,
-      assignment: worker.assignment,
+      assignment: args.assignment === undefined ? worker.assignment : (args.assignment ?? undefined),
       machineRef:
         args.clearMachineRef
           ? undefined

package/src/component/scheduler.ts

@@ -75,6 +75,11 @@ type SchedulerWorkerRow = {
   region: string | null;
 };
 
+type SchedulerConversationTarget = {
+  conversationId: string;
+  agentKey: string;
+};
+
 const PROVIDER_RECONCILE_GRACE_MS = 90_000;
 
 export const reconcileWorkerPool = action({
@@ -182,11 +187,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 +200,7 @@ async function runReconcileWorkerPool(
     allowSpawn: true,
     convexUrl,
     workspaceId,
-    activeConversationIds,
+    activeConversations,
     desiredActiveWorkers: clamp(activeConversationCount, 0, scaling.maxWorkers),
   });
   if (activeConversationCount > 0 || cycle.pending > 0) {
@@ -258,7 +263,7 @@ async function runEnforceIdleShutdowns(
     scaling: DEFAULT_CONFIG.scaling,
     allowSpawn: false,
     desiredActiveWorkers: 0,
-    activeConversationIds: [],
+    activeConversations: [],
   });
 
   if (cycle.pending > 0) {
@@ -282,7 +287,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 +343,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 +355,23 @@ 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;
         await ctx.runMutation(internal.queue.upsertWorkerState, {
           workerId,
           provider: input.providerConfig.kind,
@@ -359,6 +379,7 @@ async function runWorkerLifecycleCycle(
           load: 0,
           nowMs: input.nowMs,
           scheduledShutdownAt: input.nowMs + input.scaling.idleTimeoutMs,
+          assignment,
         });
         let created;
         try {
@@ -377,6 +398,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) {
@@ -456,6 +479,7 @@ async function runWorkerLifecycleCycle(
           machineId: created.machineId,
           appName: input.providerConfig.appName,
           region: created.region,
+          assignment,
         });
         await scheduleIdleShutdownWatch(
           ctx,
@@ -862,10 +886,46 @@ function filterScopedWorkers(workerRows: Array<SchedulerWorkerRow>, appName: str
 
 function countWorkersAvailableForActiveConversations(
   workerRows: Array<SchedulerWorkerRow>,
-  activeConversationIds: Array<string>,
+  activeConversations: Array<SchedulerConversationTarget>,
   staleHeartbeatCutoff: number,
 ) {
-  const activeConversationSet = new Set(activeConversationIds);
+  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 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 +936,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

@@ -410,6 +410,21 @@ export default defineSchema({
     .index("by_companyId_and_templateKey", ["companyId", "templateKey"])
     .index("by_companyId_and_enabled", ["companyId", "enabled"]),
 
+  messageTemplates: defineTable({
+    templateKey: v.string(),
+    title: v.string(),
+    text: v.string(),
+    tags: v.array(v.string()),
+    usageCount: v.number(),
+    enabled: v.boolean(),
+    createdBy: v.string(),
+    updatedBy: v.string(),
+    createdAt: v.number(),
+    updatedAt: v.number(),
+  })
+    .index("by_templateKey", ["templateKey"])
+    .index("by_enabled", ["enabled"]),
+
   messagePushJobs: defineTable({
     companyId: v.string(),
     consumerUserId: v.string(),