@k2b/cloud 0.6.0-rc.0

package/src/types/ambient.d.ts

@@ -0,0 +1,24 @@
+// Cross-app ambient module declarations. Every app's tsconfig includes this
+// so apps can import non-TS assets (e.g. `import md from './x.md' with { type: "text" }`)
+// without re-declaring the stub per-package.
+
+declare module "*.md" {
+  const content: string;
+  export default content;
+}
+
+declare module "*.sh" {
+  const content: string;
+  export default content;
+}
+
+declare module "*service-worker.js" {
+  const content: string;
+  export default content;
+}
+
+declare module "@babel/core" {
+  export const types: typeof import("@babel/types");
+
+  export function transformAsync(source: string, options: Record<string, unknown>): Promise<{ code?: string } | null>;
+}

package/src/workflows/ai/index.ts

@@ -0,0 +1,17 @@
+export {
+  WORKFLOW_AI_DEPENDENCY_KIND,
+  createWorkflowAiTask,
+  getWorkflowAiTask,
+  getWorkflowAiTaskByEffectKey,
+  migrateWorkflowAi,
+  wakeWorkflowAiTask,
+  workflowAiTaskExists,
+} from "./store";
+export { startWorkflowAiRuntime, stopWorkflowAiRuntime, submitWorkflowAiTask, type WorkflowAiRuntimeOptions } from "./runtime";
+export {
+  type WorkflowAiRequest,
+  type WorkflowAiRequestInput,
+  type WorkflowAiTask,
+  type WorkflowAiTaskStatus,
+  workflowAiRequestSchema,
+} from "./types";

package/src/workflows/ai/runtime.ts

