feat(entitlements): add entitlement/permission system for tasks and workflows (#370)

Adds a declarative entitlement system that mirrors the existing schema pattern:
static class-level declarations, dynamic instance overrides, change events,
and graph-level aggregation. This enables consumers to inspect required
permissions before workflow execution and optionally enforce them.

Author: sroussey

packages/ai/src/task/base/AiTask.ts

@@ -10,21 +10,24 @@
  */
 
 import { Job } from "@workglow/job-queue";
+import type {
+  IExecuteContext,
+  IExecuteReactiveContext,
+  TaskConfig,
+  TaskEntitlement,
+  TaskEntitlements,
+  TaskOutput,
+} from "@workglow/task-graph";
 import {
+  Entitlements,
   Task,
   TaskConfigSchema,
   TaskConfigurationError,
   TaskInput,
   hasStructuredOutput,
 } from "@workglow/task-graph";
-import type {
-  IExecuteContext,
-  IExecuteReactiveContext,
-  TaskConfig,
-  TaskOutput,
-} from "@workglow/task-graph";
-import type { DataPortSchema, JsonSchema } from "@workglow/util/schema";
 import type { ServiceRegistry } from "@workglow/util";
+import type { DataPortSchema, JsonSchema } from "@workglow/util/schema";
 
 import { AiJob, AiJobInput } from "../../job/AiJob";
 import { MODEL_REPOSITORY } from "../../model/ModelRegistry";
@@ -63,6 +66,35 @@ export class AiTask<
   Config extends TaskConfig<Input> = TaskConfig<Input>,
 > extends Task<Input, Output, Config> {
   public static override type: string = "AiTask";
+  public static override hasDynamicEntitlements: boolean = true;
+
+  public static override entitlements(): TaskEntitlements {
+    return {
+      entitlements: [{ id: Entitlements.AI_INFERENCE, reason: "Runs AI model inference" }],
+    };
+  }
+
+  public override entitlements(): TaskEntitlements {
+    const base: TaskEntitlement[] = [
+      { id: Entitlements.AI_INFERENCE, reason: "Runs AI model inference" },
+    ];
+    // Prefer runInputData.model (runtime) over defaults.model (construction-time)
+    const runModel = this.runInputData?.model;
+    const modelId =
+      typeof runModel === "string"
+        ? runModel
+        : typeof this.defaults.model === "string"
+          ? this.defaults.model
+          : undefined;
+    if (modelId) {
+      base.push({
+        id: Entitlements.AI_MODEL,
+        reason: `Uses model ${modelId}`,
+        resources: [modelId],
+      });
+    }
+    return { entitlements: base };
+  }
 
   public static override configSchema(): DataPortSchema {
     return aiTaskConfigSchema;

packages/task-graph/src/common.ts

@@ -7,6 +7,7 @@
 export * from "./task-graph/Dataflow";
 export * from "./task-graph/DataflowEvents";
 
+export * from "./task-graph/GraphEntitlementUtils";
 export * from "./task-graph/GraphFormatScanner";
 export * from "./task-graph/GraphSchemaUtils";
 export * from "./task-graph/ITaskGraph";

packages/task-graph/src/task-graph/GraphEntitlementUtils.ts

@@ -0,0 +1,106 @@
+/**
+ * @license
+ * Copyright 2025 Steven Roussey <sroussey@gmail.com>
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import {
+  type EntitlementId,
+  type TaskEntitlement,
+  type TaskEntitlements,
+  type TrackedTaskEntitlement,
+  type TrackedTaskEntitlements,
+  EMPTY_ENTITLEMENTS,
+  mergeEntitlementPair,
+} from "../task/TaskEntitlements";
+import type { TaskIdType } from "../task/TaskTypes";
+import { TaskStatus } from "../task/TaskTypes";
+import type { TaskGraph } from "./TaskGraph";
+
+// ========================================================================
+// Options
+// ========================================================================
+
+export interface GraphEntitlementOptions {
+  /**
+   * When true, annotate each entitlement with the source task IDs that require it.
+   */
+  readonly trackOrigins?: boolean;
+  /**
+   * Controls which ConditionalTask branches are included.
+   * - "all" (default): Include entitlements from ALL branches (conservative, pre-execution analysis)
+   * - "active": Only include entitlements from currently active branches (runtime, after conditions evaluated)
+   */
+  readonly conditionalBranches?: "all" | "active";
+}
+
+// ========================================================================
+// Graph Entitlement Computation
+// ========================================================================
+
+/**
+ * Computes the aggregated entitlements for a TaskGraph.
+ * Returns the union of all task entitlements in the graph.
+ *
+ * When `trackOrigins` is true, returns TrackedTaskEntitlements with source task IDs.
+ */
+export function computeGraphEntitlements(
+  graph: TaskGraph,
+  options?: GraphEntitlementOptions & { readonly trackOrigins: true }
+): TrackedTaskEntitlements;
+export function computeGraphEntitlements(
+  graph: TaskGraph,
+  options?: GraphEntitlementOptions
+): TaskEntitlements;
+export function computeGraphEntitlements(
+  graph: TaskGraph,
+  options?: GraphEntitlementOptions
+): TaskEntitlements | TrackedTaskEntitlements {
+  const tasks = graph.getTasks();
+  if (tasks.length === 0) return EMPTY_ENTITLEMENTS;
+
+  const trackOrigins = options?.trackOrigins ?? false;
+  const conditionalBranches = options?.conditionalBranches ?? "all";
+
+  // Accumulate entitlements by ID
+  const merged = new Map<
+    EntitlementId,
+    { entitlement: TaskEntitlement; sourceTaskIds: TaskIdType[] }
+  >();
+
+  for (const task of tasks) {
+    // For ConditionalTask with "active" mode, skip disabled tasks
+    if (conditionalBranches === "active" && task.status !== undefined) {
+      if (task.status === TaskStatus.DISABLED) continue;
+    }
+
+    const taskEntitlements = task.entitlements();
+    for (const entitlement of taskEntitlements.entitlements) {
+      const existing = merged.get(entitlement.id);
+      if (existing) {
+        // Merge: optional=false wins, resources are unioned
+        existing.entitlement = mergeEntitlementPair(existing.entitlement, entitlement);
+        if (trackOrigins) {
+          existing.sourceTaskIds.push(task.id);
+        }
+      } else {
+        merged.set(entitlement.id, {
+          entitlement,
+          sourceTaskIds: trackOrigins ? [task.id] : [],
+        });
+      }
+    }
+  }
+
+  if (merged.size === 0) return EMPTY_ENTITLEMENTS;
+
+  if (trackOrigins) {
+    const entitlements: TrackedTaskEntitlement[] = [];
+    for (const { entitlement, sourceTaskIds } of merged.values()) {
+      entitlements.push({ ...entitlement, sourceTaskIds });
+    }
+    return { entitlements };
+  }
+
+  return { entitlements: Array.from(merged.values()).map((e) => e.entitlement) };
+}

packages/task-graph/src/task-graph/TaskGraph.ts

@@ -4,17 +4,19 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { DirectedAcyclicGraph } from "@workglow/util/graph";
 import { EventEmitter, ServiceRegistry, uuid4 } from "@workglow/util";
+import { DirectedAcyclicGraph } from "@workglow/util/graph";
 import { TaskOutputRepository } from "../storage/TaskOutputRepository";
 import type { ITask } from "../task/ITask";
 import type { StreamEvent } from "../task/StreamTypes";
+import type { TaskEntitlements } from "../task/TaskEntitlements";
 import type { JsonTaskItem, TaskGraphJson, TaskGraphJsonOptions } from "../task/TaskJSON";
 import type { TaskIdType, TaskInput, TaskOutput, TaskStatus } from "../task/TaskTypes";
-import { ensureTask } from "./Conversions";
 import type { PipeFunction } from "./Conversions";
-import { Dataflow } from "./Dataflow";
+import { ensureTask } from "./Conversions";
 import type { DataflowIdType } from "./Dataflow";
+import { Dataflow } from "./Dataflow";
+import { computeGraphEntitlements } from "./GraphEntitlementUtils";
 import { addBoundaryNodesToDependencyJson, addBoundaryNodesToGraphJson } from "./GraphSchemaUtils";
 import type { ITaskGraph } from "./ITaskGraph";
 import {
@@ -27,8 +29,8 @@ import {
   TaskGraphStatusEvents,
   TaskGraphStatusListeners,
 } from "./TaskGraphEvents";
-import { CompoundMergeStrategy, GraphResult, TaskGraphRunner } from "./TaskGraphRunner";
 import type { GraphResultArray } from "./TaskGraphRunner";
+import { CompoundMergeStrategy, GraphResult, TaskGraphRunner } from "./TaskGraphRunner";
 
 /**
  * Configuration for running a task graph
@@ -57,9 +59,16 @@ export interface TaskGraphRunConfig {
    * Defaults to no limit. Set this to prevent runaway graph construction.
    */
   maxTasks?: number;
+  /**
+   * When true, check entitlements via the registered IEntitlementEnforcer before
+   * graph execution begins. Throws TaskEntitlementError if any required (non-optional)
+   * entitlements are denied. Default: false.
+   */
+  enforceEntitlements?: boolean;
 }
 
-export interface TaskGraphRunReactiveConfig extends TaskGraphRunConfig {
+export interface TaskGraphRunReactiveConfig
+  extends Omit<TaskGraphRunConfig, "enforceEntitlements" | "timeout"> {
   /** Optional service registry to use for this task graph */
   registry?: ServiceRegistry;
 }
@@ -563,6 +572,67 @@ export class TaskGraph implements ITaskGraph {
     };
   }
 
+  /**
+   * Subscribes to entitlement changes on all tasks (existing and future).
+   * When any task's entitlements change, the graph recomputes and emits its own
+   * `entitlementChange` event. Structural changes (task_added, task_removed) also trigger.
+   *
+   * @param callback - Function called with the aggregated entitlements whenever they change
+   * @returns a function to unsubscribe from all entitlement events
+   */
+  public subscribeToTaskEntitlements(
+    callback: (entitlements: TaskEntitlements) => void
+  ): () => void {
+    const globalUnsubs: (() => void)[] = [];
+    const taskUnsubs = new Map<TaskIdType, () => void>();
+
+    const emitChange = () => {
+      const entitlements = computeGraphEntitlements(this);
+      this.emit("entitlementChange", entitlements);
+      callback(entitlements);
+    };
+
+    const subscribeTask = (taskId: TaskIdType) => {
+      const task = this.getTask(taskId);
+      if (!task || typeof task.subscribe !== "function") return;
+      const unsub = task.subscribe("entitlementChange", () => emitChange());
+      taskUnsubs.set(taskId, unsub);
+    };
+
+    // Subscribe to entitlementChange events on all existing tasks
+    for (const task of this.getTasks()) {
+      subscribeTask(task.id);
+    }
+
+    // Emit the initial state immediately so subscribers don't miss the current entitlements
+    emitChange();
+
+    // Subscribe to new tasks being added
+    globalUnsubs.push(
+      this.subscribe("task_added", (taskId: TaskIdType) => {
+        subscribeTask(taskId);
+        emitChange();
+      })
+    );
+
+    globalUnsubs.push(
+      this.subscribe("task_removed", (taskId: TaskIdType) => {
+        const unsub = taskUnsubs.get(taskId);
+        if (unsub) {
+          unsub();
+          taskUnsubs.delete(taskId);
+        }
+        emitChange();
+      })
+    );
+
+    return () => {
+      globalUnsubs.forEach((unsub) => unsub());
+      taskUnsubs.forEach((unsub) => unsub());
+      taskUnsubs.clear();
+    };
+  }
+
   /**
    * Registers an event listener for the specified event
    * @param name - The event name to listen for

packages/task-graph/src/task-graph/TaskGraphEvents.ts

@@ -6,6 +6,7 @@
 
 import { EventParameters } from "@workglow/util";
 import type { StreamEvent } from "../task/StreamTypes";
+import type { TaskEntitlements } from "../task/TaskEntitlements";
 import { TaskIdType } from "../task/TaskTypes";
 import { DataflowIdType } from "./Dataflow";
 
@@ -26,6 +27,8 @@ export type TaskGraphStatusListeners = {
   task_stream_chunk: (taskId: TaskIdType, event: StreamEvent) => void;
   /** Fired when a task in the graph finishes streaming */
   task_stream_end: (taskId: TaskIdType, output: Record<string, any>) => void;
+  /** Fired when the aggregated entitlements of the graph change */
+  entitlementChange: (entitlements: TaskEntitlements) => void;
 };
 export type TaskGraphStatusEvents = keyof TaskGraphStatusListeners;
 export type TaskGraphStatusListener<Event extends TaskGraphStatusEvents> =