@microfox/ai-worker 1.0.3 → 1.0.5

package/dist/handler.d.ts

@@ -1,5 +1,29 @@
 import { SQSEvent, Context } from 'aws-lambda';
 import { ZodType } from 'zod';
+import { R as RetryContext, S as SmartRetryConfig, Q as QueueStepOutput, b as ChainContext, H as HitlResumeContext, L as LoopContext } from './queue-B5n6YVQV.js';
+export { B as BuiltInRetryPattern, C as CustomRetryPattern, a as RetryPattern } from './queue-B5n6YVQV.js';
+import './hitlConfig.js';
+
+/**
+ * Token budget tracking for workers.
+ * Workers report usage via ctx.reportTokenUsage(); the runtime accumulates
+ * and throws TokenBudgetExceededError when the limit is reached.
+ */
+interface TokenUsage {
+    inputTokens: number;
+    outputTokens: number;
+}
+interface TokenBudgetState {
+    inputTokens: number;
+    outputTokens: number;
+    /** null = no budget configured */
+    budget: number | null;
+}
+declare class TokenBudgetExceededError extends Error {
+    readonly used: number;
+    readonly budget: number;
+    constructor(used: number, budget: number);
+}
 
 /**
  * Generic Lambda handler wrapper for worker agents.
@@ -34,6 +58,7 @@ interface JobRecord {
         jobId: string;
         workerId: string;
     }>;
+    userId?: string;
     createdAt: string;
     updatedAt: string;
     completedAt?: string;
@@ -98,6 +123,8 @@ interface WorkerHandlerParams<INPUT, OUTPUT> {
         jobId: string;
         workerId: string;
         requestId?: string;
+        /** ID of the user who triggered this job. Pass via DispatchOptions.userId from your API route. */
+        userId?: string;
         /**
          * Job store interface for updating and retrieving job state.
          * Uses MongoDB directly when configured; never HTTP/origin URL.
@@ -116,49 +143,114 @@ interface WorkerHandlerParams<INPUT, OUTPUT> {
             messageId?: string;
             output?: unknown;
         }>;
+        /**
+         * Report token usage after an LLM call. Accumulates across all calls in this job.
+         * Throws TokenBudgetExceededError if the configured maxTokens budget is exceeded.
+         * Also persists usage to the job store for observability.
+         *
+         * @example
+         * ```ts
+         * const result = await anthropic.messages.create({ ... });
+         * await ctx.reportTokenUsage({
+         *   inputTokens: result.usage.input_tokens,
+         *   outputTokens: result.usage.output_tokens,
+         * });
+         * ```
+         */
+        reportTokenUsage: (usage: TokenUsage) => Promise<void>;
+        /**
+         * Get the current token usage and remaining budget for this job.
+         * Returns `{ used, budget: null, remaining: null }` when no maxTokens was set.
+         */
+        getTokenBudget: () => {
+            used: number;
+            budget: number | null;
+            remaining: number | null;
+        };
+        /**
+         * Populated on retry attempts (attempt >= 2). Contains info about the previous failure
+         * so the handler can self-correct (e.g. inject the error message into the next prompt).
+         * `undefined` on the first attempt — use `if (ctx.retryContext)` to detect retries.
+         */
+        retryContext?: RetryContext;
         [key: string]: any;
     };
 }
+
 type WorkerHandler<INPUT, OUTPUT> = (params: WorkerHandlerParams<INPUT, OUTPUT>) => Promise<OUTPUT>;
 /** Result of getNextStep for queue chaining. */
 interface QueueNextStep {
     workerId: string;
     delaySeconds?: number;
-    mapInputFromPrev?: string;
+    requiresApproval?: boolean;
+    /** Whether this step has a `chain` function (or built-in string) defined. */
+    hasChain?: boolean;
+    /** Whether this step has a `resume` function defined. */
+    hasResume?: boolean;
+    /** Optional HITL metadata from queue step config (UI/tooling only). */
+    hitl?: {
+        ui?: unknown;
+    } | unknown;
+    /** Smart retry config for this step. Overrides worker-level retry for this step only. */
+    retry?: SmartRetryConfig;
 }