@@ -0,0 +1,286 @@
+import { StructuredOutputError } from "@k2b/nessi";
+import type { Job, Worker } from "@k2b/sync";
+import { z } from "zod";
+import { lazySync } from "../../_internal/process-sync";
+import { isAiSettingsError, type RunAiStructuredInput, runAiStructured } from "../../ai";
+
+import type { WorkflowJsonValue } from "../contracts";
+import {
+  claimWorkflowAiTask,
+  completeWorkflowAiTask,
+  failWorkflowAiTask,
+  getWorkflowAiTask,
+  listRecoverableWorkflowAiTaskIds,
+  markWorkflowAiTaskCanceledIfRequested,
+  requeueWorkflowAiTask,
+  wakeWorkflowAiTask,
+  workflowAiTaskCancellationRequested,
+} from "./store";
+import type { WorkflowAiRequest, WorkflowAiTask } from "./types";
+
+const JOB_ID = "cloud.workflow-ai";
+const MAX_ATTEMPTS = 3;
+const DEFAULT_CANCEL_POLL_MS = 500;
+const HEARTBEAT_INTERVAL_MS = 10_000;
+
+type StructuredRunner = typeof runAiStructured;
+
+export type WorkflowAiRuntimeOptions = {
+  runStructured?: StructuredRunner;
+  cancelPollMs?: number;
+};
+
+/** Attempt budget shared by the job declaration and the Postgres-side settle logic. */
+export type WorkflowAiProcessOptions = Required<WorkflowAiRuntimeOptions> & { maxAttempts: number };
+
+const workflowAiJob = lazySync((sync) =>
+  sync.job<{ taskId: string }>({
+    id: JOB_ID,
+    owner: "cloud",
+    delivery: { ackWaitMs: 30_000, maxAttempts: MAX_ATTEMPTS, backoffMs: [1_000, 2_000] },
+  }),
+);
+
+class WorkflowAiAttemptError extends Error {
+  constructor(
+    readonly code: string,
+    message: string,
+    readonly retryable: boolean,
+  ) {
+    super(message);
+  }
+}
+
+type WorkflowAiAttemptFailure = Pick<WorkflowAiAttemptError, "code" | "message" | "retryable">;
+
+export const settleWorkflowAiAttemptFailure = async (
+  taskId: string,
+  error: WorkflowAiAttemptFailure,
+  failureCount: number,
+  maxAttempts: number,
+): Promise<"canceled" | "retry" | "failed"> => {
+  if (await workflowAiTaskCancellationRequested(taskId)) {
+    await markWorkflowAiTaskCanceledIfRequested(taskId);
+    return "canceled";
+  }
+  if (error.retryable && failureCount + 1 < maxAttempts) {
+    await requeueWorkflowAiTask(taskId);
+    return "retry";
+  }
+  await failWorkflowAiTask(taskId, { code: error.code, message: error.message });
+  return "failed";
+};
+
+const asJson = (value: unknown): WorkflowJsonValue => JSON.parse(JSON.stringify(value)) as WorkflowJsonValue;
+const choiceEnum = (values: string[]) => z.enum(values as [string, ...string[]]);
+const structuredOutputSchema = (fields: Extract<WorkflowAiRequest, { kind: "extract_data" }>["fields"]) => {
+  const shape: Record<string, z.ZodType> = {};
+  for (const field of fields) {
+    let schema: z.ZodType =
+      field.type === "text"
+        ? z.string().max(field.maxLength ?? 20_000)
+        : field.type === "number"
+          ? z.number()
+          : field.type === "boolean"
+            ? z.boolean()
+            : field.type === "date_time"
+              ? z.iso.datetime({ offset: true })
+              : choiceEnum(field.choices!);
+    if (!field.required) schema = schema.optional();
+    shape[field.name] = schema.describe(field.description);
+  }
+  return z.object(shape).strict();
+};
+
+export const executeWorkflowAiRequest = async (task: WorkflowAiTask, runStructured: StructuredRunner, signal: AbortSignal) => {
+  const request = task.request;
+  const common = {
+    appId: task.appId,
+    attribution: { workflowRunId: task.runId, userId: task.usageUserId },
+    requestedModelId: task.modelProfileId,
+    signal,
+    temperature: 0,
+  } satisfies Partial<RunAiStructuredInput<z.ZodType>>;
+
+  if (request.kind === "generate_text") {
+    const result = await runStructured({
+      ...common,
+      task: "workflow-generate-text",
+      systemPrompt: request.prompt,
+      input: request.input === undefined ? "Create the requested text." : JSON.stringify(request.input),
+      outputName: "generated_text",
+      output: z.object({ text: z.string().max(request.maxOutputChars) }),
+      maxOutputTokens: Math.min(8_192, Math.max(64, Math.ceil(request.maxOutputChars / 2))),
+    });
+    return { output: result.output.text as WorkflowJsonValue, usage: result.usage ? asJson(result.usage) : null };
+  }
+
+  if (request.kind === "classify") {
+    const result = await runStructured({
+      ...common,
+      task: "workflow-classify",
+      systemPrompt: `${request.prompt}\nReturn exactly one of the declared choices.`,
+      input: JSON.stringify(request.input),
+      outputName: "classification",
+      output: z.object({ choice: choiceEnum(request.choices) }),
+      maxOutputTokens: 200,
+    });
+    return { output: result.output.choice as WorkflowJsonValue, usage: result.usage ? asJson(result.usage) : null };
+  }
+
+  if (request.kind === "extract_data") {
+    const result = await runStructured({
+      ...common,
+      task: "workflow-extract-data",
+      systemPrompt: `${request.prompt}\nReturn only the declared fields. Do not invent values that are not supported by the input.`,
+      input: JSON.stringify(request.input),
+      outputName: "structured_data",
+      output: structuredOutputSchema(request.fields),
+      maxOutputTokens: 4_096,
+    });
+    return { output: asJson(result.output), usage: result.usage ? asJson(result.usage) : null };
+  }
+
+  const maximum = request.maxChoices ?? request.choices.length;
+  const result = await runStructured({
+    ...common,
+    task: "workflow-classify-many",
+    systemPrompt: `${request.prompt}\nReturn only unique values from the declared choices.`,
+    input: JSON.stringify(request.input),
+    outputName: "classifications",
+    output: z.object({ choices: z.array(choiceEnum(request.choices)).min(request.minChoices).max(maximum) }),
+    maxOutputTokens: 500,
+  });
+  const selected = new Set(result.output.choices);
+  return {
+    output: request.choices.filter((choice) => selected.has(choice)) as WorkflowJsonValue,
+    usage: result.usage ? asJson(result.usage) : null,
+  };
+};
+
+const retryableError = (error: unknown): boolean => {
+  if (isAiSettingsError(error)) return false;
+  if (error instanceof StructuredOutputError) return error.code === "loop_failed";
+  return true;
+};
+
+const errorCode = (error: unknown): string => {
+  if (isAiSettingsError(error)) return error.aiError.code.toUpperCase();
+  if (error instanceof StructuredOutputError) return `WORKFLOW_AI_${error.code.toUpperCase()}`;
+  return "WORKFLOW_AI_PROVIDER_ERROR";
+};
+
+export const processWorkflowAiTask = async (
+  taskId: string,
+  ctx: { signal: AbortSignal; heartbeat(): Promise<void> },
+  options: WorkflowAiProcessOptions,
+): Promise<void> => {
+  const task = await claimWorkflowAiTask(taskId);
+  if (!task) {
+    const terminal = await getWorkflowAiTask(taskId);
+    if (terminal && ["succeeded", "failed", "canceled"].includes(terminal.status)) await wakeWorkflowAiTask(terminal);
+    return;
+  }
+  if (task.status === "canceled") return;
+
+  const controller = new AbortController();
+  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 });
+
+  const monitor = setInterval(() => {
+    if (checking || controller.signal.aborted) return;
+    checking = true;
+    void (async () => {
+      try {
+        if (await workflowAiTaskCancellationRequested(task.id)) {
+          canceled = true;
+          controller.abort(new Error("Workflow canceled"));
+          return;
+        }
+        if (Date.now() - lastHeartbeat >= HEARTBEAT_INTERVAL_MS) {
+          await ctx.heartbeat();
+          lastHeartbeat = Date.now();
+        }
+      } catch (error) {
+        monitorError = error;
+        controller.abort(error);
+      } finally {
+        checking = false;
+      }
+    })();
+  }, options.cancelPollMs);
+
+  try {
+    const result = await executeWorkflowAiRequest(task, options.runStructured, controller.signal);
+    if (canceled || (await workflowAiTaskCancellationRequested(task.id))) {
+      await markWorkflowAiTaskCanceledIfRequested(task.id);
+      return;
+    }
+    await completeWorkflowAiTask(task.id, result.output, result.usage);
+  } catch (error) {
+    if (canceled || (await workflowAiTaskCancellationRequested(task.id))) {
+      await markWorkflowAiTaskCanceledIfRequested(task.id);
+      return;
+    }
+    if (ctx.signal.aborted || monitorError) {
+      throw new WorkflowAiAttemptError("WORKFLOW_AI_INTERRUPTED", "Workflow AI task was interrupted.", true);
+    }
+    throw new WorkflowAiAttemptError(
+      errorCode(error),
+      error instanceof Error ? error.message : "Workflow AI provider call failed.",
+      retryableError(error),
+    );
+  } finally {
+    clearInterval(monitor);
+    ctx.signal.removeEventListener("abort", abortForShutdown);
+  }
+};
+
+let activeJob: Job<{ taskId: string }> | null = null;
+let activeWorker: Worker | null = null;
+
+export const startWorkflowAiRuntime = async (input: WorkflowAiRuntimeOptions = {}): Promise<void> => {
+  if (activeJob) return;
+  const options: WorkflowAiProcessOptions = {
+    runStructured: input.runStructured ?? runAiStructured,
+    maxAttempts: MAX_ATTEMPTS,
+    cancelPollMs: input.cancelPollMs ?? DEFAULT_CANCEL_POLL_MS,
+  };
+  const next = workflowAiJob();
+  activeJob = next;
+  try {
+    activeWorker = await next.process({ concurrency: 1 }, async (context) => {
+      try {
+        await processWorkflowAiTask(context.input.taskId, context, options);
+      } catch (caught) {
+        const error =
+          caught instanceof WorkflowAiAttemptError
+            ? caught
+            : new WorkflowAiAttemptError("WORKFLOW_AI_RUNTIME_ERROR", caught instanceof Error ? caught.message : String(caught), true);
+        const outcome = await settleWorkflowAiAttemptFailure(context.input.taskId, error, context.failureCount, options.maxAttempts);
+        if (outcome === "retry") throw error;
+        // Postgres owns the terminal result, including intentional cancellation.
+      }
+    });
+    for (const taskId of await listRecoverableWorkflowAiTaskIds()) await submitWorkflowAiTask(taskId);
+  } catch (error) {
+    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 } });
+};
+
+export const stopWorkflowAiRuntime = (): void => {
+  activeWorker?.stop();
+
+  activeWorker = null;
+  activeJob = null;
+};

