@microfox/ai-worker 1.0.3 → 1.0.5

package/dist/handler.js

@@ -21,6 +21,7 @@ var __toCommonJS = (mod) => __copyProps(__defProp({}, "__esModule", { value: tru
 var handler_exports = {};
 __export(handler_exports, {
   SQS_MAX_DELAY_SECONDS: () => SQS_MAX_DELAY_SECONDS,
+  TokenBudgetExceededError: () => TokenBudgetExceededError,
   createLambdaHandler: () => createLambdaHandler,
   createWorkerLogger: () => createWorkerLogger,
   wrapHandlerForQueue: () => wrapHandlerForQueue
@@ -68,7 +69,7 @@ async function getJobById(jobId) {
     return null;
   }
 }
-function createMongoJobStore(workerId, jobId, input, metadata) {
+function createMongoJobStore(workerId, jobId, input, metadata, userId) {
   return {
     update: async (update) => {
       try {
@@ -107,6 +108,7 @@ function createMongoJobStore(workerId, jobId, input, metadata) {
             output: update.output,
             error: update.error,
             metadata: metadataUpdate,
+            ...userId ? { userId } : {},
             createdAt: now,
             updatedAt: now,
             completedAt: set.completedAt
@@ -172,7 +174,7 @@ function createMongoJobStore(workerId, jobId, input, metadata) {
     }
   };
 }
-async function upsertJob(jobId, workerId, input, metadata) {
+async function upsertJob(jobId, workerId, input, metadata, userId) {
   const coll = await getCollection();
   const now = (/* @__PURE__ */ new Date()).toISOString();
   await coll.updateOne(
@@ -185,6 +187,7 @@ async function upsertJob(jobId, workerId, input, metadata) {
         status: "queued",
         input: input ?? {},
         metadata: metadata ?? {},
+        ...userId ? { userId } : {},
         createdAt: now,
         updatedAt: now
       }
@@ -263,13 +266,14 @@ async function loadJob(jobId) {
     error: parseJson(data.error),
     metadata: parseJson(data.metadata) ?? {},
     internalJobs,
+    ...data.userId ? { userId: data.userId } : {},
     createdAt: data.createdAt,
     updatedAt: data.updatedAt,
     completedAt: data.completedAt
   };
   return record;
 }
-function createRedisJobStore(workerId, jobId, input, metadata) {
+function createRedisJobStore(workerId, jobId, input, metadata, userId) {
   return {
     update: async (update) => {
       const redis = getRedis();
@@ -327,28 +331,20 @@ function createRedisJobStore(workerId, jobId, input, metadata) {
     }
   };
 }
-async function upsertRedisJob(jobId, workerId, input, metadata) {
+async function upsertRedisJob(jobId, workerId, input, metadata, userId) {
   const redis = getRedis();
   const key = jobKey(jobId);
   const now = (/* @__PURE__ */ new Date()).toISOString();
-  const doc = {
-    jobId,
-    workerId,
-    status: "queued",
-    input,
-    metadata,
-    createdAt: now,
-    updatedAt: now
-  };
   const toSet = {
     jobId,
     workerId,
-    status: doc.status,
-    input: JSON.stringify(doc.input ?? {}),
-    metadata: JSON.stringify(doc.metadata ?? {}),
+    status: "queued",
+    input: JSON.stringify(input ?? {}),
+    metadata: JSON.stringify(metadata ?? {}),
     createdAt: now,
     updatedAt: now
   };
+  if (userId) toSet.userId = userId;
   await redis.hset(key, toSet);
   if (jobTtlSeconds > 0) {
     await redis.expire(key, jobTtlSeconds);
@@ -440,6 +436,7 @@ async function loadQueueJobRedis(queueJobId) {
     status: String(d.status ?? "running"),
     steps: stepsFromHash(d.steps),
     metadata: metadataFromHash(d.metadata),
+    ...d.userId ? { userId: String(d.userId) } : {},
     createdAt: String(d.createdAt ?? (/* @__PURE__ */ new Date()).toISOString()),
     updatedAt: String(d.updatedAt ?? (/* @__PURE__ */ new Date()).toISOString()),
     completedAt: d.completedAt != null ? String(d.completedAt) : void 0
@@ -459,9 +456,8 @@ async function saveQueueJobRedis(record) {
     createdAt: record.createdAt || now,
     updatedAt: record.updatedAt || now
   };
-  if (record.completedAt) {
-    toSet.completedAt = record.completedAt;
-  }
+  if (record.completedAt) toSet.completedAt = record.completedAt;
+  if (record.userId) toSet.userId = record.userId;
   await redis.hset(key, toSet);
   if (queueJobTtlSeconds > 0) {
     await redis.expire(key, queueJobTtlSeconds);
@@ -478,7 +474,7 @@ function preferRedis() {
   return getStoreType() !== "mongodb" && Boolean((redisUrl2 || "").trim() && (redisToken2 || "").trim());
 }
 async function upsertInitialQueueJob(options) {
-  const { queueJobId, queueId, firstWorkerId, firstWorkerJobId, metadata } = options;
+  const { queueJobId, queueId, firstWorkerId, firstWorkerJobId, metadata, userId } = options;
   const now = (/* @__PURE__ */ new Date()).toISOString();
   if (preferMongo()) {
     const coll = await getMongoQueueCollection();
@@ -515,6 +511,7 @@ async function upsertInitialQueueJob(options) {
           }
         ],
         metadata: metadata ?? {},
+        ...userId ? { userId } : {},
         createdAt: now,
         updatedAt: now
       };
@@ -553,6 +550,7 @@ async function upsertInitialQueueJob(options) {
           }
         ],
         metadata: metadata ?? {},
+        ...userId ? { userId } : {},
         createdAt: now,
         updatedAt: now
       };
@@ -665,6 +663,150 @@ async function appendQueueJobStepInStore(options) {
   }
 }
 
+// src/queue.ts
+var QUEUE_ORCHESTRATION_KEYS = [
+  "__workerQueue",
+  "__hitlInput",
+  "__hitlDecision",
+  "__hitlPending",
+  "hitl"
+];
+
+// src/retryConfig.ts
+var BUILT_IN_PATTERNS = {
+  "rate-limit": {
+    match: (err) => /rate.?limit|too.?many.?requests/i.test(err.message) || err.status === 429 || err.code === 429 || err.name === "RateLimitError",
+    delayMs: (attempt) => attempt * 1e4,
+    // 10s, 20s, 30s…
+    injectContext: false
+  },
+  "json-parse": {
+    match: (err) => err.name === "SyntaxError" || err.name === "ZodError" || /json|parse|unexpected.?token|invalid.?format/i.test(err.message),
+    delayMs: (_attempt) => 0,
+    // Immediate — model self-corrects from ctx
+    injectContext: true
+  },
+  overloaded: {
+    match: (err) => /overloaded|model.?is.?busy/i.test(err.message) || err.status === 529 || err.code === 529,
+    delayMs: (attempt) => attempt * 15e3,
+    // 15s, 30s…
+    injectContext: false
+  },
+  "server-error": {
+    match: (err) => /internal.?server.?error|service.?unavailable|bad.?gateway/i.test(err.message) || typeof err.status === "number" && err.status >= 500 && err.status < 600,
+    delayMs: (attempt) => attempt * 5e3,
+    // 5s, 10s…
+    injectContext: false
+  }
+};
+function matchesRetryPattern(err, patterns, attempt) {
+  for (const pattern of patterns) {
+    if (typeof pattern === "string") {
+      const impl = BUILT_IN_PATTERNS[pattern];
+      if (impl.match(err)) {
+        return { matched: true, delayMs: impl.delayMs(attempt), injectContext: impl.injectContext };
+      }
+    } else {
+      let matched = false;
+      if (pattern.match instanceof RegExp) {
+        matched = pattern.match.test(err.message);
+      } else {
+        try {
+          matched = pattern.match(err);
+        } catch {
+          matched = false;
+        }
+      }
+      if (matched) {
+        const delayMs = typeof pattern.delayMs === "function" ? pattern.delayMs(attempt) : pattern.delayMs ?? 0;
+        return { matched: true, delayMs, injectContext: pattern.injectContext ?? false };
+      }
+    }
+  }
+  return { matched: false, delayMs: 0, injectContext: false };
+}
+async function executeWithRetry(fn, config, onRetry) {
+  const maxAttempts = config.maxAttempts ?? 3;
+  let lastError;
+  for (let attempt = 1; attempt <= maxAttempts; attempt++) {
+    const retryCtx = attempt > 1 && lastError ? {
+      attempt,
+      maxAttempts,
+      lastError: {
+        message: lastError.message,
+        name: lastError.name,
+        stack: lastError.stack,
+        code: lastError.code ?? lastError.status
+      }
+    } : void 0;
+    try {
+      return await fn(retryCtx);
+    } catch (err) {
+      lastError = err instanceof Error ? err : new Error(String(err));
+      if (err?.name === "TokenBudgetExceededError") throw err;
+      if (attempt >= maxAttempts) throw err;
+      const retryAttemptNumber = attempt;
+      const { matched, delayMs } = matchesRetryPattern(lastError, config.on, retryAttemptNumber);
+      if (!matched) throw err;
+      const nextCtx = {
+        attempt: attempt + 1,
+        maxAttempts,
+        lastError: {
+          message: lastError.message,
+          name: lastError.name,
+          stack: lastError.stack,
+          code: lastError.code ?? lastError.status
+        }
+      };
+      onRetry?.(nextCtx, delayMs);
+      if (delayMs > 0) {
+        await new Promise((r) => setTimeout(r, delayMs));
+      }
+    }
+  }
+  throw lastError ?? new Error("executeWithRetry: unknown error");
+}
+
+// src/tokenBudget.ts
+var TokenBudgetExceededError = class extends Error {
+  constructor(used, budget) {
+    super(`Token budget exceeded: used ${used} tokens (budget: ${budget})`);
+    this.name = "TokenBudgetExceededError";
+    this.used = used;
+    this.budget = budget;
+  }
+};
+function createTokenTracker(budget) {
+  let inputTokens = 0;
+  let outputTokens = 0;
+  function checkBudget() {
+    if (budget !== null) {
+      const total = inputTokens + outputTokens;
+      if (total > budget) {
+        throw new TokenBudgetExceededError(total, budget);
+      }
+    }
+  }
+  return {
+    report(usage) {
+      inputTokens += usage.inputTokens;
+      outputTokens += usage.outputTokens;
+      checkBudget();
+    },
+    getState() {
+      return { inputTokens, outputTokens, budget };
+    },
+    getBudgetInfo() {
+      const used = inputTokens + outputTokens;
+      return {
+        used,
+        budget,
+        remaining: budget !== null ? Math.max(0, budget - used) : null
+      };
+    }
+  };
+}
+
 // src/handler.ts
 var SQS_MAX_DELAY_SECONDS = 900;
 function createWorkerLogger(jobId, workerId) {
@@ -687,6 +829,67 @@ function createWorkerLogger(jobId, workerId) {
   };
 }
 var WORKER_QUEUE_KEY = "__workerQueue";
+async function loadPreviousOutputsBeforeStep(queueRuntime, queueJobId, beforeStepIndex) {
+  if (!queueJobId || typeof queueRuntime.getQueueJob !== "function") {
+    return [];
+  }
+  try {
+    const job = await queueRuntime.getQueueJob(queueJobId);
+    if (!job?.steps) return [];
+    return job.steps.slice(0, beforeStepIndex).map((s, i) => ({ stepIndex: i, workerId: s.workerId, output: s.output }));
+  } catch (e) {
+    if (process.env.AI_WORKER_QUEUES_DEBUG === "1") {
+      console.warn("[Worker] getQueueJob failed (resume mapping):", e?.message ?? e);
+    }
+    return [];
+  }
+}
+async function maybeApplyHitlResumeMapper(params, queueRuntime) {
+  const inputObj = params.input;
+  if (!inputObj || typeof inputObj !== "object") return;
+  if (!("__hitlInput" in inputObj)) return;
+  const wq = inputObj[WORKER_QUEUE_KEY];
+  if (!wq?.id || typeof wq.stepIndex !== "number") return;
+  const queueId = wq.id;
+  const stepIndex = wq.stepIndex;
+  const initialInput = wq.initialInput;
+  const queueJobId = wq.queueJobId;
+  const previousOutputs = await loadPreviousOutputsBeforeStep(queueRuntime, queueJobId, stepIndex);
+  const pendingInput = { ...inputObj };
+  for (const key of QUEUE_ORCHESTRATION_KEYS) {
+    delete pendingInput[key];
+  }
+  delete pendingInput[WORKER_QUEUE_KEY];
+  const reviewerInput = inputObj.__hitlInput;
+  const decision = inputObj.__hitlDecision;
+  let merged;
+  if (typeof queueRuntime.invokeResume === "function") {
+    merged = await queueRuntime.invokeResume(queueId, stepIndex, {
+      initialInput,
+      previousOutputs,
+      reviewerInput,
+      pendingInput
+    });
+  } else {
+    merged = {
+      ...pendingInput,
+      ...reviewerInput !== null && typeof reviewerInput === "object" ? reviewerInput : {}
+    };
+  }
+  const mergedObj = merged !== null && typeof merged === "object" ? merged : { value: merged };
+  params.input = {
+    ...mergedObj,
+    [WORKER_QUEUE_KEY]: wq,
+    ...decision !== void 0 ? { __hitlDecision: decision } : {}
+  };
+}
+function getWorkerQueueContext(input, metadata) {
+  const fromInput = input !== null && typeof input === "object" && WORKER_QUEUE_KEY in input ? input[WORKER_QUEUE_KEY] : void 0;
+  const fromMeta = metadata !== void 0 && typeof metadata === "object" && WORKER_QUEUE_KEY in metadata ? metadata[WORKER_QUEUE_KEY] : void 0;
+  const q = fromInput ?? fromMeta;
+  if (q === null || typeof q !== "object") return void 0;
+  return q;
+}
 async function notifyQueueJobStep(queueJobId, action, params) {
   try {
     if (action === "append") {
@@ -706,7 +909,7 @@ async function notifyQueueJobStep(queueJobId, action, params) {
       return;
     }
     if (params.stepIndex === void 0) return;
-    const status = action === "start" ? "running" : action === "complete" ? "completed" : action === "fail" ? "failed" : void 0;
+    const status = action === "start" ? "running" : action === "awaiting_approval" ? "awaiting_approval" : action === "complete" ? "completed" : action === "fail" ? "failed" : void 0;
     if (!status) return;
     await updateQueueJobStepInStore({
       queueJobId,
@@ -726,6 +929,13 @@ async function notifyQueueJobStep(queueJobId, action, params) {
       status
     });
   } catch (err) {
+    if (action === "append") {
+      console.error("[Worker] Queue append failed (rethrowing):", {
+        queueJobId,
+        error: err?.message ?? String(err)
+      });
+      throw err;
+    }
     console.warn("[Worker] Queue job update error:", {
       queueJobId,
       action,
@@ -735,26 +945,152 @@ async function notifyQueueJobStep(queueJobId, action, params) {
 }
 function wrapHandlerForQueue(handler, queueRuntime) {
   return async (params) => {
-    const queueContext = params.input?.[WORKER_QUEUE_KEY];
-    const output = await handler(params);
-    if (!queueContext || typeof queueContext !== "object" || !queueContext.id) {
+    await maybeApplyHitlResumeMapper(params, queueRuntime);
+    const inputObj = params.input !== null && typeof params.input === "object" ? params.input : {};
+    const queueContextRaw = inputObj[WORKER_QUEUE_KEY];
+    const queueCtxForRetry = queueContextRaw && typeof queueContextRaw === "object" ? queueContextRaw : void 0;
+    const stepRetryConfig = queueCtxForRetry?.id && typeof queueCtxForRetry.stepIndex === "number" && typeof queueRuntime.getStepAt === "function" ? queueRuntime.getStepAt(queueCtxForRetry.id, queueCtxForRetry.stepIndex)?.retry : void 0;
+    const domainInput = { ...inputObj };
+    for (const key of QUEUE_ORCHESTRATION_KEYS) {
+      delete domainInput[key];
+    }
+    delete domainInput[WORKER_QUEUE_KEY];
+    params.input = domainInput;
+    let output;
+    if (stepRetryConfig && stepRetryConfig.on.length > 0) {
+      output = await executeWithRetry(
+        async (retryCtx) => {
+          params.ctx.retryContext = retryCtx;
+          return handler(params);
+        },
+        stepRetryConfig,
+        (retryCtx, delayMs) => {
+          const logger = params.ctx.logger;
+          if (logger?.warn) {
+            logger.warn(
+              `[queue-retry] Retrying step (attempt ${retryCtx.attempt}/${retryCtx.maxAttempts}): ${retryCtx.lastError.message}`,
+              { delayMs }
+            );
+          } else {
+            console.warn("[queue-retry] Step retry", { attempt: retryCtx.attempt, error: retryCtx.lastError.message, delayMs });
+          }
+        }
+      );
+    } else {
+      output = await handler(params);
+    }
+    if (!queueContextRaw || typeof queueContextRaw !== "object") {
+      return output;
+    }
+    const queueContext = queueContextRaw;
+    if (!queueContext.id) {
       return output;
     }
     const { id: queueId, stepIndex, initialInput, queueJobId } = queueContext;
-    const jobId = params.ctx?.jobId;
-    const workerId = params.ctx?.workerId ?? "";
+    const arrayStepIndex = queueContext.arrayStepIndex ?? stepIndex;
+    const jobId = params.ctx.jobId;
+    const workerId = params.ctx.workerId ?? "";
     const next = queueRuntime.getNextStep(queueId, stepIndex);
     const childJobId = next ? `job-${Date.now()}-${Math.random().toString(36).slice(2, 11)}` : void 0;
+    const iterationCount = queueContext.iterationCount ?? 0;
+    if (typeof queueRuntime.invokeLoop === "function") {
+      const currentStep = typeof queueRuntime.getStepAt === "function" ? queueRuntime.getStepAt(queueId, stepIndex) : void 0;
+      const maxIterations = currentStep?.loop?.maxIterations ?? 50;
+      if (iterationCount < maxIterations - 1) {
+        let previousOutputsForLoop = [];
+        let stepsLengthBeforeAppend = arrayStepIndex + 1;
+        if (queueJobId && typeof queueRuntime.getQueueJob === "function") {
+          try {
+            const job = await queueRuntime.getQueueJob(queueJobId);
+            if (job?.steps) {
+              previousOutputsForLoop = job.steps.slice(0, stepIndex).map((s, i) => ({ stepIndex: i, workerId: s.workerId, output: s.output }));
+              stepsLengthBeforeAppend = job.steps.length;
+            }
+          } catch {
+          }
+        }
+        previousOutputsForLoop = previousOutputsForLoop.concat([{ stepIndex, workerId, output }]);
+        const shouldLoop = await queueRuntime.invokeLoop(queueId, stepIndex, {
+          output,
+          stepIndex,
+          iterationCount,
+          initialInput,
+          previousOutputs: previousOutputsForLoop
+        });
+        if (shouldLoop) {
+          const loopJobId = `job-${Date.now()}-${Math.random().toString(36).slice(2, 11)}`;
+          let loopInput = output;
+          if (typeof queueRuntime.invokeChain === "function") {
+            loopInput = await queueRuntime.invokeChain(queueId, stepIndex, {
+              initialInput,
+              previousOutputs: previousOutputsForLoop
+            });
+          }
+          const loopInputWithQueue = {
+            ...loopInput !== null && typeof loopInput === "object" ? loopInput : { value: loopInput },
+            [WORKER_QUEUE_KEY]: {
+              id: queueId,
+              stepIndex,
+              // definition index stays fixed
+              arrayStepIndex: stepsLengthBeforeAppend,
+              // actual index for next iteration
+              initialInput,
+              queueJobId,
+              iterationCount: iterationCount + 1
+            }
+          };
+          if (queueJobId) {
+            await notifyQueueJobStep(queueJobId, "append", { workerJobId: loopJobId, workerId });
+          }
+          if (queueJobId && typeof arrayStepIndex === "number") {
+            await notifyQueueJobStep(queueJobId, "complete", {
+              queueId,
+              stepIndex: arrayStepIndex,
+              workerJobId: jobId,
+              workerId,
+              output
+            });
+          }
+          if (currentStep?.requiresApproval && queueJobId) {
+            const hitlUiSpec = currentStep.hitl && typeof currentStep.hitl === "object" && "ui" in currentStep.hitl ? currentStep.hitl.ui : void 0;
+            const pendingInput = {
+              ...loopInputWithQueue,
+              ...hitlUiSpec !== void 0 ? { hitl: { uiSpec: hitlUiSpec } } : {},
+              __hitlPending: {
+                queueId,
+                queueJobId,
+                stepIndex,
+                workerId,
+                createdAt: (/* @__PURE__ */ new Date()).toISOString()
+              }
+            };
+            await notifyQueueJobStep(queueJobId, "awaiting_approval", {
+              queueId,
+              stepIndex: stepsLengthBeforeAppend,
+              workerJobId: loopJobId,
+              workerId,
+              input: pendingInput
+            });
+            return output;
+          }
+          await params.ctx.dispatchWorker(workerId, loopInputWithQueue, {
+            await: false,
+            jobId: loopJobId
+          });
+          return output;
+        }
+      }
+    }
     if (next && queueJobId) {
       await notifyQueueJobStep(queueJobId, "append", {
         workerJobId: childJobId,
         workerId: next.workerId
       });
     }
-    if (queueJobId && typeof stepIndex === "number") {
+    if (queueJobId && typeof arrayStepIndex === "number") {
       await notifyQueueJobStep(queueJobId, "complete", {
         queueId,
-        stepIndex,
+        stepIndex: arrayStepIndex,
         workerJobId: jobId,
         workerId,
         output
@@ -764,7 +1100,7 @@ function wrapHandlerForQueue(handler, queueRuntime) {
       return output;
     }
     let nextInput = output;
-    if (next.mapInputFromPrev && typeof queueRuntime.invokeMapInput === "function") {
+    if (typeof queueRuntime.invokeChain === "function") {
       let previousOutputs = [];
       if (queueJobId && typeof queueRuntime.getQueueJob === "function") {
         try {
@@ -772,7 +1108,7 @@ function wrapHandlerForQueue(handler, queueRuntime) {
           if (job?.steps) {
             const fromStore = job.steps.slice(0, stepIndex).map((s, i) => ({ stepIndex: i, workerId: s.workerId, output: s.output }));
             previousOutputs = fromStore.concat([
-              { stepIndex, workerId: params.ctx?.workerId ?? "", output }
+              { stepIndex, workerId: params.ctx.workerId ?? "", output }
             ]);
           }
         } catch (e) {
@@ -781,12 +1117,10 @@ function wrapHandlerForQueue(handler, queueRuntime) {
           }
         }
       }
-      nextInput = await queueRuntime.invokeMapInput(
-        queueId,
-        stepIndex + 1,
+      nextInput = await queueRuntime.invokeChain(queueId, stepIndex + 1, {
         initialInput,
         previousOutputs
-      );
+      });
     }
     const nextInputWithQueue = {
       ...nextInput !== null && typeof nextInput === "object" ? nextInput : { value: nextInput },
@@ -806,6 +1140,37 @@ function wrapHandlerForQueue(handler, queueRuntime) {
         delaySeconds: next.delaySeconds
       });
     }
+    if (next.requiresApproval && queueJobId && typeof stepIndex === "number") {
+      const hitlUiSpec = next.hitl && typeof next.hitl === "object" && "ui" in next.hitl ? next.hitl.ui : void 0;
+      const pendingInput = {
+        ...nextInputWithQueue,
+        ...hitlUiSpec !== void 0 ? { hitl: { uiSpec: hitlUiSpec } } : {},
+        __hitlPending: {
+          queueId,
+          queueJobId,
+          stepIndex: stepIndex + 1,
+          workerId: next.workerId,
+          createdAt: (/* @__PURE__ */ new Date()).toISOString()
+        }
+      };
+      await notifyQueueJobStep(queueJobId, "awaiting_approval", {
+        queueId,
+        stepIndex: stepIndex + 1,
+        workerJobId: childJobId,
+        workerId: next.workerId,
+        input: pendingInput
+      });
+      if (debug) {
+        console.log("[Worker] Queue chain paused for HITL approval:", {
+          queueId,
+          queueJobId,
+          nextStep: stepIndex + 1,
+          nextWorkerId: next.workerId,
+          pendingWorkerJobId: childJobId
+        });
+      }
+      return output;
+    }
     await params.ctx.dispatchWorker(next.workerId, nextInputWithQueue, {
       await: false,
       delaySeconds: next.delaySeconds,
@@ -829,6 +1194,7 @@ function createDispatchWorker(parentJobId, parentWorkerId, parentContext, jobSto
     const metadata = options?.metadata ?? {};
     const serializedContext = {};
     if (parentContext.requestId) serializedContext.requestId = parentContext.requestId;
+    if (parentContext.userId) serializedContext.userId = parentContext.userId;
     const messageBody = {
       workerId: calleeWorkerId,
       jobId: childJobId,
@@ -918,13 +1284,14 @@ async function sendWebhook(webhookUrl, payload) {
     });
   }
 }
-function createLambdaHandler(handler, outputSchema) {
+function createLambdaHandler(handler, outputSchema, options) {
   return async (event, lambdaContext) => {
     const promises = event.Records.map(async (record) => {
       let messageBody = null;
       try {
         messageBody = JSON.parse(record.body);
-        const { workerId, jobId, input, context, webhookUrl, metadata = {} } = messageBody;
+        const { workerId, jobId, input, context, webhookUrl, metadata = {}, userId: messageUserId, maxTokens } = messageBody;
+        const userId = context.userId ?? messageUserId;
         const raw = (process.env.WORKER_DATABASE_TYPE || "upstash-redis").toLowerCase();
         const jobStoreType = raw === "mongodb" ? "mongodb" : "upstash-redis";
         if (jobStoreType === "upstash-redis" && isRedisJobStoreConfigured()) {
@@ -950,33 +1317,45 @@ function createLambdaHandler(handler, outputSchema) {
         }
         let jobStore;
         if (jobStoreType === "upstash-redis" && isRedisJobStoreConfigured()) {
-          await upsertRedisJob(jobId, workerId, input, metadata);
-          jobStore = createRedisJobStore(workerId, jobId, input, metadata);
+          await upsertRedisJob(jobId, workerId, input, metadata, userId);
+          jobStore = createRedisJobStore(workerId, jobId, input, metadata, userId);
         } else if (jobStoreType === "mongodb" || isMongoJobStoreConfigured()) {
-          await upsertJob(jobId, workerId, input, metadata);
-          jobStore = createMongoJobStore(workerId, jobId, input, metadata);
+          await upsertJob(jobId, workerId, input, metadata, userId);
+          jobStore = createMongoJobStore(workerId, jobId, input, metadata, userId);
+        }
+        if (userId) {
+          console.log(`[WORKER_USER:${userId}]`, { jobId, workerId, timestamp: (/* @__PURE__ */ new Date()).toISOString() });
         }
         const baseContext = {
           jobId,
           workerId,
           requestId: context.requestId || lambdaContext.awsRequestId,
+          ...userId ? { userId } : {},
           ...context
         };
+        const tokenTracker = createTokenTracker(maxTokens ?? null);
+        const logger = createWorkerLogger(jobId, workerId);
         const handlerContext = {
           ...baseContext,
           ...jobStore ? { jobStore } : {},
-          logger: createWorkerLogger(jobId, workerId),
-          dispatchWorker: createDispatchWorker(
-            jobId,
-            workerId,
-            baseContext,
-            jobStore
-          )
+          logger,
+          dispatchWorker: createDispatchWorker(jobId, workerId, baseContext, jobStore),
+          reportTokenUsage: async (usage) => {
+            tokenTracker.report(usage);
+            const state = tokenTracker.getState();
+            if (jobStore) {
+              await jobStore.update({ metadata: { tokenUsage: state } }).catch((e) => {
+                logger.warn("Failed to persist tokenUsage to job store", { error: e?.message });
+              });
+            }
+          },
+          getTokenBudget: () => tokenTracker.getBudgetInfo(),
+          retryContext: void 0
         };
         if (jobStore) {
           try {
             await jobStore.update({ status: "running" });
-            const queueCtxForLog = input?.__workerQueue ?? metadata?.__workerQueue;
+            const queueCtxForLog = getWorkerQueueContext(input, metadata);
             console.log("[Worker] Job status updated to running:", {
               jobId,
               workerId,
@@ -991,7 +1370,7 @@ function createLambdaHandler(handler, outputSchema) {
             });
           }
         }
-        const queueCtx = input?.__workerQueue ?? metadata?.__workerQueue;
+        const queueCtx = getWorkerQueueContext(input, metadata);
         if (queueCtx?.queueJobId && typeof queueCtx.stepIndex === "number") {
           if (queueCtx.stepIndex === 0) {
             try {
@@ -1000,7 +1379,8 @@ function createLambdaHandler(handler, outputSchema) {
                 queueId: queueCtx.id,
                 firstWorkerId: workerId,
                 firstWorkerJobId: jobId,
-                metadata
+                metadata,
+                userId
               });
             } catch (e) {
               console.warn("[Worker] Failed to upsert initial queue job:", {
@@ -1012,7 +1392,9 @@ function createLambdaHandler(handler, outputSchema) {
           }
           await notifyQueueJobStep(queueCtx.queueJobId, "start", {
             queueId: queueCtx.id,
-            stepIndex: queueCtx.stepIndex,
+            // Use arrayStepIndex when set — it tracks the actual steps[] position for
+            // looping steps where the definition index stays fixed across iterations.
+            stepIndex: queueCtx.arrayStepIndex ?? queueCtx.stepIndex,
             workerJobId: jobId,
             workerId,
             input
@@ -1020,12 +1402,21 @@ function createLambdaHandler(handler, outputSchema) {
         }
         let output;
         try {
-          output = await handler({
-            input,
-            ctx: handlerContext
-          });
-          if (outputSchema) {
-            output = outputSchema.parse(output);
+          const workerRetryConfig = options?.retry;
+          const executeHandler = async (retryCtx) => {
+            handlerContext.retryContext = retryCtx;
+            const result = await handler({ input, ctx: handlerContext });
+            return outputSchema ? outputSchema.parse(result) : result;
+          };
+          if (workerRetryConfig && workerRetryConfig.on.length > 0) {
+            output = await executeWithRetry(executeHandler, workerRetryConfig, (retryCtx, delayMs) => {
+              logger.warn(
+                `[worker-retry] Retrying handler (attempt ${retryCtx.attempt}/${retryCtx.maxAttempts}): ${retryCtx.lastError.message}`,
+                { delayMs }
+              );
+            });
+          } else {
+            output = await executeHandler(void 0);
           }
         } catch (error) {
           const errorPayload = {
@@ -1057,7 +1448,7 @@ function createLambdaHandler(handler, outputSchema) {
               });
             }
           }
-          const queueCtxFail = input?.__workerQueue ?? metadata?.__workerQueue;
+          const queueCtxFail = getWorkerQueueContext(input, metadata);
           if (queueCtxFail?.queueJobId && typeof queueCtxFail.stepIndex === "number") {
             await notifyQueueJobStep(queueCtxFail.queueJobId, "fail", {
               queueId: queueCtxFail.id,
@@ -1121,6 +1512,7 @@ function createLambdaHandler(handler, outputSchema) {
 // Annotate the CommonJS export names for ESM import in node:
 0 && (module.exports = {
   SQS_MAX_DELAY_SECONDS,
+  TokenBudgetExceededError,
   createLambdaHandler,
   createWorkerLogger,
   wrapHandlerForQueue