@corgelabs/corge 0.0.1 → 1.0.1

package/src/tasks/schema.ts

@@ -0,0 +1,310 @@
+import type { DocumentReference } from 'firebase-admin/firestore'
+import { z } from 'zod'
+import type { AppBindings, DbRole } from '../config'
+import type { EntryId } from '../db'
+import type { AuthLevel, BaseRouteConfig, RouteMode } from '../router'
+import { StaleWorkerEpochError, type Duration } from '../utils'
+
+export { StaleWorkerEpochError }
+
+export type TaskStatus =
+  | 'running'
+  | 'sleeping'
+  | 'waiting_signal'
+  | 'compensating'
+  | 'completed'
+  | 'errored'
+  | 'cancelled'
+
+export interface TaskExecutionFence {
+  /** Monotonically increasing integer incremented on every worker claim, yield, hedge, or cancel transition */
+  epoch: number
+  /** Unique UUID for the currently owning worker execution attempt */
+  activeRunId: string | null
+  /** Timestamp of the latest worker heartbeat */
+  lastHeartbeatAt: number
+}
+
+export interface TaskProgressUpdate {
+  percent?: number
+  stage?: string
+  message?: string
+  data?: Record<string, unknown>
+}
+
+export interface TaskStatusRecord<O = unknown> {
+  taskId: string
+  workflowName: string
+  status: TaskStatus
+  userId?: EntryId
+  epoch?: number
+  activeRunId?: string | null
+  lastHeartbeatAt?: number
+  compensatedSteps?: Record<string, boolean>
+  progress?: TaskProgressUpdate
+  response?: O
+  error?: string
+  createdAt?: string
+  updatedAt?: string
+}
+
+export interface WorkflowContext {
+  taskId: string
+  workflowName: string
+  attempt: number
+  startedAt: number
+  deadlineMs: number
+  signal: AbortSignal
+}
+
+export interface StepRetryPolicy {
+  maxAttempts?: number
+  initialInterval?: Duration
+  backoffCoefficient?: number
+  maxInterval?: Duration
+  jitter?: boolean
+  nonRetryableErrorTypes?: string[]
+}
+
+export interface StepContext {
+  signal: AbortSignal
+  attempt: number
+  stepId: string
+}
+
+export interface StepOptions<T = unknown> {
+  retry?: number | StepRetryPolicy
+  timeout?: Duration
+  rollback?: (result: T) => Promise<unknown>
+}
+
+export interface StepMapOptions<T, R> extends StepOptions<R> {
+  concurrency?: number
+  shortCircuit?: (item: T, index: number) => Promise<R | undefined> | R | undefined
+}
+
+export type StepSpec<T> =
+  | ((ctx: StepContext) => Promise<T>)
+  | {
+      fn: (ctx: StepContext) => Promise<T>
+      rollback?: (result: T) => Promise<unknown>
+      retry?: number | StepRetryPolicy
+      timeout?: Duration
+    }
+
+export type StepFn = <T>(name: string, fnOrSpec: StepSpec<T>, options?: StepOptions<T>) => Promise<T>
+
+export interface WorkflowTargetRef<I = any, O = any> {
+  workflowName?: string
+  timeout?: Duration
+  handler?: (ctx: any) => Promise<O>
+  __input?: I
+}
+
+export interface StepRunner extends StepFn {
+  run: StepFn
+  map: <T, R>(
+    name: string,
+    items: readonly T[],
+    fn: (item: T, index: number, ctx: StepContext) => Promise<R>,
+    options?: StepMapOptions<T, R>
+  ) => Promise<R[]>
+  invoke: <I = unknown, O = unknown>(
+    name: string,
+    target: string | WorkflowTargetRef<I, O>,
+    input: I,
+    options?: { timeout?: Duration; dedupeKey?: string }
+  ) => Promise<O>
+  spawn: <I = unknown>(
+    name: string,
+    target: string | WorkflowTargetRef<I, unknown>,
+    input: I,
+    options?: { dedupeKey?: string; timeout?: Duration }
+  ) => Promise<{ taskId: string }>
+  sleep: (name: string, duration: Duration) => Promise<void>
+  waitForSignal: <T = unknown>(
+    name: string,
+    signalName: string,
+    options?: { timeout?: Duration }
+  ) => Promise<T>
+  continueAsNew: <I = unknown>(newInput: I) => Promise<never>
+  progress: (update: TaskProgressUpdate) => Promise<void>
+  flush: () => Promise<PromiseSettledResult<unknown>[]>
+  compensate: () => Promise<void>
+  onError: () => Promise<void>
+}
+
+export interface RollbackLeg<T> {
+  fn: () => Promise<T>
+  rollback?: (result: T) => Promise<unknown>
+}
+
+export interface WorkflowTaskPayload {
+  taskId: string
+  workflowName: string
+  params: Record<string, unknown>
+  isDev: boolean
+  expectedEpoch?: number
+}
+
+export interface WorkflowRuntimeEntry {
+  mode: RouteMode
+  deadlineSeconds?: number
+  runTask: (params: any, step: StepRunner) => Promise<unknown>
+}
+
+export class TaskYieldSignal extends Error {
+  readonly isTaskYield = true
+  constructor(
+    public readonly reason: 'sleeping' | 'waiting_signal' | 'continue_as_new' | 'yield',
+    public readonly meta: {
+      delaySeconds?: number
+      signalName?: string
+      timeoutMs?: number
+      newInput?: unknown
+    } = {}
+  ) {
+    super(`TaskYield:${reason}`)
+    this.name = 'TaskYieldSignal'
+  }
+}
+
+export interface ClaimCheckPointer {
+  __claimCheck: string
+}
+
+export class ReplayCompleteSignal extends Error {
+  readonly isReplayComplete = true as const
+  constructor() {
+    super('ReplayComplete')
+    this.name = 'ReplayCompleteSignal'
+  }
+}
+
+export interface ActiveTaskExecution {
+  taskId: string
+  step: StepRunner
+  controller: AbortController
+  claimedEpoch?: number
+}
+
+export interface ActiveWorkflowStore {
+  workflow: WorkflowContext
+  step: StepRunner
+  env?: AppBindings
+}
+
+export interface CreateStepRunnerOptions {
+  isDev?: boolean
+  signals?: Record<string, unknown>
+  workflowContext?: WorkflowContext
+  canYield?: boolean
+  claimedEpoch?: number
+  abortController?: AbortController
+  compensatedSteps?: Record<string, boolean>
+  replayOnly?: boolean
+}
+
+export interface DispatchWorkflowOptions {
+  mode?: RouteMode
+  timeout?: Duration
+  dedupeKey?: string
+  wait?: boolean
+  maxWaitMs?: number
+  signal?: AbortSignal
+  userId?: EntryId
+  userRole?: DbRole
+}
+
+export interface ExecuteWorkflowTaskOptions {
+  canYield?: boolean
+}
+
+export type TaskRecord<O = unknown> =
+  | { status: 'completed'; response: O }
+  | { status: 'errored'; error: string; errorCode?: string; nonRetryable?: boolean }
+  | { status: 'cancelled'; error?: string }
+
+export class TaskWaitDetachedSignal extends Error {
+  readonly isTaskWaitDetached = true as const
+  constructor(public readonly taskId: string) {
+    super(`TaskWaitDetached:${taskId}`)
+    this.name = 'TaskWaitDetachedSignal'
+  }
+}
+
+export interface PendingEntry {
+  taskId: string
+  workflowName: string
+  docRef: DocumentReference
+  createdAt: number
+  hedged: boolean
+  hedgeable?: boolean
+  keepAlive?: boolean
+  unsub?: () => void
+  dispatch: () => Promise<void>
+  finish: (res: TaskRecord<any>) => void
+}
+
+export interface WatchTaskOptions {
+  isDev?: boolean
+  docRef?: DocumentReference
+  workflowName?: string
+  hedgeable?: boolean
+  keepAlive?: boolean
+  timeoutMs?: number
+  maxWaitMs?: number
+  hedgeAfterMs?: number
+  pollIntervalMs?: number
+  signal?: AbortSignal
+  dispatch?: () => Promise<void>
+}
+
+export interface RouteRegistryItem {
+  id: string
+  mode: RouteMode
+  method: string
+  path: string
+  workflowName?: string
+  auth?: AuthLevel
+  description: string
+  mcp?: boolean
+  tool?: boolean
+  input?: z.ZodType<any>
+  output?: z.ZodType<any>
+  guard?: BaseRouteConfig['guard']
+}
+
+export interface RouteIntrospectionItem {
+  id: string
+  method: string
+  path: string
+  mode: 'rest' | 'workflow'
+  workflowName?: string
+  auth: string | false
+  description: string
+  mcp: boolean
+  tool: boolean
+  guards: string[]
+  inputSchema: Record<string, unknown> | null
+  outputSchema: Record<string, unknown> | null
+}
+
+export interface RoutesIntrospectionResponse {
+  appName: string
+  generatedAt: string
+  routes: RouteIntrospectionItem[]
+  collections?: Record<
+    string,
+    {
+      relationType: string
+      rbac?: Record<string, unknown>
+      schema: Record<string, unknown> | null
+    }
+  >
+}
+
+export const signalTaskBody = z.object({
+  signal: z.string().min(1),
+  payload: z.unknown().optional()
+})