package/src/workflows/ai/store.ts

@@ -0,0 +1,314 @@
+import { type SQL, sql } from "bun";
+import type { WorkflowJsonValue } from "../contracts";
+import { wakeWorkflowRunsWaitingOn } from "../store/runs";
+import {
+  type WorkflowAiRequest,
+  type WorkflowAiRequestInput,
+  type WorkflowAiTask,
+  type WorkflowAiTaskStatus,
+  workflowAiRequestSchema,
+} from "./types";
+
+export const WORKFLOW_AI_DEPENDENCY_KIND = "ai.workflow-task";
+
+type WorkflowAiTaskRow = {
+  usage_user_id: string | null;
+  id: string;
+  app_id: string;
+  run_id: string;
+  step_key: string;
+  effect_key: string;
+  kind: WorkflowAiRequest["kind"];
+  request: unknown;
+  input_hash: string;
+  model_profile_id: string;
+  status: WorkflowAiTaskStatus;
+  output: unknown;
+  usage: unknown;
+  attempts: number | string;
+  error_code: string | null;
+  error_message: string | null;
+  cancel_requested_at: Date | string | null;
+  created_at: Date | string;
+  started_at: Date | string | null;
+  completed_at: Date | string | null;
+  updated_at: Date | string;
+};
+
+const iso = (value: Date | string): string => (value instanceof Date ? value.toISOString() : new Date(value).toISOString());
+const optionalIso = (value: Date | string | null): string | null => (value === null ? null : iso(value));
+const json = (value: unknown): unknown => {
+  if (typeof value !== "string") return value;
+  try {
+    return JSON.parse(value);
+  } catch {
+    return value;
+  }
+};
+
+const taskFromRow = (row: WorkflowAiTaskRow): WorkflowAiTask => ({
+  id: row.id,
+  appId: row.app_id,
+  runId: row.run_id,
+  usageUserId: row.usage_user_id ?? undefined,
+  stepKey: row.step_key,
+  effectKey: row.effect_key,
+  kind: row.kind,
+  request: workflowAiRequestSchema.parse(json(row.request)) as WorkflowAiRequest,
+  inputHash: row.input_hash,
+  modelProfileId: row.model_profile_id,
+  status: row.status,
+  output: json(row.output) as WorkflowJsonValue | null,
+  usage: json(row.usage) as WorkflowJsonValue | null,
+  attempts: Number(row.attempts),
+  errorCode: row.error_code,
+  errorMessage: row.error_message,
+  cancelRequestedAt: optionalIso(row.cancel_requested_at),
+  createdAt: iso(row.created_at),
+  startedAt: optionalIso(row.started_at),
+  completedAt: optionalIso(row.completed_at),
+  updatedAt: iso(row.updated_at),
+});
+
+const TASK_COLUMNS = `
+  (SELECT u.id::text FROM workflows.run r JOIN auth.users u ON u.id::text = r.authorization_snapshot #>> '{actor,userId}' WHERE r.id = task.run_id) AS usage_user_id,
+  task.id::text,
+  task.app_id,
+  task.run_id::text,
+  task.step_key,
+  task.effect_key,
+  task.kind,
+  task.request,
+  task.input_hash,
+  task.model_profile_id,
+  task.status,
+  task.output,
+  task.usage,
+  task.attempts,
+  task.error_code,
+  task.error_message,
+  task.cancel_requested_at,
+  task.created_at,
+  task.started_at,
+  task.completed_at,
+  task.updated_at
+`;
+
+export const migrateWorkflowAi = async (db: SQL = sql): Promise<void> => {
+  await db`CREATE SCHEMA IF NOT EXISTS ai`.simple();
+  await db`
+    CREATE TABLE IF NOT EXISTS ai.workflow_task (
+      id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
+      app_id TEXT NOT NULL,
+      run_id UUID NOT NULL REFERENCES workflows.run(id) ON DELETE CASCADE,
+      step_key TEXT NOT NULL CHECK (char_length(step_key) BETWEEN 1 AND 1000),
+      effect_key TEXT NOT NULL UNIQUE CHECK (char_length(effect_key) BETWEEN 1 AND 500),
+      kind TEXT NOT NULL CHECK (kind IN ('generate_text', 'classify', 'classify_many', 'extract_data')),
+      request JSONB NOT NULL CHECK (jsonb_typeof(request) = 'object' AND octet_length(request::text) <= 100000),
+      input_hash TEXT NOT NULL CHECK (char_length(input_hash) = 64),
+      model_profile_id TEXT NOT NULL CHECK (char_length(model_profile_id) BETWEEN 1 AND 120),
+      status TEXT NOT NULL DEFAULT 'queued' CHECK (status IN ('queued', 'running', 'succeeded', 'failed', 'canceled')),
+      output JSONB,
+      usage JSONB,
+      attempts INTEGER NOT NULL DEFAULT 0 CHECK (attempts >= 0),
+      error_code TEXT,
+      error_message TEXT CHECK (error_message IS NULL OR char_length(error_message) <= 2000),
+      cancel_requested_at TIMESTAMPTZ,
+      created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
+      started_at TIMESTAMPTZ,
+      completed_at TIMESTAMPTZ,
+      updated_at TIMESTAMPTZ NOT NULL DEFAULT now(),
+      CONSTRAINT ai_workflow_task_terminal_check CHECK (
+        (status IN ('succeeded', 'failed', 'canceled')) = (completed_at IS NOT NULL)
+      )
+    )
+  `.simple();
+  await db`
+    ALTER TABLE ai.workflow_task DROP CONSTRAINT IF EXISTS workflow_task_kind_check;
+    ALTER TABLE ai.workflow_task ADD CONSTRAINT workflow_task_kind_check
+      CHECK (kind IN ('generate_text', 'classify', 'classify_many', 'extract_data'))
+  `.simple();
+  await db`
+    CREATE INDEX IF NOT EXISTS idx_ai_workflow_task_recovery
+    ON ai.workflow_task(status, updated_at, id)
+    WHERE status IN ('queued', 'running')
+  `.simple();
+};
+
+const requestHash = (request: WorkflowAiRequest): string => new Bun.CryptoHasher("sha256").update(JSON.stringify(request)).digest("hex");
+
+export const getWorkflowAiTask = async (id: string, db: SQL = sql): Promise<WorkflowAiTask | null> => {
+  const rows = await db<WorkflowAiTaskRow[]>`
+    SELECT ${db.unsafe(TASK_COLUMNS)}
+    FROM ai.workflow_task AS task
+    WHERE task.id = ${id}::uuid
+  `;
+  return rows[0] ? taskFromRow(rows[0]) : null;
+};
+
+export const getWorkflowAiTaskByEffectKey = async (effectKey: string, db: SQL = sql): Promise<WorkflowAiTask | null> => {
+  const rows = await db<WorkflowAiTaskRow[]>`
+    SELECT ${db.unsafe(TASK_COLUMNS)}
+    FROM ai.workflow_task AS task
+    WHERE task.effect_key = ${effectKey}
+  `;
+  return rows[0] ? taskFromRow(rows[0]) : null;
+};
+
+export const workflowAiTaskExists = async (effectKey: string, db: SQL = sql): Promise<boolean> => {
+  const [row] = await db<{ exists: boolean }[]>`
+    SELECT EXISTS (SELECT 1 FROM ai.workflow_task WHERE effect_key = ${effectKey}) AS exists
+  `;
+  return Boolean(row?.exists);
+};
+
+export const createWorkflowAiTask = async (
+  input: { runId: string; stepKey: string; effectKey: string; request: WorkflowAiRequestInput; modelProfileId: string },
+  db: SQL = sql,
+): Promise<{ task: WorkflowAiTask; created: boolean }> => {
+  const request = workflowAiRequestSchema.parse(input.request) as WorkflowAiRequest;
+  const rows = await db<WorkflowAiTaskRow[]>`
+    INSERT INTO ai.workflow_task (app_id, run_id, step_key, effect_key, kind, request, input_hash, model_profile_id)
+    SELECT run.app_id, run.id, ${input.stepKey}, ${input.effectKey}, ${request.kind}, ${request}, ${requestHash(request)}, ${input.modelProfileId}
+    FROM workflows.run AS run
+    WHERE run.id = ${input.runId}::uuid
+    ON CONFLICT (effect_key) DO NOTHING
+    RETURNING ${db.unsafe(TASK_COLUMNS.replaceAll("task.", ""))}
+  `;
+  if (rows[0]) return { task: taskFromRow(rows[0]), created: true };
+  const existing = await getWorkflowAiTaskByEffectKey(input.effectKey, db);
+  if (!existing) throw new Error("Workflow AI task could not be created because its run is unavailable.");
+  return { task: existing, created: false };
+};
+
+export const listRecoverableWorkflowAiTaskIds = async (db: SQL = sql): Promise<string[]> => {
+  const rows = await db<{ id: string }[]>`
+    SELECT id::text
+    FROM ai.workflow_task
+    WHERE status IN ('queued', 'running')
+    ORDER BY updated_at, id
+  `;
+  return rows.map((row) => row.id);
+};
+
+export const claimWorkflowAiTask = async (id: string, db: SQL = sql): Promise<WorkflowAiTask | null> => {
+  const canceled = await markWorkflowAiTaskCanceledIfRequested(id, db);
+  if (canceled) return canceled;
+  const rows = await db<WorkflowAiTaskRow[]>`
+    UPDATE ai.workflow_task AS task
+    SET status = 'running',
+        attempts = task.attempts + 1,
+        started_at = COALESCE(task.started_at, now()),
+        error_code = NULL,
+        error_message = NULL,
+        updated_at = now()
+    FROM workflows.run AS run
+    WHERE task.id = ${id}::uuid
+      AND run.id = task.run_id
+      AND run.cancel_requested_at IS NULL
+      AND run.state <> 'canceled'
+      AND task.status IN ('queued', 'running')
+    RETURNING ${db.unsafe(TASK_COLUMNS)}
+  `;
+  return rows[0] ? taskFromRow(rows[0]) : markWorkflowAiTaskCanceledIfRequested(id, db);
+};
+
+export const workflowAiTaskCancellationRequested = async (id: string, db: SQL = sql): Promise<boolean> => {
+  const [row] = await db<{ canceled: boolean }[]>`
+    SELECT run.cancel_requested_at IS NOT NULL OR run.state = 'canceled' AS canceled
+    FROM ai.workflow_task AS task
+    JOIN workflows.run AS run ON run.id = task.run_id
+    WHERE task.id = ${id}::uuid
+  `;
+  return Boolean(row?.canceled);
+};
+
+export const wakeWorkflowAiTask = async (task: Pick<WorkflowAiTask, "appId" | "id">): Promise<void> => {
+  await wakeWorkflowRunsWaitingOn({ appId: task.appId, kind: WORKFLOW_AI_DEPENDENCY_KIND, key: task.id });
+};
+
+export const markWorkflowAiTaskCanceledIfRequested = async (id: string, db: SQL = sql): Promise<WorkflowAiTask | null> => {
+  const rows = await db<WorkflowAiTaskRow[]>`
+    UPDATE ai.workflow_task AS task
+    SET status = 'canceled',
+        output = NULL,
+        usage = NULL,
+        error_code = 'WORKFLOW_AI_CANCELED',
+        error_message = 'Workflow AI task was canceled.',
+        cancel_requested_at = COALESCE(task.cancel_requested_at, now()),
+        completed_at = now(),
+        updated_at = now()
+    FROM workflows.run AS run
+    WHERE task.id = ${id}::uuid
+      AND run.id = task.run_id
+      AND (run.cancel_requested_at IS NOT NULL OR run.state = 'canceled')
+      AND task.status IN ('queued', 'running')
+    RETURNING ${db.unsafe(TASK_COLUMNS)}
+  `;
+  const task = rows[0] ? taskFromRow(rows[0]) : null;
+  if (task) await wakeWorkflowAiTask(task);
+  return task;
+};
+
+export const requeueWorkflowAiTask = async (id: string, db: SQL = sql): Promise<void> => {
+  await db`
+    UPDATE ai.workflow_task
+    SET status = 'queued', updated_at = now()
+    WHERE id = ${id}::uuid AND status = 'running'
+  `;
+};
+
+export const completeWorkflowAiTask = async (
+  id: string,
+  output: WorkflowJsonValue,
+  usage: WorkflowJsonValue | null,
+  db: SQL = sql,
+): Promise<WorkflowAiTask | null> => {
+  const rows = await db<WorkflowAiTaskRow[]>`
+    UPDATE ai.workflow_task AS task
+    SET status = 'succeeded', output = ${output}, usage = ${usage}, completed_at = now(), updated_at = now()
+    FROM workflows.run AS run
+    WHERE task.id = ${id}::uuid
+      AND run.id = task.run_id
+      AND run.cancel_requested_at IS NULL
+      AND run.state <> 'canceled'
+      AND task.status = 'running'
+    RETURNING ${db.unsafe(TASK_COLUMNS)}
+  `;
+  const task = rows[0] ? taskFromRow(rows[0]) : null;
+  if (task) {
+    await wakeWorkflowAiTask(task);
+    return task;
+  }
+  return markWorkflowAiTaskCanceledIfRequested(id, db);
+};
+
+export const failWorkflowAiTask = async (
+  id: string,
+  error: { code: string; message: string },
+  db: SQL = sql,
+): Promise<WorkflowAiTask | null> => {
+  const rows = await db<WorkflowAiTaskRow[]>`
+    UPDATE ai.workflow_task AS task
+    SET status = 'failed',
+        output = NULL,
+        usage = NULL,
+        error_code = ${error.code},
+        error_message = ${error.message.slice(0, 2_000)},
+        completed_at = now(),
+        updated_at = now()
+    FROM workflows.run AS run
+    WHERE task.id = ${id}::uuid
+      AND run.id = task.run_id
+      AND run.cancel_requested_at IS NULL
+      AND run.state <> 'canceled'
+      AND task.status IN ('queued', 'running')
+    RETURNING ${db.unsafe(TASK_COLUMNS)}
+  `;
+  const task = rows[0] ? taskFromRow(rows[0]) : null;
+  if (task) {
+    await wakeWorkflowAiTask(task);
+    return task;
+  }
+  return markWorkflowAiTaskCanceledIfRequested(id, db);
+};