TS SDK: replace registerTask with Dag and registerDags

airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst

@@ -86,13 +86,14 @@ value routes the task to the Node.js coordinator.
 TypeScript implementation
 ~~~~~~~~~~~~~~~~~~~~~~~~~
 
-A task is an ordinary (usually ``async``) function receiving ``TaskHandlerArgs``. Register it with the
-``dag_id`` and ``task_id`` it implements, then start the coordinator runtime; the registrations and the
-top-level ``await startCoordinator()`` make the module a runnable bundle entry point.
+A task is an ordinary (usually ``async``) function receiving ``TaskHandlerArgs``. Create a ``Dag`` with
+the ``dag_id`` it implements, attach each handler with ``dag.task``, collect the Dags in a ``DagRegistry``,
+then serve them to Airflow with ``serveDags``; that top-level ``await`` makes the module a runnable bundle
+entry point.
 
 .. code-block:: typescript
 
-    import { registerTask, startCoordinator, type TaskHandlerArgs } from "@apache-airflow/ts-sdk";
+    import { Dag, DagRegistry, serveDags, type TaskHandlerArgs } from "@apache-airflow/ts-sdk";
 
     export async function buildMessage({ ctx, client }: TaskHandlerArgs) {
       const upstream = await client.getXCom<string>({
@@ -103,12 +104,22 @@ top-level ``await startCoordinator()`` make the module a runnable bundle entry p
       return `${greeting ?? "hello from TypeScript"}; upstream=${upstream ?? "missing"}`;
     }
 
-    registerTask({ dagId: "typescript_example", taskId: "build_message" }, buildMessage);
+    const dag = new Dag("typescript_example");
+    dag.task("build_message", buildMessage);
 
-    await startCoordinator();
+    await serveDags(new DagRegistry(dag));
 
-The ``dagId`` passed to ``registerTask`` must match the ``dag_id`` of the Python Dag, and each ``taskId``
-must match a ``@task.stub`` function in that Dag.
+The ``dagId`` passed to ``new Dag(...)`` must match the ``dag_id`` of the Python Dag, and each ``taskId``
+passed to ``dag.task`` must match a ``@task.stub`` function in that Dag. The registry passed to
+``serveDags`` is the bundle's complete set of Dags; a second ``serveDags`` call is rejected. A Dag left out
+of the registry is not part of the packed bundle, and its tasks are marked removed at runtime.
+
+``DagRegistry`` holds no sockets and starts nothing, so a unit test can build one and dispatch a handler
+through ``registry.getTaskHandler(dagId, taskId)`` without a coordinator runtime. A bundle that collects
+its Dags across several modules can add them incrementally with ``registry.register(...)``.
+
+``new Dag`` and ``dag.task`` take a trailing options object — ``spec`` on both, plus ``inputs`` on a task.
+These are not used yet; do not set them. Any other key is rejected.
 
 .. note::
 

ts-sdk/README.md

@@ -33,14 +33,17 @@ runtime used to execute registered TypeScript handlers from Airflow.
 ## Task Handlers
 
 ```ts
-import { registerTask, type TaskHandlerArgs } from "@apache-airflow/ts-sdk";
+import { Dag, DagRegistry, serveDags, type TaskHandlerArgs } from "@apache-airflow/ts-sdk";
 
 export async function sayHello({ ctx, client }: TaskHandlerArgs) {
   const greeting = await client.getVariable("greeting");
   return { message: `Hello from ${ctx.taskId}: ${greeting}` };
 }
 
-registerTask({ dagId: "example_dag", taskId: "say_hello" }, sayHello);
+const dag = new Dag("example_dag");
+dag.task("say_hello", sayHello);
+
+await serveDags(new DagRegistry(dag));
 ```
 
 Non-`undefined` return values are pushed to XCom under the `"return_value"`
@@ -95,7 +98,7 @@ Airflow metadata in the bundle itself.
 TypeScript entrypoint:
 
 ```ts
-import { registerTask, startCoordinator, type TaskHandlerArgs } from "@apache-airflow/ts-sdk";
+import { Dag, DagRegistry, serveDags, type TaskHandlerArgs } from "@apache-airflow/ts-sdk";
 
 export async function extract({ client }: TaskHandlerArgs) {
   const connection = await client.getConnection("sales_db");
@@ -118,33 +121,49 @@ export async function transform({ client }: TaskHandlerArgs) {
   };
 }
 
-registerTask({ dagId: "sales_pipeline", taskId: "extract" }, extract);
-registerTask({ dagId: "sales_pipeline", taskId: "transform" }, transform);
+const salesPipeline = new Dag("sales_pipeline");
+salesPipeline.task("extract", extract);
+salesPipeline.task("transform", transform);
 
-await startCoordinator();
+await serveDags(new DagRegistry(salesPipeline));
 ```
 
 The Python stub defines the Dag dependency graph. The TypeScript handler does
-the work and uses `TaskClient` for task-time Airflow data access. Register each
-handler with the Python Dag's `dag_id` and the stub task's `task_id`. The
-handler function is the reusable task implementation; `registerTask` binds that
-handler to a Python stub Dag/task identity for coordinator mode.
+the work and uses `TaskClient` for task-time Airflow data access. Create a
+`Dag` with the Python Dag's `dag_id` and attach each handler with the stub
+task's `task_id`. The handler function is the reusable task implementation;
+`dag.task` binds that handler to a Python stub task identity, a `DagRegistry`
+collects the Dags this bundle can execute, and `serveDags` serves them to
+Airflow.
+
+`serveDags` is the entrypoint, and the registry it is given is the whole bundle:
+a Dag left out of the registry is not part of the bundle, and its tasks are
+marked removed at runtime. The registry itself holds no sockets and starts
+nothing, so a unit test can build one and dispatch through
+`registry.getTaskHandler(dagId, taskId)` without any runtime involved.
 
-For larger projects, keep one Airflow entrypoint that imports every module that
-registers tasks, then starts the coordinator:
+`new Dag` and `dag.task` take a trailing options object — `spec` on both, plus
+`inputs` on a task. These are not used yet; do not set them.
+
+For larger projects, declare each Dag in its own module and keep one Airflow
+entrypoint that serves them all:
 
 ```ts
-import "./sales/tasks";
-import "./billing/tasks";
-import { startCoordinator } from "@apache-airflow/ts-sdk";
+import { salesDag } from "./sales/dag";
+import { billingDag } from "./billing/dag";
+import { DagRegistry, serveDags } from "@apache-airflow/ts-sdk";
 
-await startCoordinator();
+await serveDags(new DagRegistry(salesDag, billingDag));
 ```
 
+A bundle that collects its Dags across several modules can add them
+incrementally with `registry.register(...)` instead of passing them all to the
+constructor.
+
 Airflow launches the bundled entrypoint with `--comm=host:port` and
-`--logs=host:port`. `startCoordinator()` connects to those sockets, receives
-the task startup message, finds the registered handler for the Dag/task pair,
-and reports the terminal task state back to Airflow.
+`--logs=host:port`. `serveDags()` connects to those sockets, receives the task
+startup message, finds the registered handler for the Dag/task pair, and
+reports the terminal task state back to Airflow.
 
 See [`example/`](example/) for a coordinator-runtime example that packs a
 bundle with `airflow-ts-pack` and uses a Python stub Dag.

ts-sdk/docs/index.md

@@ -34,19 +34,23 @@ The SDK is currently distributed as source in the `ts-sdk/` directory of the
 Apache Airflow repository. Build it there and add it as a local dependency of
 your task bundle; it is not yet published to a public npm registry.
 
-Register a task handler. Handlers receive a `TaskContext` and a `TaskClient`;
-any non-`undefined` return value is pushed to XCom under the `"return_value"`
-key by the active runtime, matching Python `@task` behavior:
+Define a Dag and register its task handlers. Handlers receive a `TaskContext`
+and a `TaskClient`; any non-`undefined` return value is pushed to XCom under
+the `"return_value"` key by the active runtime, matching Python `@task`
+behavior:
 
 ```ts
-import { registerTask, type TaskHandlerArgs } from "@apache-airflow/ts-sdk";
+import { Dag, DagRegistry, serveDags, type TaskHandlerArgs } from "@apache-airflow/ts-sdk";
 
 export async function sayHello({ ctx, client }: TaskHandlerArgs) {
   const greeting = await client.getVariable("greeting");
   return { message: `Hello from ${ctx.taskId}: ${greeting}` };
 }
 
-registerTask({ dagId: "example_dag", taskId: "say_hello" }, sayHello);
+const dag = new Dag("example_dag");
+dag.task("say_hello", sayHello);
+
+await serveDags(new DagRegistry(dag));
 ```
 
 ## Coordinators

ts-sdk/example/src/main.ts

@@ -17,9 +17,9 @@
  * under the License.
  */
 
-import { registerTask, startCoordinator, type TaskHandlerArgs } from "@apache-airflow/ts-sdk";
+import { Dag, DagRegistry, serveDags, type TaskHandlerArgs } from "@apache-airflow/ts-sdk";
 
-const DAG_ID = "typescript_example";
+const dag = new Dag("typescript_example");
 
 export async function buildMessage({ client }: TaskHandlerArgs) {
   const upstream = await client.getXCom<string>({
@@ -49,7 +49,7 @@ export async function readConnection({ client }: TaskHandlerArgs) {
   };
 }
 
-registerTask({ dagId: DAG_ID, taskId: "build_message" }, buildMessage);
-registerTask({ dagId: DAG_ID, taskId: "read_connection" }, readConnection);
+dag.task("build_message", buildMessage);
+dag.task("read_connection", readConnection);
 
-await startCoordinator();
+await serveDags(new DagRegistry(dag));

ts-sdk/src/cli/pack.ts

@@ -129,9 +129,16 @@ function readBundleManifest(bundlePath: string): BundleManifest {
       maxBuffer: MANIFEST_MAX_BUFFER_BYTES,
     });
   } catch (error) {
-    throw new Error(`Running the bundle with ${AIRFLOW_METADATA_FLAG} failed: ${String(error)}`, {
-      cause: error,
-    });
+    const stderr = (error as { stderr?: string }).stderr ?? "";
+    const reported = stderr
+      .split("\n")
+      .reverse()
+      .find((line) => /^\w*Error: /.test(line.trim()))
+      ?.trim();
+    throw new Error(
+      reported ?? `Running the bundle with ${AIRFLOW_METADATA_FLAG} failed: ${String(error)}`,
+      { cause: error },
+    );
   }
 
   // Import-time logging from user code lands on stdout too; pick the sentinel line.
@@ -143,20 +150,52 @@ function readBundleManifest(bundlePath: string): BundleManifest {
     throw new Error(`Bundle produced no ${AIRFLOW_METADATA_FLAG} output`);
   }
 
-  let manifest: BundleManifest;
+  let parsed: unknown;
   try {
-    manifest = JSON.parse(line.slice(AIRFLOW_METADATA_SENTINEL.length)) as BundleManifest;
+    parsed = JSON.parse(line.slice(AIRFLOW_METADATA_SENTINEL.length));
   } catch (error) {
     throw new Error(`Bundle produced invalid ${AIRFLOW_METADATA_FLAG} output: ${String(error)}`, {
       cause: error,
     });
   }
-  if (!manifest.supervisor_schema_version || !manifest.dags || typeof manifest.dags !== "object") {
+  if (!isBundleManifest(parsed)) {
     throw new Error(`Bundle produced incomplete ${AIRFLOW_METADATA_FLAG} output`);
   }
+  const manifest = parsed;
+  // The line is whatever the bundle printed and nothing downstream re-validates
+  // it, so check each Dag entry down to the task-id element.
+  for (const [dagId, dag] of Object.entries(manifest.dags)) {
+    if (dag == null || !isTaskIdList(dag.tasks)) {
+      throw new Error(
+        `Bundle produced ${AIRFLOW_METADATA_FLAG} output with a malformed entry for Dag "${dagId}"`,
+      );
+    }
+  }
   return manifest;
 }
 
+// The document is checked before anything is read off it: JSON.parse also yields
+// null and primitives, and `null.supervisor_schema_version` would surface as a
+// raw TypeError rather than a report about the bundle.
+function isBundleManifest(value: unknown): value is BundleManifest {
+  if (typeof value !== "object" || value === null || Array.isArray(value)) return false;
+  const { supervisor_schema_version: version, dags } = value as Partial<BundleManifest>;
+  return (
+    // Rendered into the manifest verbatim, where the schema requires a non-empty
+    // string, so a truthy number or boolean would travel to Airflow as-is.
+    typeof version === "string" &&
+    version.length > 0 &&
+    typeof dags === "object" &&
+    dags !== null &&
+    // An array would pass the typeof check and yield Dags named "0", "1", ...
+    !Array.isArray(dags)
+  );
+}
+
+function isTaskIdList(value: unknown): value is string[] {
+  return Array.isArray(value) && value.every((item) => typeof item === "string" && item.length > 0);
+}
+
 // esbuild keeps an entry hashbang as line 1, where the metadata comment must go;
 // NodeCoordinator always runs the bundle through `node`, so drop it.
 function stripShebang(bundle: string): string {
@@ -193,10 +232,16 @@ export async function runPack(argv: readonly string[]): Promise<void> {
     });
 
     const manifest = readBundleManifest(stagingPath);
-    if (Object.keys(manifest.dags).length === 0) {
-      throw new Error(
-        `${args.entry} registered no tasks; call registerTask(...) before startCoordinator()`,
-      );
+    const dagEntries = Object.entries(manifest.dags);
+    if (dagEntries.length === 0) {
+      throw new Error(`${args.entry} served no Dags; pass them to serveDags(new DagRegistry(...))`);
+    }
+    // Warn rather than fail, as airflow-go-pack does: the shared schema allows a
+    // Dag with no tasks.
+    for (const [dagId, dag] of dagEntries) {
+      if (dag.tasks.length === 0) {
+        process.stderr.write(`warning: dag ${JSON.stringify(dagId)} has no tasks\n`);
+      }
     }
 
     const metadataYaml = renderMetadataYaml({

ts-sdk/src/coordinator/index.ts

@@ -21,8 +21,11 @@
 //
 // TaskClient and related types are exported from the package root. This
 // barrel only exports coordinator-specific entry points.
+//
+// `startCoordinator` is deliberately not exported: Dag authors reach the
+// runtime through `serveDags()`, and never name the coordinator itself.
 
-export { startCoordinator, type StartCoordinatorOptions } from "./runtime.js";
+export { serveDags } from "./runtime.js";
 /** Cadwyn schema version this SDK was generated against. Not sent on
  *  the wire — exposed so callers can read it for bundle metadata,
  *  health checks, or to confirm which schema their build is pinned to. */

ts-sdk/src/coordinator/manifest.ts

@@ -18,27 +18,57 @@
  */
 
 import { SUPERVISOR_API_VERSION } from "./protocol.js";
-import { listRegisteredTasks, type TaskRegistration } from "../sdk/registry.js";
+import { listRegistryDags, type DagRegistry } from "../sdk/registry.js";
 
 export const AIRFLOW_METADATA_FLAG = "--airflow-metadata";
 
 /** Marks the manifest line on stdout, which import-time logging may also reach. */
 export const AIRFLOW_METADATA_SENTINEL = "__AIRFLOW_METADATA__ ";
 
+// Mirrors the Python task-SDK KEY_REGEX and validate_key in airflow.sdk.definitions._internal.node.
+// Checked here rather than in Dag()/task() registration: that runs on every
+// bundle module load, including at task-execution-runtime startup, whereas
+// this only runs once, when the bundle is packed.
+const KEY_REGEX = /^[\p{L}\p{N}_.-]+$/u;
+const MAX_KEY_LENGTH = 250;
+
+function validateKey(label: string, value: string): void {
+  if (typeof value !== "string" || !KEY_REGEX.test(value)) {
+    throw new Error(
+      `${label} must be made of alphanumeric characters, dashes, dots, and underscores`,
+    );
+  }
+  if (value.length > MAX_KEY_LENGTH) {
+    throw new Error(`${label} must be less than ${MAX_KEY_LENGTH} characters, not ${value.length}`);
+  }
+}
+
 /** Bundle manifest fields only the built bundle itself knows: the schema
  *  version it was compiled against and the Dag/task pairs it registered.
- *  `airflow-ts-pack` runs `node bundle.mjs --airflow-metadata` to read this. */
+ *  Registered Dags without tasks appear with an empty `tasks` list so
+ *  `airflow-ts-pack` (which runs `node bundle.mjs --airflow-metadata` to
+ *  read this) can warn about them instead of silently dropping them. */
 export interface BundleManifest {
   supervisor_schema_version: string;
   dags: Record<string, { tasks: string[] }>;
 }
 
-export function buildBundleManifest(
-  registrations: readonly TaskRegistration[] = listRegisteredTasks(),
-): BundleManifest {
+export function buildBundleManifest(registry: DagRegistry): BundleManifest {
   const dags: BundleManifest["dags"] = {};
-  for (const { dagId, taskId } of registrations) {
-    (dags[dagId] ??= { tasks: [] }).tasks.push(taskId);
+  for (const { dagId, tasks } of listRegistryDags(registry)) {
+    validateKey(`Dag "${dagId}"`, dagId);
+    for (const taskId of tasks) {
+      validateKey(`Task "${taskId}" of Dag "${dagId}"`, taskId);
+    }
+    Object.defineProperty(dags, dagId, {
+      configurable: true,
+      enumerable: true,
+      value: { tasks: [...tasks] },
+      writable: true,
+    });
   }
-  return { supervisor_schema_version: SUPERVISOR_API_VERSION, dags };
+  return {
+    supervisor_schema_version: SUPERVISOR_API_VERSION,
+    dags,
+  };
 }