Validate Dag and task IDs in the ts-sdk task registry (#69400)

guan404ming · web-flow · commit 5ac8c33b2182 · 2026-07-10T16:04:29.000+08:00
* Validate Dag and task IDs in the ts-sdk task registry

* Validate Dag and task ID length in the ts-sdk task registry
diff --git a/ts-sdk/src/sdk/registry.ts b/ts-sdk/src/sdk/registry.ts
@@ -19,6 +19,21 @@
 
 import type { TaskHandler } from "./task.js";
 
+// Mirrors the Python task-SDK KEY_REGEX and validate_key in airflow.sdk.definitions._internal.node.
+const KEY_REGEX = /^[\p{L}\p{N}_.-]+$/u;
+const MAX_KEY_LENGTH = 250;
+
+function validateKey(name: string, value: string): void {
+  if (typeof value !== "string" || !KEY_REGEX.test(value)) {
+    throw new Error(
+      `${name} must be made of alphanumeric characters, dashes, dots, and underscores`,
+    );
+  }
+  if (value.length > MAX_KEY_LENGTH) {
+    throw new Error(`${name} must be less than ${MAX_KEY_LENGTH} characters, not ${value.length}`);
+  }
+}
+
 /** Identifies the Airflow task handled by a TypeScript function. */
 export interface TaskRegistration {
   /** Identifier of the Dag containing this task. */
@@ -39,12 +54,8 @@ export class TaskRegistry {
    */
   register<TReturn = unknown>(registration: TaskRegistration, handler: TaskHandler<TReturn>): void {
     const { dagId, taskId } = registration;
-    if (!dagId || typeof dagId !== "string") {
-      throw new Error("dagId must be a non-empty string");
-    }
-    if (!taskId || typeof taskId !== "string") {
-      throw new Error("taskId must be a non-empty string");
-    }
+    validateKey("dagId", dagId);
+    validateKey("taskId", taskId);
     if (typeof handler !== "function") {
       throw new Error(`handler for Dag "${dagId}" task "${taskId}" must be a function`);
     }
diff --git a/ts-sdk/tests/sdk/registry.test.ts b/ts-sdk/tests/sdk/registry.test.ts
@@ -73,14 +73,51 @@ describe("registry", () => {
     const registry = new TaskRegistry();
     expect(() =>
       registry.register({ dagId: "", taskId: "my_task" }, async () => undefined),
-    ).toThrowError(/dagId must be a non-empty string/);
+    ).toThrowError(/dagId must be made of alphanumeric/);
   });
 
   it("rejects an empty taskId", () => {
     const registry = new TaskRegistry();
     expect(() =>
       registry.register({ dagId: "example_dag", taskId: "" }, async () => undefined),
-    ).toThrowError(/taskId must be a non-empty string/);
+    ).toThrowError(/taskId must be made of alphanumeric/);
+  });
+
+  it.each(["   ", "\t", "my dag", "a/b", "task@1"])(
+    "rejects a dagId with characters no Python dag_id allows: %j",
+    (dagId) => {
+      const registry = new TaskRegistry();
+      expect(() =>
+        registry.register({ dagId, taskId: "my_task" }, async () => undefined),
+      ).toThrowError(/dagId must be made of alphanumeric/);
+    },
+  );
+
+  it.each(["   ", "\t", "my task", "a/b", "task@1"])(
+    "rejects a taskId with characters no Python task_id allows: %j",
+    (taskId) => {
+      const registry = new TaskRegistry();
+      expect(() =>
+        registry.register({ dagId: "example_dag", taskId }, async () => undefined),
+      ).toThrowError(/taskId must be made of alphanumeric/);
+    },
+  );
+
+  it.each([
+    ["dagId", { dagId: "d".repeat(251), taskId: "my_task" }],
+    ["taskId", { dagId: "example_dag", taskId: "t".repeat(251) }],
+  ])("rejects a %s longer than 250 characters", (name, registration) => {
+    const registry = new TaskRegistry();
+    expect(() => registry.register(registration, async () => undefined)).toThrowError(
+      new RegExp(`${name} must be less than 250 characters, not 251`),
+    );
+  });
+
+  it("accepts a Unicode dagId that Python's word-character rule allows", () => {
+    const registry = new TaskRegistry();
+    const handler = async () => undefined;
+    registry.register({ dagId: "café_dag", taskId: "任務" }, handler);
+    expect(registry.get("café_dag", "任務")).toBe(handler);
   });
 
   it("rejects non-function handlers", () => {
