@k2b/cloud 0.43.0 → 0.45.0

package/src/workflows/ai/runtime.ts

@@ -1,15 +1,22 @@
 import { StructuredOutputError } from "@k2b/nessi";
-import type { Job, Worker } from "@k2b/sync";
+import { type Job, StoreFullError, type Worker } from "@k2b/sync";
 import { lazySync } from "../../_internal/process-sync";
-import { isAiSettingsError, runAiStructured } from "../../ai";
 import { AiBackgroundAdmissionError, AiBackgroundCostError } from "../../ai/inference-calls";
+import { AiQuotaError } from "../../ai/quotas";
+import { resolveAiWorkflowModel, runAiStructured } from "../../ai/structured";
 import { executeAiTask } from "../../ai/task-execution";
+import { resolveApplicationAiTaskModel } from "../../ai/task-model";
+import { isAiSettingsError } from "../../ai/validate";
+import { env } from "../../config/env";
+import { WorkflowAiAdmissionError, workflowAiModelHash } from "./policy";
 
 import {
   claimWorkflowAiTask,
   completeWorkflowAiTask,
+  expireWorkflowAiTasks,
   failWorkflowAiTask,
   getWorkflowAiTask,
+  heartbeatWorkflowAiTask,
   listRecoverableWorkflowAiTaskIds,
   markWorkflowAiTaskCanceledIfRequested,
   requeueWorkflowAiTask,
@@ -18,7 +25,7 @@ import {
 } from "./store";
 import type { WorkflowAiTask } from "./types";
 
-const JOB_ID = "cloud.workflow-ai";
+export const WORKFLOW_AI_JOB_ID = "cloud.workflow-ai.v2";
 const MAX_ATTEMPTS = 3;
 const DEFAULT_CANCEL_POLL_MS = 500;
 const HEARTBEAT_INTERVAL_MS = 10_000;
@@ -28,16 +35,26 @@ type StructuredRunner = typeof runAiStructured;
 export type WorkflowAiRuntimeOptions = {
   runStructured?: StructuredRunner;
   cancelPollMs?: number;
+  attemptTimeoutMs?: number;
 };
 
 /** Attempt budget shared by the job declaration and the Postgres-side settle logic. */
-export type WorkflowAiProcessOptions = Required<WorkflowAiRuntimeOptions> & { maxAttempts: number };
+export type WorkflowAiProcessOptions = Required<Omit<WorkflowAiRuntimeOptions, "attemptTimeoutMs">> & {
+  maxAttempts: number;
+  attemptTimeoutMs?: number;
+};
 
 const workflowAiJob = lazySync((sync) =>
   sync.job<{ taskId: string }>({
-    id: JOB_ID,
+    id: WORKFLOW_AI_JOB_ID,
+    whenFull: "reject",
     owner: "cloud",
-    delivery: { ackWaitMs: 30_000, maxAttempts: MAX_ATTEMPTS, backoffMs: [1_000, 2_000] },
+    delivery: {
+      maxInFlight: env.WORKFLOW_AI_BUDGET.globalConcurrent,
+      ackWaitMs: 30_000,
+      maxAttempts: MAX_ATTEMPTS,
+      backoffMs: [1_000, 2_000],
+    },
   }),
 );
 
@@ -76,24 +93,40 @@ export const executeWorkflowAiRequest = async (task: WorkflowAiTask, runStructur
     task.request,
     {
       appId: task.appId,
-      attribution: { workflowRunId: task.runId, stepKey: task.stepKey, userId: task.usageUserId },
-      requestedModelId: task.modelProfileId,
+      attribution: {
+        workflowRunId: task.runId ?? undefined,
+        stepKey: task.stepKey,
+        userId: task.usageSubject?.type === "user" ? task.usageSubject.userId : task.usageUserId,
+      },
+      usageSubject: task.usageSubject,
+      model: { id: task.modelId },
+      resolveModel: async () => {
+        const resolved = task.runId
+          ? await resolveAiWorkflowModel({ id: task.modelId })
+          : await resolveApplicationAiTaskModel(task.usageSubject, { id: task.modelId });
+        if (task.modelConfigHash && workflowAiModelHash(resolved.profile) !== task.modelConfigHash)
+          throw new WorkflowAiAdmissionError("WORKFLOW_AI_MODEL_CHANGED");
+        return { ...resolved, model: task.modelRef ?? resolved.model };
+      },
       signal,
-      taskPrefix: "workflow",
+      taskPrefix: task.runId ? "workflow" : "application",
     },
     runStructured,
   );
 
 const retryableError = (error: unknown): boolean => {
+  if (error instanceof WorkflowAiAdmissionError) return error.retryable;
   if (error instanceof AiBackgroundAdmissionError) return error.retryable;
-  if (error instanceof AiBackgroundCostError) return false;
+  if (error instanceof AiBackgroundCostError || error instanceof AiQuotaError) return false;
   if (isAiSettingsError(error)) return false;
   if (error instanceof StructuredOutputError) return error.code === "loop_failed";
   return true;
 };
 
 const errorCode = (error: unknown): string => {
-  if (error instanceof AiBackgroundCostError || error instanceof AiBackgroundAdmissionError) return error.code;
+  if (error instanceof WorkflowAiAdmissionError) return error.code;
+  if (error instanceof AiBackgroundCostError || error instanceof AiBackgroundAdmissionError || error instanceof AiQuotaError)
+    return error.code;
   if (isAiSettingsError(error)) return error.aiError.code.toUpperCase();
   if (error instanceof StructuredOutputError) return `WORKFLOW_AI_${error.code.toUpperCase()}`;
   return "WORKFLOW_AI_PROVIDER_ERROR";
@@ -113,12 +146,18 @@ export const processWorkflowAiTask = async (
   if (task.status === "canceled") return;
 
   const controller = new AbortController();
+  let timedOut = false;
+  const deadline = setTimeout(() => {
+    timedOut = true;
+    controller.abort(new Error("Workflow AI attempt deadline exceeded. Retry the workflow or inspect the provider."));
+  }, options.attemptTimeoutMs ?? env.WORKFLOW_AI_ATTEMPT_TIMEOUT_MS);
   let canceled = false;
   let monitorError: unknown;
   let checking = false;
   let lastHeartbeat = Date.now();
   const abortForShutdown = () => controller.abort(ctx.signal.reason ?? new Error("Workflow AI runtime stopped"));
   ctx.signal.addEventListener("abort", abortForShutdown, { once: true });
+  if (ctx.signal.aborted) abortForShutdown();
 
   const monitor = setInterval(() => {
     if (checking || controller.signal.aborted) return;
@@ -132,6 +171,7 @@ export const processWorkflowAiTask = async (
         }
         if (Date.now() - lastHeartbeat >= HEARTBEAT_INTERVAL_MS) {
           await ctx.heartbeat();
+          await heartbeatWorkflowAiTask(task.id);
           lastHeartbeat = Date.now();
         }
       } catch (error) {
@@ -144,7 +184,13 @@ export const processWorkflowAiTask = async (
   }, options.cancelPollMs);
 
   try {
-    const result = await executeWorkflowAiRequest(task, options.runStructured, controller.signal);
+    controller.signal.throwIfAborted();
+    const aborted = new Promise<never>((_resolve, reject) => {
+      const onAbort = () => reject(controller.signal.reason);
+      if (controller.signal.aborted) onAbort();
+      else controller.signal.addEventListener("abort", onAbort, { once: true });
+    });
+    const result = await Promise.race([executeWorkflowAiRequest(task, options.runStructured, controller.signal), aborted]);
     if (canceled || (await workflowAiTaskCancellationRequested(task.id))) {
       await markWorkflowAiTaskCanceledIfRequested(task.id);
       return;
@@ -155,6 +201,13 @@ export const processWorkflowAiTask = async (
       await markWorkflowAiTaskCanceledIfRequested(task.id);
       return;
     }
+    if (timedOut) {
+      throw new WorkflowAiAttemptError(
+        "WORKFLOW_AI_TIMEOUT",
+        "Workflow AI attempt deadline exceeded. Retry the workflow or inspect the provider.",
+        true,
+      );
+    }
     if (ctx.signal.aborted || monitorError) {
       throw new WorkflowAiAttemptError("WORKFLOW_AI_INTERRUPTED", "Workflow AI task was interrupted.", true);
     }
@@ -165,12 +218,23 @@ export const processWorkflowAiTask = async (
     );
   } finally {
     clearInterval(monitor);
+    clearTimeout(deadline);
     ctx.signal.removeEventListener("abort", abortForShutdown);
   }
 };
 
 let activeJob: Job<{ taskId: string }> | null = null;
 let activeWorker: Worker | null = null;
+let recoveryTimer: ReturnType<typeof setInterval> | null = null;
+let recovery: Promise<void> | null = null;
+let shutdown: AbortController | null = null;
+
+const recover = (): Promise<void> => {
+  recovery ??= expireWorkflowAiTasks(env.WORKFLOW_AI_QUEUE_TIMEOUT_MS).finally(() => {
+    recovery = null;
+  });
+  return recovery;
+};
 
 export const startWorkflowAiRuntime = async (input: WorkflowAiRuntimeOptions = {}): Promise<void> => {
   if (activeJob) return;
@@ -178,14 +242,26 @@ export const startWorkflowAiRuntime = async (input: WorkflowAiRuntimeOptions = {
     runStructured: input.runStructured ?? runAiStructured,
     maxAttempts: MAX_ATTEMPTS,
     cancelPollMs: input.cancelPollMs ?? DEFAULT_CANCEL_POLL_MS,
+    attemptTimeoutMs: input.attemptTimeoutMs ?? env.WORKFLOW_AI_ATTEMPT_TIMEOUT_MS,
   };
   const next = workflowAiJob();
   activeJob = next;
+  const controller = new AbortController();
+  shutdown = controller;
   try {
-    activeWorker = await next.process({ concurrency: 1 }, async (context) => {
+    activeWorker = await next.process({ concurrency: env.WORKFLOW_AI_BUDGET.globalConcurrent }, async (context) => {
       try {
-        await processWorkflowAiTask(context.input.taskId, context, options);
+        await processWorkflowAiTask(
+          context.input.taskId,
+          { signal: AbortSignal.any([context.signal, controller.signal]), heartbeat: () => context.heartbeat() },
+          options,
+        );
       } catch (caught) {
+        if (controller.signal.aborted) {
+          // Shutdown is not a provider failure. Persist before drain lets the process exit.
+          if (!(await markWorkflowAiTaskCanceledIfRequested(context.input.taskId))) await requeueWorkflowAiTask(context.input.taskId);
+          throw caught;
+        }
         const error =
           caught instanceof WorkflowAiAttemptError
             ? caught
@@ -195,21 +271,55 @@ export const startWorkflowAiRuntime = async (input: WorkflowAiRuntimeOptions = {
         // Postgres owns the terminal result, including intentional cancellation.
       }
     });
-    for (const taskId of await listRecoverableWorkflowAiTaskIds()) await submitWorkflowAiTask(taskId);
+    await recover();
+    for (const taskId of await listRecoverableWorkflowAiTaskIds()) {
+      try {
+        await submitWorkflowAiTask(taskId);
+      } catch (error) {
+        // Submission already persists and wakes the failed task; recover the others.
+        console.error("Workflow AI task resubmission failed", { taskId, error });
+      }
+    }
+    recoveryTimer = setInterval(() => {
+      void recover().catch((error) => console.error("Workflow AI recovery failed", error));
+    }, HEARTBEAT_INTERVAL_MS);
   } catch (error) {
-    stopWorkflowAiRuntime();
+    await stopWorkflowAiRuntime();
     throw error;
   }
 };
 
 export const submitWorkflowAiTask = async (taskId: string): Promise<void> => {
   if (!activeJob) throw new Error("Workflow AI runtime is not started.");
-  await activeJob.submit({ key: `task:${taskId}`, coalesce: true, input: { taskId } });
+  try {
+    await activeJob.submit({ key: `task:${taskId}`, coalesce: true, input: { taskId } });
+  } catch (error) {
+    const failure = new WorkflowAiAttemptError(
+      error instanceof StoreFullError ? "WORKFLOW_AI_QUEUE_FULL" : "WORKFLOW_AI_SUBMISSION_FAILED",
+      "Workflow AI queue could not accept the task. Retry the workflow after capacity or connectivity recovers.",
+      true,
+    );
+    await failWorkflowAiTask(taskId, failure);
+    throw failure;
+  }
 };
 
-export const stopWorkflowAiRuntime = (): void => {
-  activeWorker?.stop();
+export const stopWorkflowAiRuntime = async (): Promise<void> => {
+  const worker = activeWorker;
+  worker?.stop();
+  shutdown?.abort(new Error("Workflow AI runtime stopped"));
+  if (recoveryTimer) clearInterval(recoveryTimer);
+  recoveryTimer = null;
 
-  activeWorker = null;
-  activeJob = null;
+  // Sync's existing drain budget bounds settlement; do not exit before DB requeue.
+  try {
+    await worker?.drain();
+  } finally {
+    // The recovery timer already reports its failure. Always allow a later start,
+    // including when drain rejects, while preserving the drain error for callers.
+    await recovery?.catch(() => {});
+    activeWorker = null;
+    activeJob = null;
+    shutdown = null;
+  }
 };