package/src/tasks/service.ts

@@ -0,0 +1,254 @@
+import { Timestamp, type DocumentReference } from 'firebase-admin/firestore'
+import { CORE_CONFIG } from '../config'
+import type { Core } from '../index'
+import {
+  StaleWorkerEpochError,
+  type StepRunner,
+  type TaskStatusRecord,
+  type WorkflowContext,
+  type WorkflowRuntimeEntry
+} from './schema'
+import { TaskStepsEngine } from './steps'
+import { resolveTaskDeadline } from './utils'
+import { TaskWatcherEngine } from './watcher'
+import { TaskWorkerEngine } from './worker'
+
+export class TasksService {
+  readonly #core: Core
+  readonly #worker: TaskWorkerEngine
+  readonly #steps: TaskStepsEngine
+  readonly #watcher: TaskWatcherEngine
+
+  constructor(coreClient: Core) {
+    this.#core = coreClient
+    this.#worker = new TaskWorkerEngine(coreClient)
+    this.#steps = new TaskStepsEngine(coreClient, this.#worker.invokeChildWorkflow)
+    this.#watcher = new TaskWatcherEngine(coreClient, this.#steps)
+    this.#worker.bindEngines(this.#steps, this.#watcher)
+  }
+
+  #useStep = (): StepRunner => {
+    const store = this.#worker.getActiveWorkflowStore()
+    return store?.step ?? this.#steps.createStepRunner('workflow')
+  }
+
+  #assertTaskOwnership = (data: Record<string, unknown> | undefined): void => {
+    if (!data) throw this.#core.errors.httpError(404, 'Task not found', 'NOT_FOUND')
+    const caller = this.#core.caller
+    if (!caller || caller.role === 'admin' || caller.userId === CORE_CONFIG.auth.adminUserId) return
+    if (typeof data.userId === 'string' && data.userId !== caller.userId) {
+      throw this.#core.errors.httpError(403, 'Forbidden', 'FORBIDDEN')
+    }
+  }
+
+  status = async <O = unknown>(taskId: string): Promise<TaskStatusRecord<O>> => {
+    const isDev = Boolean(this.#core.env.isDev)
+    const snap = await this.#core.db.firestore(isDev).collection('_tasks').doc(taskId).get()
+    const data = snap.exists ? (snap.data() as Record<string, unknown>) : undefined
+    this.#assertTaskOwnership(data)
+    const response =
+      'response' in data! ? await this.#steps.maybeHydrateStepValue<O>(data!.response, isDev) : undefined
+    return {
+      taskId,
+      workflowName: String(data!.workflowName ?? ''),
+      status: (data!.status as TaskStatusRecord['status']) ?? 'running',
+      ...(typeof data!.userId === 'string' ? { userId: data!.userId } : {}),
+      ...(typeof data!.epoch === 'number' ? { epoch: data!.epoch } : {}),
+      ...('activeRunId' in data! ? { activeRunId: (data!.activeRunId as string | null) ?? null } : {}),
+      ...(typeof data!.lastHeartbeatAt === 'number' ? { lastHeartbeatAt: data!.lastHeartbeatAt } : {}),
+      ...(data!.compensatedSteps
+        ? { compensatedSteps: data!.compensatedSteps as Record<string, boolean> }
+        : {}),
+      ...(data!.progress ? { progress: data!.progress as TaskStatusRecord['progress'] } : {}),
+      ...('response' in data! ? { response: response as O } : {}),
+      ...(typeof data!.error === 'string' ? { error: data!.error } : {}),
+      ...(typeof data!.createdAt === 'string' ? { createdAt: data!.createdAt } : {}),
+      ...(typeof data!.updatedAt === 'string' ? { updatedAt: data!.updatedAt } : {})
+    }
+  }
+
+  signal = async (
+    taskId: string,
+    signalName: string,
+    payload?: unknown
+  ): Promise<{ taskId: string; signaled: true }> => {
+    const isDev = Boolean(this.#core.env.isDev)
+    const db = this.#core.db.firestore(isDev)
+    const docRef = db.collection('_tasks').doc(taskId)
+
+    const outcome = await db.runTransaction(async (tx) => {
+      const snap = await tx.get(docRef)
+      const data = snap.exists ? (snap.data() as Record<string, unknown>) : undefined
+      this.#assertTaskOwnership(data)
+
+      const workflowName = String(data!.workflowName ?? '')
+      const params = (data!.params as Record<string, unknown> | undefined) ?? {}
+      const signals = {
+        ...((data!.signals as Record<string, unknown> | undefined) ?? {}),
+        [signalName]: payload ?? true
+      }
+      const wasWaiting = data!.status === 'waiting_signal'
+      const currentEpoch = typeof data!.epoch === 'number' ? data!.epoch : 0
+      const nextEpoch = wasWaiting ? currentEpoch + 1 : currentEpoch
+
+      tx.set(
+        docRef,
+        {
+          signals,
+          ...(wasWaiting ? { status: 'running', waitingSignal: null, epoch: nextEpoch, activeRunId: null } : {}),
+          updatedAt: new Date().toISOString()
+        },
+        { merge: true }
+      )
+      return { workflowName, params, wasWaiting, nextEpoch }
+    })
+
+    if (outcome.wasWaiting) {
+      const wf = this.#worker.getRuntime().registry[outcome.workflowName]
+      const deadline = resolveTaskDeadline(outcome.workflowName, wf?.deadlineSeconds)
+      if (CORE_CONFIG.tasks.useCloudTasks) {
+        await this.#watcher.enqueueBatched(
+          {
+            taskId,
+            workflowName: outcome.workflowName,
+            params: outcome.params,
+            isDev,
+            expectedEpoch: outcome.nextEpoch
+          },
+          deadline
+        )
+      } else if (wf) {
+        void this.#worker
+          .execute({
+            taskId,
+            workflowName: outcome.workflowName,
+            params: outcome.params,
+            isDev,
+            expectedEpoch: outcome.nextEpoch
+          })
+          .catch(() => {})
+      }
+    }
+    return { taskId, signaled: true }
+  }
+
+  #replayUncompensatedSteps = async (
+    taskId: string,
+    docRef: DocumentReference,
+    freshData: Record<string, unknown>,
+    isDev: boolean
+  ): Promise<void> => {
+    const stepsMap = { ...((freshData.steps as Record<string, unknown> | undefined) ?? {}) }
+    const compMap = { ...((freshData.compensatedSteps as Record<string, boolean> | undefined) ?? {}) }
+    const workflowName = String(freshData.workflowName ?? '')
+    const wf: WorkflowRuntimeEntry | undefined = this.#worker.getRuntime().registry[workflowName]
+    if (!wf || !Object.keys(stepsMap).some((k) => !compMap[k])) return
+
+    const effectiveEnv = this.#core.env
+    const expiresAt = Timestamp.fromMillis(Date.now() + CORE_CONFIG.tasks.ttlMs)
+    const now = Date.now()
+    const wfCtx: WorkflowContext = {
+      taskId,
+      workflowName,
+      attempt: 1,
+      startedAt: now,
+      deadlineMs: now + 60_000,
+      signal: new AbortController().signal
+    }
+    const replayStep = this.#steps.createStepRunner(wf.mode, docRef, stepsMap, expiresAt, {
+      isDev,
+      workflowContext: wfCtx,
+      canYield: false,
+      compensatedSteps: compMap,
+      replayOnly: true
+    })
+    const params = (freshData.params as Record<string, unknown> | undefined) ?? {}
+    try {
+      await this.#worker.runWithWorkflowContext({ workflow: wfCtx, step: replayStep, env: effectiveEnv }, () =>
+        this.#core.internals.context.run(
+          { env: effectiveEnv, userId: this.#core.userId, role: this.#core.role, step: replayStep },
+          () => wf.runTask(params, replayStep)
+        )
+      )
+    } catch {}
+    await replayStep.compensate()
+  }
+
+  cancel = async (taskId: string): Promise<{ taskId: string; status: 'cancelled' }> => {
+    const isDev = Boolean(this.#core.env.isDev)
+    const db = this.#core.db.firestore(isDev)
+    const docRef = db.collection('_tasks').doc(taskId)
+    const preActiveExec = this.#worker.getActiveTaskExecution(taskId)
+
+    const txResult = await db.runTransaction(async (tx) => {
+      const snap = await tx.get(docRef)
+      const data = snap.exists ? (snap.data() as Record<string, unknown>) : undefined
+      this.#assertTaskOwnership(data)
+
+      const currentStatus = data!.status as string | undefined
+      if (currentStatus === 'completed' || currentStatus === 'errored' || currentStatus === 'cancelled') {
+        return { alreadyTerminal: true as const, data: data! }
+      }
+
+      const currentEpoch = typeof data!.epoch === 'number' ? data!.epoch : 0
+      const nextEpoch = currentEpoch + 1
+      tx.set(
+        docRef,
+        {
+          status: 'compensating',
+          epoch: nextEpoch,
+          activeRunId: null,
+          updatedAt: new Date().toISOString()
+        },
+        { merge: true }
+      )
+      return { alreadyTerminal: false as const, nextEpoch, data: data! }
+    })
+
+    if (!txResult.alreadyTerminal) {
+      const activeExec = preActiveExec ?? this.#worker.getActiveTaskExecution(taskId)
+      if (activeExec) {
+        activeExec.controller.abort(
+          new StaleWorkerEpochError(`Task ${taskId} cancelled by caller`, {
+            taskId,
+            claimedEpoch: activeExec.claimedEpoch,
+            actualEpoch: txResult.nextEpoch
+          })
+        )
+        await activeExec.step.compensate()
+      }
+
+      const freshSnap = await docRef.get().catch(() => null)
+      const freshData = (freshSnap?.data() as Record<string, unknown> | undefined) ?? txResult.data
+      await this.#replayUncompensatedSteps(taskId, docRef, freshData, isDev)
+
+      await docRef.set(
+        {
+          status: 'cancelled',
+          activeRunId: null,
+          error: 'Cancelled by caller',
+          updatedAt: new Date().toISOString()
+        },
+        { merge: true }
+      )
+    }
+
+    return { taskId, status: 'cancelled' }
+  }
+
+  get internals() {
+    return {
+      createStepRunner: this.#steps.createStepRunner,
+      allWithRollback: this.#steps.allWithRollback,
+      getContext: this.#worker.getContext,
+      useStep: this.#useStep,
+      getActiveWorkflowStore: this.#worker.getActiveWorkflowStore,
+      dispatch: this.#worker.dispatch,
+      execute: this.#worker.execute,
+      registerRuntime: this.#worker.registerRuntime,
+      getRuntime: this.#worker.getRuntime,
+      watchTask: this.#watcher.watchTask,
+      getActiveWatcherCount: this.#watcher.getActiveWatcherCount
+    }
+  }
+}