-/** One previous step's output (for mapInputFromPrev context). */
-interface QueueStepOutput {
-    stepIndex: number;
-    workerId: string;
-    output: unknown;
+/**
+ * @deprecated Use {@link ChainContext} for the normal chain path and
+ * {@link HitlResumeContext} for the HITL resume path instead.
+ * Kept for backwards compatibility with queue files written against the old API.
+ */
+interface MapStepInputContext {
+    initialInput: unknown;
+    previousOutputs: QueueStepOutput[];
+    /** @deprecated Use HitlResumeContext.reviewerInput instead. */
+    hitlInput?: unknown;
+    /** @deprecated Use HitlResumeContext.pendingInput instead. */
+    pendingStepInput?: Record<string, unknown>;
 }
+
 /** Runtime helpers for queue-aware wrappers (provided by generated registry). */
 interface QueueRuntime {
     getNextStep(queueId: string, stepIndex: number): QueueNextStep | undefined;
-    /** Optional: when provided, mapping can use outputs from any previous step. */
+    /** Step config at `stepIndex`. */
+    getStepAt?(queueId: string, stepIndex: number): QueueNextStep | undefined;
+    /** Optional: when provided, mappers can use outputs from any previous step. */
     getQueueJob?(queueJobId: string): Promise<{
         steps: Array<{
             workerId: string;
             output?: unknown;
         }>;
     } | null>;
-    /** (initialInput, previousOutputs) – previousOutputs includes outputs for steps 0..stepIndex-1 and current step. */
-    invokeMapInput?(queueId: string, stepIndex: number, initialInput: unknown, previousOutputs: QueueStepOutput[]): Promise<unknown> | unknown;
+    /**
+     * Build the input for a step when the queue advances normally (no HITL resume).
+     * Calls the step's `chain` function, or the built-in passthrough/continueFromPrevious.
+     */
+    invokeChain?(queueId: string, stepIndex: number, ctx: ChainContext): Promise<unknown> | unknown;
+    /**
+     * Build the domain input for a step when it resumes after HITL approval.
+     * Calls the step's `resume` function, or merges pendingInput + reviewerInput by default.
+     */
+    invokeResume?(queueId: string, stepIndex: number, ctx: HitlResumeContext): Promise<unknown> | unknown;
+    /**
+     * Evaluate whether a looping step should run again.
+     * Calls the step's `loop.shouldContinue` function. Returns false if none defined.
+     */
+    invokeLoop?(queueId: string, stepIndex: number, ctx: LoopContext): Promise<boolean> | boolean;
 }
 /**
- * Wraps a user handler so that when the job has __workerQueue context (from
- * dispatchQueue or queue cron), it dispatches the next worker in the sequence
- * after the handler completes. Uses literal worker IDs so the CLI env injection
- * picks up WORKER_QUEUE_URL_* for next-step workers.
+ * Wraps a user handler so that when the job has `__workerQueue` context (from
+ * `dispatchQueue` or queue cron), it dispatches the next worker in the sequence
+ * **after** the handler completes.
+ *
+ * All queue/HITL envelope keys (`__workerQueue`, `__hitlInput`, `__hitlDecision`,
+ * `__hitlPending`, `hitl`) are **stripped from `params.input` before the user handler
+ * runs** — workers receive clean domain input and do not need to accept these keys
+ * in their Zod schemas.
+ *
+ * **HITL resume:** When `__hitlInput` is present, `invokeResume` is called first to
+ * produce the merged domain input. **Chain advancement:** After a step completes,
+ * `invokeChain` is called to compute the next step's input.
  */
-declare function wrapHandlerForQueue<INPUT, OUTPUT>(handler: WorkerHandler<INPUT, OUTPUT>, queueRuntime: QueueRuntime): WorkerHandler<INPUT & {
-    __workerQueue?: {
-        id: string;
-        stepIndex: number;
-        initialInput: unknown;
-        queueJobId?: string;
-    };
-}, OUTPUT>;
+declare function wrapHandlerForQueue<INPUT, OUTPUT>(handler: WorkerHandler<INPUT, OUTPUT>, queueRuntime: QueueRuntime): WorkerHandler<INPUT, OUTPUT>;
 interface SQSMessageBody {
     workerId: string;
     jobId: string;
@@ -169,6 +261,10 @@ interface SQSMessageBody {
     jobStoreUrl?: string;
     metadata?: Record<string, any>;
     timestamp: string;
+    /** ID of the user who triggered this job. Forwarded from dispatch options. */
+    userId?: string;
+    /** Maximum total tokens (input + output) for this job. Forwarded from DispatchOptions.maxTokens. */
+    maxTokens?: number;
 }
 interface WebhookPayload {
     jobId: string;
@@ -190,6 +286,8 @@ interface WebhookPayload {
  * @param outputSchema - Optional Zod schema for output validation
  * @returns A Lambda handler function
  */
-declare function createLambdaHandler<INPUT, OUTPUT>(handler: WorkerHandler<INPUT, OUTPUT>, outputSchema?: ZodType<OUTPUT>): (event: SQSEvent, context: Context) => Promise<void>;
+declare function createLambdaHandler<INPUT, OUTPUT>(handler: WorkerHandler<INPUT, OUTPUT>, outputSchema?: ZodType<OUTPUT>, options?: {
+    retry?: SmartRetryConfig;
+}): (event: SQSEvent, context: Context) => Promise<void>;
 
-export { type DispatchWorkerOptions, type JobRecord, type JobStore, type JobStoreUpdate, type QueueNextStep, type QueueRuntime, type QueueStepOutput, type SQSMessageBody, SQS_MAX_DELAY_SECONDS, type WebhookPayload, type WorkerHandler, type WorkerHandlerParams, type WorkerLogger, createLambdaHandler, createWorkerLogger, wrapHandlerForQueue };
+export { ChainContext, type DispatchWorkerOptions, HitlResumeContext, type JobRecord, type JobStore, type JobStoreUpdate, LoopContext, type MapStepInputContext, type QueueNextStep, type QueueRuntime, QueueStepOutput, RetryContext, type SQSMessageBody, SQS_MAX_DELAY_SECONDS, SmartRetryConfig, TokenBudgetExceededError, type TokenBudgetState, type TokenUsage, type WebhookPayload, type WorkerHandler, type WorkerHandlerParams, type WorkerLogger, createLambdaHandler, createWorkerLogger, wrapHandlerForQueue };