TS SDK: answer TaskHandlerParseRequest from bundle.serve

airflow-core/tests/unit/dag_processing/test_task_handler_processor_ts.py

@@ -0,0 +1,146 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Pack the TypeScript SDK example bundle and probe it for its task handlers."""
+
+from __future__ import annotations
+
+import json
+import os
+import shutil
+import subprocess
+from typing import TYPE_CHECKING
+from unittest import mock
+
+import pytest
+import structlog
+
+from airflow.dag_processing.processor import TaskHandlerDeclaration, TaskHandlerParsingResult
+from airflow.dag_processing.task_handler_processor import LangSDKTaskHandlerProcessorProcess
+from airflow.sdk.coordinators._subprocess import supports_task_handler_parsing
+from airflow.sdk.execution_time import supervisor
+from airflow.sdk.execution_time.coordinator import get_coordinator_manager, reset_coordinator_manager
+
+from tests_common.test_utils.config import conf_vars
+from tests_common.test_utils.paths import AIRFLOW_ROOT_PATH
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+    from pathlib import Path
+
+TS_SDK_PATH = AIRFLOW_ROOT_PATH / "ts-sdk"
+# The SDK's `engines` requirement.
+MIN_NODE_MAJOR = 22
+COORDINATORS = {
+    "ts": {
+        "classpath": "airflow.sdk.coordinators.node.NodeCoordinator",
+        "kwargs": {"node_executable": shutil.which("node")},
+    }
+}
+
+pytestmark = pytest.mark.skipif(
+    os.environ.get("AIRFLOW_LANG_SDK_REAL_PROBE_TESTS") != "1",
+    reason="set AIRFLOW_LANG_SDK_REAL_PROBE_TESTS=1 to pack and probe a real TypeScript bundle",
+)
+
+
+def _get_toolchain_problem() -> str | None:
+    """Return why this machine cannot pack the example bundle, or ``None`` when it can."""
+    if not TS_SDK_PATH.is_dir():
+        return "the TypeScript SDK sources are absent"
+    node = shutil.which("node")
+    if node is None or shutil.which("pnpm") is None:
+        return "needs node and pnpm"
+    try:
+        version = subprocess.run(
+            [node, "--version"], capture_output=True, text=True, check=True
+        ).stdout.strip()
+        major = int(version.lstrip("v").split(".")[0])
+    except (OSError, subprocess.CalledProcessError, ValueError) as e:
+        return f"cannot read the Node.js version: {e}"
+    if major < MIN_NODE_MAJOR:
+        return f"needs Node.js {MIN_NODE_MAJOR} or later, found {version}"
+    return None
+
+
+@pytest.fixture(scope="module")
+def example_bundle(tmp_path_factory) -> Path:
+    """Pack ``ts-sdk/example`` from a copy of the SDK sources, so the checkout gets no build output."""
+    if (problem := _get_toolchain_problem()) is not None:
+        pytest.skip(problem)
+    sdk = tmp_path_factory.mktemp("ts-sdk") / "ts-sdk"
+    shutil.copytree(TS_SDK_PATH, sdk, ignore=shutil.ignore_patterns("node_modules", "dist", ".pnpm-store"))
+    env = {**os.environ, "CI": "true", "COREPACK_ENABLE_DOWNLOAD_PROMPT": "0"}
+    # The example links the SDK, so it is installed again once the SDK is built.
+    for cwd, args in [
+        (sdk, ["install", "--frozen-lockfile"]),
+        (sdk, ["run", "build"]),
+        (sdk / "example", ["install"]),
+        (sdk / "example", ["run", "build"]),
+    ]:
+        subprocess.run(["pnpm", *args], cwd=cwd, env=env, check=True)
+    return sdk / "example" / "dist" / "bundle.min.mjs"
+
+
+@pytest.fixture
+def fresh_coordinator_manager() -> Iterator[None]:
+    reset_coordinator_manager()
+    yield
+    reset_coordinator_manager()
+
+
+def _declare(task_id: str) -> TaskHandlerDeclaration:
+    return TaskHandlerDeclaration(task_id=task_id, binding="named", params=None)
+
+
+@pytest.mark.usefixtures("fresh_coordinator_manager")
+@conf_vars({("sdk", "coordinators"): json.dumps(COORDINATORS)})
+@mock.patch.object(supervisor, "_should_use_exec", autospec=True, return_value=False)
+def test_probes_the_task_handlers_of_a_packed_typescript_bundle(mock_should_use_exec, example_bundle):
+    coordinator = get_coordinator_manager().get_coordinator("ts")
+    artifact = coordinator._find_task_handler_artifact(
+        bundle_path=example_bundle.parent, dag_id="typescript_example"
+    )
+    assert supports_task_handler_parsing(artifact.schema_version)
+    assert artifact.path == example_bundle.resolve()
+
+    result = LangSDKTaskHandlerProcessorProcess.run(
+        coordinator="ts",
+        path=artifact.path,
+        bundle_path=example_bundle.parent,
+        bundle_name="ts-task-handlers",
+        artifact_rel_path=os.fspath(artifact.path.relative_to(example_bundle.parent.resolve())),
+        logger=structlog.get_logger(),
+    )
+
+    # Types are erased, so each handler lists no params and only its presence is checked.
+    assert result == TaskHandlerParsingResult(
+        fileloc=os.fspath(artifact.path),
+        task_handlers={
+            "typescript_example": [
+                _declare("build_message"),
+                _declare("read_connection"),
+                _declare("write_and_delete_variable"),
+            ],
+            "typescript_taskflow_example": [
+                _declare("summarize"),
+                _declare("report"),
+                _declare("build_message"),
+            ],
+        },
+    )
+    assert list(result.task_handlers) == ["typescript_example", "typescript_taskflow_example"]

contributing-docs/30_new_language_sdk.rst

@@ -146,8 +146,8 @@ runs and its supervisor schema version, found in the Dag bundle at
 *bundle_path* itself, because ``self._get_scan_roots()`` is not set there. It
 raises ``FileNotFoundError`` when the task would find no artifact it can run,
 or ``ValueError`` when the artifact it found uses a supervisor schema version
-this Task SDK does not know (``ExecutableCoordinator`` skips such bundles, so
-it raises ``FileNotFoundError`` instead).
+this Task SDK does not know (``ExecutableCoordinator`` and ``NodeCoordinator``
+skip such bundles, so they raise ``FileNotFoundError`` instead).
 
 ``_build_parse_task_handler_command`` builds the command that starts the
 runtime for the artifact at *path*, which ``_find_task_handler_artifact``
@@ -162,7 +162,8 @@ not started. The runtime answers as described in
 Both defaults raise ``NotImplementedError``, so the stub tasks of a coordinator
 that does not implement them are not checked. ``ExecutableCoordinator``
 implements both for executable bundles. ``JavaCoordinator`` implements both for
-JARs, and its probe runs the same command as a task.
+JARs, and its probe runs the same command as a task. ``NodeCoordinator``
+implements both for packed ``*.min.mjs`` bundles.
 
 Supervisor Schema
 ~~~~~~~~~~~~~~~~~

task-sdk/src/airflow/sdk/coordinators/node/coordinator.py

@@ -111,6 +111,8 @@ class NodeCoordinator(SubprocessCoordinator):
     :param task_handler_bundle_name: Name of the Dag bundle searched recursively for the first
         verified ``*.min.mjs`` bundle declaring the task instance's Dag. It must be registered in
         ``[dag_processor] dag_bundle_config_list``. If unset, the task's own Dag bundle is used.
+        The ``*.min.mjs`` bundle found also answers a task handler parse request with the task
+        handlers it registers.
     :param task_startup_timeout: Maximum time the coordinator waits for a task
         process to start, in seconds. The default is 10 seconds.
 
@@ -140,3 +142,17 @@ def _build_dag_file_command(
 
     def _build_parse_dag_command(self, *, path: pathlib.Path) -> tuple[list[str], str | None]:
         return self._build_bundle_command(path)
+
+    def _find_task_handler_artifact(self, *, bundle_path: pathlib.Path, dag_id: str) -> ResolvedBundle:
+        """
+        Return the bundle a task of *dag_id* runs, with symlinks in its path resolved.
+
+        Compute its path in the Dag bundle relative to ``bundle_path.resolve()``, since *bundle_path*
+        may contain symlinks. A bundle with a supervisor schema version this Task SDK does not know
+        is skipped like any unusable bundle, so it leads to ``FileNotFoundError``, never ``ValueError``.
+        """
+        bundle = _Bundle.find([bundle_path], dag_id)
+        return ResolvedBundle(bundle.path.resolve(), bundle.schema_version)
+
+    def _build_parse_task_handler_command(self, *, path: pathlib.Path) -> tuple[list[str], str | None]:
+        return self._build_bundle_command(path)

task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py

@@ -34,7 +34,7 @@
 from uuid6 import uuid7
 
 from airflow.sdk.api.datamodels._generated import BundleInfo, TaskInstance
-from airflow.sdk.coordinators._subprocess import _PopenActivitySubprocess
+from airflow.sdk.coordinators._subprocess import TASK_HANDLER_PARSING_SCHEMA_VERSION, _PopenActivitySubprocess
 from airflow.sdk.coordinators.node import _bundle_reader as _reader
 from airflow.sdk.coordinators.node._bundle_reader import _digest_cache
 from airflow.sdk.coordinators.node.coordinator import NodeCoordinator, _Bundle
@@ -218,6 +218,128 @@ def test_tampered_bundle_raises(self, tmp_path):
             NodeCoordinator()._build_parse_dag_command(path=bundle)
 
 
+class TestNodeCoordinatorFindTaskHandlerArtifact:
+    def test_finds_the_bundle_a_task_of_the_dag_runs(self, tmp_path):
+        write_bundle(tmp_path, "other_dag", name="a.min.mjs")
+        write_bundle(tmp_path, "sales", name="b.min.mjs", schema_version="2026-10-30")
+        write_bundle(tmp_path, "sales", name="c.min.mjs")
+        coordinator = NodeCoordinator(node_executable="/opt/node/bin/node")
+        with coordinator._set_scan_roots([tmp_path]):
+            command, schema_version = coordinator._build_execute_task_command(what=_make_ti(dag_id="sales"))
+
+        artifact = coordinator._find_task_handler_artifact(bundle_path=tmp_path, dag_id="sales")
+
+        assert artifact.path == pathlib.Path(command[1]).resolve()
+        assert artifact.schema_version == schema_version == "2026-10-30"
+
+    def test_resolves_a_symlinked_bundle_root(self, tmp_path):
+        target = write_bundle(tmp_path / "real" / "team-a", "sales")
+        root = tmp_path / "bundle"
+        try:
+            root.symlink_to(tmp_path / "real", target_is_directory=True)
+        except (OSError, NotImplementedError):
+            pytest.skip("symlinks not supported on this platform")
+
+        artifact = NodeCoordinator()._find_task_handler_artifact(bundle_path=root, dag_id="sales")
+
+        assert artifact.path == target.resolve()
+        assert artifact.path.relative_to(root.resolve()) == pathlib.Path("team-a", BUNDLE_NAME)
+
+    @pytest.mark.parametrize(
+        ("dag_id", "schema_version", "message"),
+        [
+            pytest.param(
+                "other_dag",
+                SCHEMA_VERSION,
+                r"verified bundle declares dag_ids=\['other_dag'\]",
+                id="no-bundle-declares-the-dag",
+            ),
+            pytest.param(
+                "sales",
+                "1999-01-01",
+                "Version '1999-01-01' not found",
+                id="unknown-schema-version",
+            ),
+        ],
+    )
+    def test_raises_file_not_found_when_no_bundle_can_run_the_dag(
+        self, tmp_path, dag_id, schema_version, message
+    ):
+        write_bundle(tmp_path, dag_id, schema_version=schema_version)
+
+        with pytest.raises(FileNotFoundError, match=message):
+            NodeCoordinator()._find_task_handler_artifact(bundle_path=tmp_path, dag_id="sales")
+
+
+def _write_plain_file(root: pathlib.Path) -> pathlib.Path:
+    path = root / BUNDLE_NAME
+    path.write_bytes(b"export {};\n")
+    return path
+
+
+def _write_tampered_bundle(root: pathlib.Path) -> pathlib.Path:
+    bundle = write_bundle(root, "sales")
+    mutate_section(bundle, "code")
+    return bundle
+
+
+class TestNodeCoordinatorParseTaskHandlerCommand:
+    def test_returns_node_and_the_bundle_schema_version(self, tmp_path):
+        # A bundle that declares no Dag: the Dag processor names the bundle, so its Dags play no part.
+        bundle = write_bundle(tmp_path)
+        coordinator = NodeCoordinator(node_executable="/opt/node/bin/node")
+
+        command, schema_version = coordinator._build_parse_task_handler_command(path=bundle)
+
+        assert command == ["/opt/node/bin/node", str(bundle)]
+        assert schema_version == SCHEMA_VERSION
+
+    @pytest.mark.parametrize(
+        ("write", "error"),
+        [
+            pytest.param(_write_plain_file, "has no airflow bundle layout", id="not-a-bundle"),
+            pytest.param(_write_tampered_bundle, "code SHA-256 mismatch", id="tampered-code"),
+        ],
+    )
+    def test_rejects_a_bundle_that_fails_verification(self, tmp_path, write, error):
+        path = write(tmp_path)
+
+        with pytest.raises(ValueError, match=error):
+            NodeCoordinator()._build_parse_task_handler_command(path=path)
+
+    @patch("airflow.sdk.coordinators._subprocess._set_close_on_exec_above_stderr", autospec=True)
+    @patch("airflow.sdk.coordinators._subprocess._set_parent_death_signal", autospec=True)
+    @patch("airflow.sdk.coordinators._subprocess.signal.signal", autospec=True)
+    @patch(
+        "airflow.sdk.coordinators._subprocess.os.execvpe", autospec=True, side_effect=OSError("exec failed")
+    )
+    def test_the_probe_execs_what_a_task_of_the_dag_runs(
+        self, mock_execvpe, mock_signal, mock_death_signal, mock_close_on_exec, tmp_path
+    ):
+        write_bundle(tmp_path, "sales", schema_version=TASK_HANDLER_PARSING_SCHEMA_VERSION)
+        coordinator = NodeCoordinator(node_executable="/opt/node/bin/node")
+        with coordinator._set_scan_roots([tmp_path]):
+            command, _ = coordinator._build_execute_task_command(what=_make_ti(dag_id="sales"))
+        artifact = coordinator._find_task_handler_artifact(bundle_path=tmp_path, dag_id="sales")
+        reported: list[str | None] = []
+
+        with pytest.raises(OSError, match="exec failed"):
+            coordinator.parse_task_handler(
+                path=artifact.path,
+                bundle_path=tmp_path,
+                comm_address=("127.0.0.1", 1001),
+                logs_address=("127.0.0.1", 1002),
+                report_schema_version=reported.append,
+            )
+
+        assert reported == [TASK_HANDLER_PARSING_SCHEMA_VERSION]
+        # Node runs a symlinked bundle under its real path, so the resolved path runs the same module.
+        argv = [command[0], str(pathlib.Path(command[1]).resolve())]
+        mock_execvpe.assert_called_once_with(
+            argv[0], [*argv, "--comm=127.0.0.1:1001", "--logs=127.0.0.1:1002"], mock.ANY
+        )
+
+
 class TestBundleFind:
     @pytest.mark.parametrize(
         "name",

ts-sdk/README.md

@@ -287,7 +287,13 @@ Its values are Python names, which `tsc` cannot see and does not check.
 Airflow launches the bundled entrypoint with `--comm=host:port` and
 `--logs=host:port`. `bundle.serve()` 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.
+reports the terminal task state back to Airflow. To check the stub tasks of a
+Python Dag, the Dag processor also starts the bundle a stub task would run and
+asks which task handlers it registers. `bundle.serve()` then answers with every
+`TaskHandler`, keyed by Dag id, without running any handler. Types are erased,
+so the Dag processor checks only that each stub task has a handler, not its
+arguments. The tasks of a Dag declared in TypeScript are not task handlers, so
+the answer leaves them out.
 
 See [`example/`](https://github.com/apache/airflow/tree/main/ts-sdk/example) for
 a coordinator-runtime example that packs a bundle with `airflow-ts-pack` and

ts-sdk/src/coordinator/comm-channel.ts

@@ -22,8 +22,8 @@
 //
 // The channel is the sole reader on the socket. The task sends
 // requests and awaits id-correlated replies. The supervisor's only
-// unprompted frame is the greeting (StartupDetails /
-// DagFileParseRequest), pre-caught into the `greeting` promise that
+// unprompted frame is the greeting (StartupDetails, DagFileParseRequest
+// or TaskHandlerParseRequest), pre-caught into the `greeting` promise that
 // `connect()` awaits — the protocol sends nothing else
 // supervisor-initiated (comms.py: "No messages are sent to task
 // process except in response to a request").
@@ -35,9 +35,9 @@ import { Deferred } from "./deferred.js";
 import type { LogChannel } from "./log-channel.js";
 
 /** What `CommChannel.connect` resolves to: the live channel plus the
- *  supervisor's first frame (StartupDetails / DagFileParseRequest),
- *  already in hand so the caller never has to manage a "frame arrived
- *  with no consumer" window. */
+ *  supervisor's first frame (StartupDetails, DagFileParseRequest or
+ *  TaskHandlerParseRequest), already in hand so the caller never has to
+ *  manage a "frame arrived with no consumer" window. */
 export interface CommConnection {
   channel: CommChannel;
   firstFrame: Frame;