This is an automated email from the ASF dual-hosted git repository.
jason810496 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 5433ffe893c TS SDK: replace TaskHandlerArgs with getContext and
getClient (#73122)
5433ffe893c is described below
commit 5433ffe893c82af33abf07a70c57e0a61e754ef6
Author: Jason(Zhe-You) Liu <[email protected]>
AuthorDate: Thu Sep 17 10:03:14 2026 +0800
TS SDK: replace TaskHandlerArgs with getContext and getClient (#73122)
---
.../language-sdks/typescript.rst | 59 ++++---
ts-sdk/README.md | 53 ++++---
ts-sdk/api-docs/dag-authoring-api.ts | 5 +-
ts-sdk/docs/index.md | 17 ++-
ts-sdk/example/src/main.ts | 9 +-
ts-sdk/scripts/verify-package.mjs | 12 +-
ts-sdk/src/coordinator/runtime.ts | 17 ++-
ts-sdk/src/index.ts | 3 +-
ts-sdk/src/sdk/dag.ts | 28 ++--
ts-sdk/src/sdk/registry.ts | 14 +-
ts-sdk/src/sdk/task.ts | 97 ++++++++++--
ts-sdk/tests/coordinator/integration.test.ts | 38 ++---
ts-sdk/tests/public-api.test.ts | 27 +++-
ts-sdk/tests/sdk/dag.test.ts | 2 +-
ts-sdk/tests/sdk/task-scope.test.ts | 170 +++++++++++++++++++++
15 files changed, 421 insertions(+), 130 deletions(-)
diff --git
a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst
b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst
index 38616832436..a1d1137dd7f 100644
--- a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst
+++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst
@@ -37,7 +37,7 @@ The SDK is the ``apache-airflow-ts-sdk`` package (ESM-only).
It is currently in
.. seealso::
- For the full TypeScript API reference (``Dag``, ``DagRegistry``,
``serveDags``, task handlers,
+ For the full TypeScript API reference (``Dag``, ``DagRegistry``,
``serveDags``, the task handler getters,
``TaskClient``, supporting types, and exceptions),
see the `TypeScript SDK API reference
<https://airflow.apache.org/docs/ts-sdk/stable/>`__.
@@ -93,16 +93,19 @@ value routes the task to the Node.js coordinator.
TypeScript implementation
~~~~~~~~~~~~~~~~~~~~~~~~~
-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.
+A task is an ordinary (usually ``async``) function taking no arguments:
+``getContext()`` and ``getClient()`` reach the runtime from inside the call,
so nothing the SDK supplies is a parameter.
+
+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 { Dag, DagRegistry, serveDags, type TaskHandlerArgs } from
"apache-airflow-ts-sdk";
+ import { Dag, DagRegistry, getClient, serveDags } from
"apache-airflow-ts-sdk";
- export async function buildMessage({ ctx, client }: TaskHandlerArgs) {
+ export async function buildMessage() {
+ const client = getClient();
const upstream = await client.getXCom<string>({
key: "return_value",
taskId: "python_start",
@@ -125,14 +128,14 @@ of the registry is not part of the packed bundle, and its
tasks are marked remov
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.
+``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::
As with the other language SDKs, XCom *dependencies* are declared in the
Python stub Dag (they define task
- order). The value must still be read explicitly in TypeScript via
``client.getXCom``, and produced either
- by the task's return value or by ``client.setXCom``.
+ order). The value must still be read explicitly in TypeScript via
``getClient().getXCom``, and produced
+ either by the task's return value or by ``getClient().setXCom``.
Coordinator configuration
~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -169,22 +172,30 @@ task instance.
Writing tasks
-------------
-Every task handler receives a single ``TaskHandlerArgs`` object:
+A task handler takes no SDK-supplied argument. Two getters, valid for as long
as the handler runs, supply what it needs:
.. list-table::
:header-rows: 1
- :widths: 15 85
+ :widths: 20 80
- * - Field
+ * - Getter
- Value
- * - ``ctx``
- - The task's execution context: ``dagId``, ``taskId`` (including any
TaskGroup prefix), ``runId``,
- ``tryNumber``, ``mapIndex`` (``-1`` for an unmapped task), and
``signal`` — an ``AbortSignal`` that
- fires when Airflow terminates the task. Pass ``signal`` to ``fetch()``,
timers, or other APIs that
- accept an ``AbortSignal`` for cooperative cancellation.
- * - ``client``
+ * - ``getContext()``
+ - The task's execution context (a ``TaskContext``): ``dagId``, ``taskId``
(including any TaskGroup
+ prefix), ``runId``, ``tryNumber``, ``mapIndex`` (``-1`` for an unmapped
task), and ``signal``, an
+ ``AbortSignal`` that fires when Airflow terminates the task. Pass
``signal`` to ``fetch()``, timers,
+ or other APIs that accept an ``AbortSignal`` for cooperative
cancellation.
+ * - ``getClient()``
- A ``TaskClient`` for Airflow Variables, Connections, and XCom.
+Both read a store the runtime installs around the handler call,
+which follows the handler across every ``await`` and into every promise it
creates,
+so a helper several frames deep reads them without being passed anything. Both
throw outside a handler.
+
+Work that outlives the handler is the one gap.
+A promise the handler never awaits still resolves, but it runs after Airflow
has been told the task's terminal state,
+so await everything a handler starts.
+
A non-``undefined`` return value becomes the task's ``return_value`` XCom,
matching Python ``@task``
behavior. An uncaught exception (or rejected promise) marks the task instance
failed in Airflow, triggering
retries if configured on the stub.
@@ -192,18 +203,18 @@ retries if configured on the stub.
The ``TaskClient`` surface
~~~~~~~~~~~~~~~~~~~~~~~~~~
-* ``getVariable(key)`` — returns the Variable as a string, or ``null`` when it
is missing;
+* ``getVariable(key)`` returns the Variable as a string, or ``null`` when it
is missing;
``getVariableOrThrow(key)`` throws ``VariableNotFoundError`` instead,
matching Python ``Variable.get``
with no default.
-* ``getConnection(connId)`` — returns a ``ConnectionResult`` with fields
``id`` and ``type``, plus the
+* ``getConnection(connId)`` returns a ``ConnectionResult`` with fields ``id``
and ``type``, plus the
optional fields ``host``, ``schema``, ``login``, ``password``, ``port``, and
``extra`` (each may be
missing or ``null``), or ``null`` when the connection does not exist;
``getConnectionOrThrow(connId)`` throws ``ConnectionNotFoundError`` instead,
matching Python
``BaseHook.get_connection``.
-* ``getXCom<T>({key, ...})`` — reads an XCom value, or ``null`` when it is
missing. The locator fields
+* ``getXCom<T>({key, ...})`` reads an XCom value, or ``null`` when it is
missing. The locator fields
(``dagId``, ``runId``, ``taskId``, ``mapIndex``) default to the current
task; pass ``taskId`` to read an
upstream task's XCom. See :ref:`typescript-sdk/types` for how the stored
JSON maps to JavaScript types.
-* ``setXCom({key, value, ...})`` — publishes an XCom value.
+* ``setXCom({key, value, ...})`` publishes an XCom value.
Logging
-------
@@ -262,7 +273,7 @@ Building and packaging
``airflow-ts-pack`` (shipped with the SDK) bundles the entry module and all of
its imports with esbuild into
a single self-contained ESM file, ``bundle.mjs``, and embeds the manifest (the
``dag_id`` and ``task_id``
map plus the supervisor schema version) after a leading compact JSON ``//#
airflowBundle=...`` layout header.
-The layout records the byte ranges and SHA-256 digests of the manifest and
executable code — one file to
+The layout records the byte ranges and SHA-256 digests of the manifest and
executable code, so one file to
deploy, with no separate manifest or ``node_modules``.
``esbuild`` is an optional peer dependency: packing is build-time only, so the
runtime install of
diff --git a/ts-sdk/README.md b/ts-sdk/README.md
index 33b25fea65d..2c5b9345748 100644
--- a/ts-sdk/README.md
+++ b/ts-sdk/README.md
@@ -35,11 +35,11 @@ npm install [email protected]
## Task Handlers
```ts
-import { Dag, DagRegistry, serveDags, type TaskHandlerArgs } from
"apache-airflow-ts-sdk";
+import { Dag, DagRegistry, getClient, getContext, serveDags } 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}` };
+export async function sayHello() {
+ const greeting = await getClient().getVariable("greeting");
+ return { message: `Hello from ${getContext().taskId}: ${greeting}` };
}
const dag = new Dag("example_dag");
@@ -48,6 +48,9 @@ dag.task("say_hello", sayHello);
await serveDags(new DagRegistry(dag));
```
+A handler is a plain function. `getContext()` and `getClient()` reach the
runtime from inside the call,
+so a handler takes no SDK-supplied parameter.
+
Non-`undefined` return values are pushed to XCom under the `"return_value"`
key by the active runtime, matching Python `@task` behavior.
@@ -100,9 +103,10 @@ Airflow metadata in the bundle itself.
TypeScript entrypoint:
```ts
-import { Dag, DagRegistry, serveDags, type TaskHandlerArgs } from
"apache-airflow-ts-sdk";
+import { Dag, DagRegistry, getClient, serveDags } from "apache-airflow-ts-sdk";
-export async function extract({ client }: TaskHandlerArgs) {
+export async function extract() {
+ const client = getClient();
const connection = await client.getConnection("sales_db");
const rowCount = Number((await client.getVariable("daily_row_count")) ??
"0");
@@ -112,8 +116,8 @@ export async function extract({ client }: TaskHandlerArgs) {
};
}
-export async function transform({ client }: TaskHandlerArgs) {
- const extracted = await client.getXCom<{ rowCount: number }>({
+export async function transform() {
+ const extracted = await getClient().getXCom<{ rowCount: number }>({
key: "return_value",
taskId: "extract",
});
@@ -144,8 +148,8 @@ 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.
-`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.
+`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:
@@ -195,13 +199,12 @@ one deployable file with no hand-written metadata sidecar.
Options:
-- `--outdir <dir>` — output directory (default `dist`)
-- `--source <name>` — display name of the primary source file shown in the
- Airflow UI (default: entry basename)
+- `--outdir <dir>`: output directory (default `dist`)
+- `--source <name>`: display name of the primary source file shown in the
Airflow UI (default: entry basename)
## TaskClient
-Every task handler receives a `TaskClient` for task-time Airflow data access:
+`getClient()` returns a `TaskClient` for task-time Airflow data access, for as
long as a handler is running:
| Method | Description |
| ------------------------------------------------ | ------------------- |
@@ -214,10 +217,9 @@ current task context when omitted.
## Cancellation
-`ctx.signal` is an `AbortSignal` controlled by the active runtime. Pass it to
-`fetch()`, timers, database clients, child processes, or any other API that
-accepts an abort signal so tasks can clean up cooperatively when Airflow
-terminates the task subprocess with SIGTERM or SIGINT.
+`getContext().signal` is an `AbortSignal` controlled by the active runtime.
+Pass it to `fetch()`, timers, database clients, child processes, or any other
API that accepts an abort signal,
+so tasks can clean up cooperatively when Airflow terminates the task
subprocess.
## Compatibility matrix
@@ -225,8 +227,7 @@ Which Airflow TaskInstance states and capabilities this SDK
supports. This table
[`capabilities.yaml`](https://github.com/apache/airflow/blob/main/ts-sdk/capabilities.yaml);
the conformance dimensions are defined in the
[Language SDK conformance
spec](https://github.com/apache/airflow/blob/main/contributing-docs/30_new_language_sdk.rst).
-Do not edit the table by hand — update the manifest and run the
-`update-ts-sdk-readme-matrix` prek hook.
+Do not edit the table by hand. Update the manifest and run the
`update-ts-sdk-readme-matrix` prek hook.
<!-- BEGIN AUTO-GENERATED LANG-SDK COMPAT MATRIX -->
@@ -275,13 +276,11 @@ Do not edit the table by hand — update the manifest and
run the
## Links
- [TypeScript SDK guide (staged
docs)](https://airflow.staged.apache.org/docs/apache-airflow/stable/authoring-and-scheduling/language-sdks/typescript.html)
- — how Airflow runs TypeScript task handlers
+ (how Airflow runs TypeScript task handlers)
- [API reference
(staged)](https://airflow.staged.apache.org/docs/ts-sdk/stable/)
- — generated from the TypeScript sources
-- [Source](https://github.com/apache/airflow/tree/main/ts-sdk) — the `ts-sdk/`
- directory of the Apache Airflow monorepo
-- [Issues](https://github.com/apache/airflow/issues) — bug reports and feature
- requests
+ (generated from the TypeScript sources)
+- [Source](https://github.com/apache/airflow/tree/main/ts-sdk): the `ts-sdk/`
directory of the Apache Airflow monorepo
+- [Issues](https://github.com/apache/airflow/issues): bug reports and feature
requests
- [Website](https://airflow.apache.org) ·
[Slack](https://s.apache.org/airflow-slack)
- [Developing this
package](https://github.com/apache/airflow/blob/main/ts-sdk/DEVELOPMENT.md)
- — local build, docs, and the release workflow
+ (local build, docs, and the release workflow)
diff --git a/ts-sdk/api-docs/dag-authoring-api.ts
b/ts-sdk/api-docs/dag-authoring-api.ts
index f91f824eabd..065cd0ae689 100644
--- a/ts-sdk/api-docs/dag-authoring-api.ts
+++ b/ts-sdk/api-docs/dag-authoring-api.ts
@@ -19,13 +19,12 @@
/** @module Authoring */
-export { Dag, DagRegistry, serveDags } from "../src/index.js";
+export { Dag, DagRegistry, getClient, getContext, serveDags } from
"../src/index.js";
export type {
DagSpec,
TaskClient,
TaskContext,
- TaskHandler,
- TaskHandlerArgs,
+ TaskFunction,
TaskInputs,
TaskOptions,
TaskRef,
diff --git a/ts-sdk/docs/index.md b/ts-sdk/docs/index.md
index 09c44272848..5082acd4a5f 100644
--- a/ts-sdk/docs/index.md
+++ b/ts-sdk/docs/index.md
@@ -35,17 +35,18 @@ Install the beta package from npm:
npm install [email protected]
```
-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:
+Define a Dag and register its task handlers.
+A handler is a plain function: `getContext()` returns the `TaskContext` and
`getClient()` the `TaskClient`
+for as long as it runs, so neither is a parameter.
+Any non-`undefined` return value is pushed to XCom under the `"return_value"`
key by the active runtime,
+matching Python `@task` behavior:
```ts
-import { Dag, DagRegistry, serveDags, type TaskHandlerArgs } from
"apache-airflow-ts-sdk";
+import { Dag, DagRegistry, getClient, getContext, serveDags } 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}` };
+export async function sayHello() {
+ const greeting = await getClient().getVariable("greeting");
+ return { message: `Hello from ${getContext().taskId}: ${greeting}` };
}
const dag = new Dag("example_dag");
diff --git a/ts-sdk/example/src/main.ts b/ts-sdk/example/src/main.ts
index 563f474b422..e98be36b92e 100644
--- a/ts-sdk/example/src/main.ts
+++ b/ts-sdk/example/src/main.ts
@@ -17,11 +17,12 @@
* under the License.
*/
-import { Dag, DagRegistry, serveDags, type TaskHandlerArgs } from
"apache-airflow-ts-sdk";
+import { Dag, DagRegistry, getClient, serveDags } from "apache-airflow-ts-sdk";
const dag = new Dag("typescript_example");
-export async function buildMessage({ client }: TaskHandlerArgs) {
+export async function buildMessage() {
+ const client = getClient();
const upstream = await client.getXCom<string>({
key: "return_value",
taskId: "python_start",
@@ -37,8 +38,8 @@ export async function buildMessage({ client }:
TaskHandlerArgs) {
};
}
-export async function readConnection({ client }: TaskHandlerArgs) {
- const connection = await client.getConnection("typescript_example_http");
+export async function readConnection() {
+ const connection = await
getClient().getConnection("typescript_example_http");
return {
id: connection?.id ?? null,
diff --git a/ts-sdk/scripts/verify-package.mjs
b/ts-sdk/scripts/verify-package.mjs
index 975f05f033f..5655310a646 100644
--- a/ts-sdk/scripts/verify-package.mjs
+++ b/ts-sdk/scripts/verify-package.mjs
@@ -32,7 +32,9 @@ const COMMAND_TIMEOUT_MS = 10 * 60 * 1000;
export const REQUIRED_ROOT_FILES = ["LICENSE", "NOTICE", "README.md",
"package.json"];
// The Dag-authoring entrypoints a consumer must be able to reach from the
package root.
-export const REQUIRED_ROOT_EXPORTS = ["Dag", "DagRegistry", "serveDags"];
+// The two getters are here because a handler cannot reach the runtime without
them: a
+// published build that dropped them would still import, and fail at the first
task.
+export const REQUIRED_ROOT_EXPORTS = ["Dag", "DagRegistry", "getClient",
"getContext", "serveDags"];
function run(command, args, options = {}) {
const result = spawnSync(command, args, {
@@ -59,7 +61,7 @@ export function isAllowedFile(path) {
}
/**
- * Every path a consumer can resolve — the `exports` conditions plus the `bin`
targets. Derived
+ * Every path a consumer can resolve: the `exports` conditions plus the `bin`
targets. Derived
* from package.json so adding an export subpath extends the check instead of
silently escaping it.
*/
export function collectEntryPoints(packageJson) {
@@ -72,7 +74,7 @@ export function collectEntryPoints(packageJson) {
/**
* The `exports` subpath keys as specifier suffixes: `"."` becomes `""` and
`"./coordinator"`
- * becomes `"/coordinator"`. Presence in the tarball is not resolution — a
subpath can ship and
+ * becomes `"/coordinator"`. Presence in the tarball is not resolution, since
a subpath can ship and
* still fail to import if its conditions are wrong or an internal import is
broken.
*/
export function collectExportSubpaths(packageJson) {
@@ -84,7 +86,7 @@ export function collectExportSubpaths(packageJson) {
/** The trailing JSON object of `pnpm pack --json` stdout, past the `prepack`
build banner. */
export function parsePackMetadata(stdout) {
// `--silent` does not suppress the `prepack` lifecycle banner, so stdout is
build log followed
- // by the JSON report — take the trailing object rather than parsing the
whole stream.
+ // by the JSON report, so take the trailing object rather than parsing the
whole stream.
const metadataMatch = stdout.match(/(?:^|\n)(\{[\s\S]*\})\s*$/);
if (!metadataMatch) {
throw new Error("pnpm pack did not return package metadata");
@@ -166,7 +168,7 @@ function verifyPackage() {
});
// Spawning the installed bin shim proves it exists, is executable, and
resolves its imports.
- // Asserting only "exited non-zero, and blamed itself" keeps this a
packaging check — the CLI's
+ // Asserting only "exited non-zero, and blamed itself" keeps this a
packaging check, since the CLI's
// own argument handling is covered by its unit tests.
const [binName] = Object.keys(packageJson.bin ?? {});
if (!binName) {
diff --git a/ts-sdk/src/coordinator/runtime.ts
b/ts-sdk/src/coordinator/runtime.ts
index 5f8c1afaf25..79a273e4278 100644
--- a/ts-sdk/src/coordinator/runtime.ts
+++ b/ts-sdk/src/coordinator/runtime.ts
@@ -26,7 +26,8 @@
// where `my-bundle.mjs` is a user-bundled Node script that imports
// the SDK, creates `Dag` objects, attaches a handler per task with
// `dag.task(...)`, collects them in a `DagRegistry`, then awaits
-// `serveDags(registry)`.
+// `serveDags(registry)`. Each handler runs inside a task scope, which is what
+// `getContext()` and `getClient()` read.
//
// Lifecycle:
// 1. Parse --comm / --logs from argv
@@ -54,7 +55,7 @@ import {
} from "./protocol.js";
import { DagRegistry, isDagRegistry, listRegistryTasks } from
"../sdk/registry.js";
import { DUPLICATE_COPY_HINT } from "../sdk/brand.js";
-import type { TaskContext, TaskHandlerArgs } from "../sdk/task.js";
+import { runInTaskScope, type TaskContext } from "../sdk/task.js";
import type { JsonValue } from "../sdk/client-types.js";
export const ABORT_GRACE_PERIOD_MS = 30_000;
@@ -82,8 +83,8 @@ function serveLatch(): Record<symbol, boolean | undefined> {
* ```
*
* The registry is the bundle's complete set of Dags: this process serves one
- * supervisor request, so a second call — which would connect a second pair of
- * sockets — is rejected. A call that fails is not a serve, and may be retried.
+ * supervisor request, so a second call, which would connect a second pair of
+ * sockets, is rejected. A call that fails is not a serve, and may be retried.
* Resolves when Airflow's supervisor has been sent the terminal frame for the
* work this process was started for; the same call also answers the build-time
* `--airflow-metadata` query `airflow-ts-pack` makes.
@@ -361,18 +362,20 @@ async function handleTask(
const ctx = buildContext(details, signal);
const client = createCoordinatorClient(comm, ctx, clientLogs);
- const args: TaskHandlerArgs = { ctx, client };
// Startup-details fields already logged above (`Received task
// startup details`); this line just marks the handler-call boundary.
logs.debug("Dispatching to handler", { task_id: ctx.taskId });
try {
- const result = await handler(args);
+ // The scope is installed around the call, not awaited inside it: the store
+ // follows the handler across every `await` it makes, so `getContext()` and
+ // `getClient()` work at any depth without the handler being handed either.
+ const result = await runInTaskScope({ ctx, client }, () => handler());
if (result !== undefined) {
await client.setXCom({ key: "return_value", value: result as JsonValue
});
}
// SucceedTask MUST include task_outlets and outlet_events as
- // empty lists — the Execution API's TISuccessStatePayload
+ // empty lists, since the Execution API's TISuccessStatePayload
// tagged-union validator rejects null for these fields.
const response: RuntimeSucceedTask = {
type: "SucceedTask",
diff --git a/ts-sdk/src/index.ts b/ts-sdk/src/index.ts
index a7c8f4960b0..b4a412f172e 100644
--- a/ts-sdk/src/index.ts
+++ b/ts-sdk/src/index.ts
@@ -19,9 +19,10 @@
export { Dag } from "./sdk/dag.js";
export { DagRegistry } from "./sdk/registry.js";
+export { getClient, getContext } from "./sdk/task.js";
export { ConnectionNotFoundError, VariableNotFoundError } from
"./sdk/client.js";
export { serveDags, SUPERVISOR_API_VERSION } from "./coordinator/index.js";
export type { DagSpec, TaskInputs, TaskOptions, TaskRef, TaskSpec } from
"./sdk/dag.js";
export type { TaskClient } from "./sdk/client.js";
export type { ConnectionResult, GetXComOpts, JsonValue, SetXComOpts } from
"./sdk/client-types.js";
-export type { TaskContext, TaskHandler, TaskHandlerArgs } from "./sdk/task.js";
+export type { TaskContext, TaskFunction } from "./sdk/task.js";
diff --git a/ts-sdk/src/sdk/dag.ts b/ts-sdk/src/sdk/dag.ts
index b00492ffd87..a329aa8c869 100644
--- a/ts-sdk/src/sdk/dag.ts
+++ b/ts-sdk/src/sdk/dag.ts
@@ -20,7 +20,7 @@
// The Dag authoring surface: `new Dag(dagId)` plus `dag.task(taskId,
handler)`.
import { brand, hasBrand } from "./brand.js";
-import type { TaskHandler } from "./task.js";
+import type { TaskFunction } from "./task.js";
function isPlainRecord(value: unknown): value is Record<string, unknown> {
if (typeof value !== "object" || value === null || Array.isArray(value))
return false;
@@ -38,7 +38,7 @@ function validateEmptySpec(name: string, value: unknown):
void {
* Dag-level options.
*
* No fields yet, so only `{}` is accepted: a field that would be silently
- * dropped — `new Dag("d", { schedule: "@daily" })` — is a compile error.
+ * dropped, such as `new Dag("d", { schedule: "@daily" })`, is a compile error.
*
* Native Dag declaration will add optional fields here, generated from the
* serialized-Dag JSON schema as `src/generated/supervisor.ts` is.
@@ -73,7 +73,7 @@ export interface TaskRef {
* Task references keyed by input name.
*
* Stored and validated, but they do not create dependencies or pass values yet
- * — see {@link TaskOptions.inputs}. Each must identify an earlier task in the
+ * (see {@link TaskOptions.inputs}). Each must identify an earlier task in the
* same Dag. Literal values are not supported.
*/
export type TaskInputs = Readonly<Record<string, TaskRef>>;
@@ -89,23 +89,23 @@ export interface TaskOptions {
/**
* References to the upstream tasks this task consumes.
*
- * Not used yet: a handler receives `{ctx, client}` only, and the Python stub
- * Dag defines task order. Read an upstream return value explicitly instead —
- * `client.getXCom({ key: "return_value", taskId: "extract" })`, where
+ * Not used yet: a handler takes no arguments, and the Python stub Dag
+ * defines task order. Read an upstream return value explicitly instead, with
+ * `getClient().getXCom({ key: "return_value", taskId: "extract" })`, where
* omitting `taskId` reads the *running* task's own XCom, not the upstream.
*
* In the future these will declare dependencies in native TypeScript Dags.
*/
readonly inputs?: TaskInputs;
- /** Task-level options. Stored, but not used yet — see {@link TaskSpec}. */
+ /** Task-level options. Stored, but not used yet (see {@link TaskSpec}). */
readonly spec?: TaskSpec;
}
-/** Per-task record a Dag retains: the handle, the handler, its spec, and the
+/** Per-task record a Dag retains: the handle, the function, its spec, and the
* upstream handles feeding it. */
export interface TaskRecord {
readonly task: TaskRef;
- readonly handler: TaskHandler;
+ readonly fn: TaskFunction;
readonly spec: TaskSpec;
/** Upstream handles keyed by input name; empty when the task has no inputs.
*/
readonly inputs: TaskInputs;
@@ -150,7 +150,7 @@ export class Dag {
// Copied and frozen, as task specs and inputs are: nothing reads a spec
// until the bundle manifest is built, long after the user's module has
run,
// so a later mutation of their object would silently change what is
packed.
- // Shallow — a nested value in a future generated spec stays mutable.
+ // Shallow, so a nested value in a future generated spec stays mutable.
this.spec = Object.freeze({ ...spec });
}
@@ -167,7 +167,7 @@ export class Dag {
*/
task<TReturn = unknown>(
taskId: string,
- handler: TaskHandler<TReturn>,
+ handler: TaskFunction<TReturn>,
options: TaskOptions = {},
): TaskRef {
if (typeof handler !== "function") {
@@ -183,15 +183,15 @@ export class Dag {
const task: TaskRef = Object.freeze({ dagId: this.dagId, taskId });
this.#tasks.set(taskId, {
task,
- handler: handler as TaskHandler,
+ fn: handler as TaskFunction,
spec: Object.freeze({ ...spec }),
inputs: Object.freeze({ ...inputs }),
});
return task;
}
- // TypeScript is bypassable — from plain JavaScript, or an `as TaskOptions`
- // cast — so an unknown key is rejected rather than silently ignored.
+ // TypeScript is bypassable from plain JavaScript or an `as TaskOptions`
+ // cast, so an unknown key is rejected rather than silently ignored.
#validateOptions(taskId: string, options: TaskOptions): void {
const value: unknown = options;
if (!isPlainRecord(value)) {
diff --git a/ts-sdk/src/sdk/registry.ts b/ts-sdk/src/sdk/registry.ts
index ee4286de926..9d9259ad540 100644
--- a/ts-sdk/src/sdk/registry.ts
+++ b/ts-sdk/src/sdk/registry.ts
@@ -19,7 +19,7 @@
import { brand, DUPLICATE_COPY_HINT, hasBrand } from "./brand.js";
import { Dag, getDagTaskRecords, isDag, type TaskRef } from "./dag.js";
-import type { TaskHandler } from "./task.js";
+import type { TaskFunction } from "./task.js";
// Assigned inside DagRegistry's static block, as Dag does for its tasks.
let dagsOf: (registry: DagRegistry) => ReadonlyMap<string, Dag>;
@@ -53,7 +53,7 @@ export interface RegisteredDag {
* handler through {@link getTaskHandler} without any runtime in scope.
*
* Lookups delegate live to each Dag's task map, so tasks added to a Dag
- * after registration are visible — the registry records Dag identity, not
+ * after registration are visible: the registry records Dag identity, not
* a snapshot of its tasks.
*/
export class DagRegistry {
@@ -80,8 +80,8 @@ export class DagRegistry {
// Typed as Dag, so narrowing it would collapse to never; these guard
// callers reaching this from plain JavaScript.
const candidate: unknown = dag;
- // Another copy's Dag cannot be registered — lookups read a private task
- // map keyed to this copy's class — so it is rejected, but by its cause.
+ // Another copy's Dag cannot be registered, since lookups read a private
+ // task map keyed to this copy's class, so it is rejected by its cause.
if (!(candidate instanceof Dag)) {
throw new Error(
isDag(candidate)
@@ -101,14 +101,14 @@ export class DagRegistry {
/** Look up a registered handler, the way the runtime dispatches a task.
* Returns `undefined` when no handler exists. */
- getTaskHandler(dagId: string, taskId: string): TaskHandler | undefined {
+ getTaskHandler(dagId: string, taskId: string): TaskFunction | undefined {
const dag = this.#dags.get(dagId);
- return dag ? getDagTaskRecords(dag).get(taskId)?.handler : undefined;
+ return dag ? getDagTaskRecords(dag).get(taskId)?.fn : undefined;
}
}
/** Internal: the task handles across a registry's Dags. Not re-exported from
the
- * package root — enumerating what the runtime dispatches is the runtime's
job. */
+ * package root: enumerating what the runtime dispatches is the runtime's
job. */
export function listRegistryTasks(registry: DagRegistry): TaskRef[] {
return [...dagsOf(registry).values()].flatMap((dag) =>
[...getDagTaskRecords(dag).values()].map((record) => record.task),
diff --git a/ts-sdk/src/sdk/task.ts b/ts-sdk/src/sdk/task.ts
index 40d5c7ee0a7..3686691120b 100644
--- a/ts-sdk/src/sdk/task.ts
+++ b/ts-sdk/src/sdk/task.ts
@@ -17,7 +17,12 @@
* under the License.
*/
-// The task-handler call surface — types every user task handler sees.
+// The task-handler call surface: what a task handler sees while it runs.
+//
+// Everything the SDK supplies comes from a getter rather than a parameter, so
+// nothing it injects shares a namespace with an author's arguments.
+
+import { AsyncLocalStorage } from "node:async_hooks";
import type { TaskClient } from "./client.js";
@@ -43,19 +48,93 @@ export interface TaskContext {
readonly signal: AbortSignal;
}
-/** Arguments passed to every task handler. */
-export interface TaskHandlerArgs {
- /** Runtime metadata for the current task invocation. */
+/** What the runtime puts in scope for the duration of one handler call. */
+export interface TaskScope {
readonly ctx: TaskContext;
- /** Client for reading and writing Airflow task-time data. */
readonly client: TaskClient;
}
+// Keyed on a global symbol, as the serve latch is: two resolved copies of the
+// package would otherwise hold one storage each, and a handler reaching for
+// `getClient()` through the copy that is not running the task would find
+// nothing in scope.
+const SCOPE_STORAGE = Symbol.for("airflow.ts-sdk.task-scope");
+
+function scopeStorage(): AsyncLocalStorage<TaskScope> {
+ const holder = globalThis as unknown as Record<symbol,
AsyncLocalStorage<TaskScope> | undefined>;
+ return (holder[SCOPE_STORAGE] ??= new AsyncLocalStorage<TaskScope>());
+}
+
+/**
+ * Internal: call `fn` with `scope` in place, as the runtime does per task.
+ *
+ * `AsyncLocalStorage` carries the store across every `await` and into every
+ * promise created inside `fn`, so a handler's helpers see it without being
+ * passed anything. Not re-exported from the package root: a handler reads the
+ * scope, it does not install one.
+ */
+export function runInTaskScope<T>(scope: TaskScope, fn: () => T): T {
+ return scopeStorage().run(scope, fn);
+}
+
+type ScopeAccessor = "getContext" | "getClient";
+
+function currentScope(accessor: ScopeAccessor): TaskScope {
+ const scope = scopeStorage().getStore();
+ if (!scope) {
+ throw new Error(
+ `${accessor}() is only available inside a task handler. ` +
+ "The scope is in place only for the duration of the handler call, so "
+
+ "this ran either at module top level or in work that outlived it.",
+ );
+ }
+ return scope;
+}
+
+/**
+ * Runtime metadata for the task currently running.
+ *
+ * ```ts
+ * async function transform() {
+ * throw new Error(`task ${getContext().taskId} has nothing to transform`);
+ * }
+ * ```
+ *
+ * @throws when called outside a task handler.
+ */
+export function getContext(): TaskContext {
+ return currentScope("getContext").ctx;
+}
+
+/**
+ * Client for reading and writing Airflow task-time data for the task currently
+ * running.
+ *
+ * ```ts
+ * async function transform() {
+ * const client = getClient();
+ * return await client.getXCom<number>({ key: "return_value", taskId:
"extract" });
+ * }
+ * ```
+ *
+ * Work that outlives the handler is the one gap. Async context propagates into
+ * a promise created inside the handler, so one it never awaits still resolves,
+ * but it runs after the task's terminal state has been reported and writes to
+ * a finished task. Await everything a handler starts.
+ *
+ * @throws when called outside a task handler.
+ */
+export function getClient(): TaskClient {
+ return currentScope("getClient").client;
+}
+
/**
* Function signature for a TypeScript task handler.
*
- * Non-`undefined` return values are automatically pushed to XCom under
- * the `"return_value"` key, matching Python `@task` behavior. Return
- * `undefined` or omit a return value to skip the automatic XCom push.
+ * A handler takes no SDK-supplied parameter: {@link getContext} and
+ * {@link getClient} supply the runtime's side of the call.
+ *
+ * Non-`undefined` return values are automatically pushed to XCom under the
+ * `"return_value"` key, matching Python `@task` behavior.
*/
-export type TaskHandler<TReturn = unknown> = (args: TaskHandlerArgs) =>
TReturn | Promise<TReturn>;
+export type TaskFunction<TReturn = unknown> = () => TReturn | Promise<TReturn>;
diff --git a/ts-sdk/tests/coordinator/integration.test.ts
b/ts-sdk/tests/coordinator/integration.test.ts
index 040de00fa96..a3293a57847 100644
--- a/ts-sdk/tests/coordinator/integration.test.ts
+++ b/ts-sdk/tests/coordinator/integration.test.ts
@@ -25,7 +25,7 @@
// drives the runtime through task success, task failure, retry, abort
signaling,
// task-time RPCs, missing handlers, and parse-mode responses.
//
-// No Python, no Airflow install — but exercises the same wire format
+// No Python, no Airflow install, but exercises the same wire format
// the real coordinator speaks.
import { afterEach, describe, expect, it, vi } from "vitest";
@@ -37,6 +37,7 @@ import {
} from "../../src/coordinator/runtime.js";
import { Dag } from "../../src/sdk/dag.js";
import { DagRegistry } from "../../src/sdk/registry.js";
+import { getClient, getContext } from "../../src/sdk/task.js";
const testDag = new Dag("test_dag");
const otherDag = new Dag("other_dag");
@@ -54,7 +55,7 @@ interface MockResult {
/** Callback used by `driveSupervisor` to answer runtime-initiated
* requests. Return `{ body, error? }` to reply with that arity-3
* frame, or `null` to ignore the request (the runtime will hang
- * waiting for a response — only useful for negative tests). */
+ * waiting for a response, which is only useful for negative tests). */
type Responder = (
msgType: string,
body: Record<string, unknown>,
@@ -157,7 +158,7 @@ async function driveSupervisor(initialFrame: unknown,
responder?: Responder): Pr
const [commSock, logsSock] = await Promise.all([commAccept, logsAccept]);
- // Send the kickoff frame as a _ResponseFrame (arity 3) — matches what
+ // Send the kickoff frame as a _ResponseFrame (arity 3), matching what
// Airflow's `_send_startup_details` actually emits on the wire.
commSock.write(frameBytes(0, initialFrame, true));
@@ -239,8 +240,8 @@ describe("coordinator runtime integration", () => {
it("dispatches StartupDetails to a registered handler and emits
SucceedTask", async () => {
let observedCtx: unknown = null;
- testDag.task("say_hello", async ({ ctx }) => {
- observedCtx = ctx;
+ testDag.task("say_hello", async () => {
+ observedCtx = getContext();
return "ok";
});
@@ -265,7 +266,7 @@ describe("coordinator runtime integration", () => {
// Logger names should be hierarchical (`ts-sdk.<subsystem>`) so the
// Python supervisor's ConsoleRenderer prints them as a distinct
- // `[name]` column — not hardcoded to "task" (which collides with
+ // `[name]` column, not hardcoded to "task" (which collides with
// user task logs).
const loggers = new Set(result.logRecords.map((r) => r["logger"]));
expect(loggers.has("ts-sdk.runtime")).toBe(true);
@@ -355,11 +356,11 @@ describe("coordinator runtime integration", () => {
});
});
- it("aborts ctx.signal on SIGTERM and reports a thrown task error", async ()
=> {
+ it("aborts the context signal on SIGTERM and reports a thrown task error",
async () => {
let sawAbort = false;
- testDag.task("aborted_then_failed", async ({ ctx }) => {
+ testDag.task("aborted_then_failed", async () => {
process.emit("SIGTERM");
- sawAbort = ctx.signal.aborted;
+ sawAbort = getContext().signal.aborted;
throw new Error("interrupted");
});
@@ -375,9 +376,9 @@ describe("coordinator runtime integration", () => {
it("returns RetryTask with the thrown error when a task fails after
SIGTERM", async () => {
let sawAbort = false;
- testDag.task("aborted_then_failed_retry", async ({ ctx }) => {
+ testDag.task("aborted_then_failed_retry", async () => {
process.emit("SIGTERM");
- sawAbort = ctx.signal.aborted;
+ sawAbort = getContext().signal.aborted;
throw new Error("interrupted");
});
@@ -398,9 +399,9 @@ describe("coordinator runtime integration", () => {
it("does not discard a completed task result after SIGTERM", async () => {
let sawAbort = false;
- testDag.task("completed_after_sigterm", async ({ ctx }) => {
+ testDag.task("completed_after_sigterm", async () => {
process.emit("SIGTERM");
- sawAbort = ctx.signal.aborted;
+ sawAbort = getContext().signal.aborted;
return "completed";
});
@@ -434,9 +435,10 @@ describe("coordinator runtime integration", () => {
const xcomStore = new Map<string, unknown>();
let observedGreeting: string | null = "<unset>";
- testDag.task("say_hello_client", async ({ ctx, client }) => {
- // The coordinator-mode handler MUST receive a client.
- if (!client) throw new Error("client missing in coordinator mode");
+ testDag.task("say_hello_client", async () => {
+ // The coordinator-mode handler MUST be able to reach a client.
+ const ctx = getContext();
+ const client = getClient();
observedGreeting = await client.getVariable("e6_greeting");
await client.setXCom({
@@ -503,8 +505,8 @@ describe("coordinator runtime integration", () => {
it("returns null from getVariable when the supervisor signals NOT_FOUND",
async () => {
let observed: string | null = "<unset>";
- testDag.task("missing_variable", async ({ client }) => {
- observed = await client.getVariable("missing_key");
+ testDag.task("missing_variable", async () => {
+ observed = await getClient().getVariable("missing_key");
});
const responder: Responder = (msgType) => {
diff --git a/ts-sdk/tests/public-api.test.ts b/ts-sdk/tests/public-api.test.ts
index c1722b85c89..9f0f732f838 100644
--- a/ts-sdk/tests/public-api.test.ts
+++ b/ts-sdk/tests/public-api.test.ts
@@ -26,7 +26,7 @@ import type {
SetXComOpts,
TaskClient,
TaskContext,
- TaskHandler,
+ TaskFunction,
TaskInputs,
TaskOptions,
TaskRef,
@@ -37,6 +37,8 @@ import {
ConnectionNotFoundError,
Dag,
DagRegistry,
+ getClient,
+ getContext,
serveDags,
SUPERVISOR_API_VERSION,
VariableNotFoundError,
@@ -153,6 +155,27 @@ describe("public API", () => {
expectTypeOf<typeof sdk>().not.toHaveProperty("startCoordinator");
});
+ describe("the task-handler getters", () => {
+ it("throw outside a handler, naming the accessor", () => {
+ // The full scope behaviour is covered in tests/sdk/task-scope.test.ts;
+ // this pins that both reach the package root and say what went wrong.
+ expect(() => getContext()).toThrow(/^getContext\(\) is only available
inside a task handler/);
+ expect(() => getClient()).toThrow(/^getClient\(\) is only available
inside a task handler/);
+ });
+
+ it("are the only way a handler reaches the runtime", () => {
+ // A handler is a plain function of its own data, so the SDK hands it no
+ // parameter at all and the scope is not something an author installs.
+ expectTypeOf<TaskFunction>().toEqualTypeOf<() => unknown |
Promise<unknown>>();
+ expectTypeOf<typeof getContext>().toEqualTypeOf<() => TaskContext>();
+ expectTypeOf<typeof getClient>().toEqualTypeOf<() => TaskClient>();
+ for (const name of ["TaskHandlerArgs", "runInTaskScope", "TaskScope"]) {
+ expect(name in sdk).toBe(false);
+ }
+ expectTypeOf<typeof sdk>().not.toHaveProperty("runInTaskScope");
+ });
+ });
+
it("exports public error classes", () => {
const err = new VariableNotFoundError("missing");
expect(err).toBeInstanceOf(Error);
@@ -185,7 +208,7 @@ describe("public API", () => {
expectTypeOf<Dag["task"]>().toEqualTypeOf<
<TReturn = unknown>(
taskId: string,
- handler: TaskHandler<TReturn>,
+ handler: TaskFunction<TReturn>,
options?: TaskOptions,
) => TaskRef
>();
diff --git a/ts-sdk/tests/sdk/dag.test.ts b/ts-sdk/tests/sdk/dag.test.ts
index e76c24ab4c1..6177cf28b02 100644
--- a/ts-sdk/tests/sdk/dag.test.ts
+++ b/ts-sdk/tests/sdk/dag.test.ts
@@ -115,7 +115,7 @@ describe("Dag", () => {
expect(dag.spec).toEqual(dagSpec);
expect(Object.isFrozen(dag.spec)).toBe(true);
const record = getDagTaskRecords(dag).get("my_task");
- expect(record?.handler).toBe(handler);
+ expect(record?.fn).toBe(handler);
expect(record?.spec).toEqual(taskSpec);
expect(Object.isFrozen(record!.spec)).toBe(true);
});
diff --git a/ts-sdk/tests/sdk/task-scope.test.ts
b/ts-sdk/tests/sdk/task-scope.test.ts
new file mode 100644
index 00000000000..dc6299a4f6f
--- /dev/null
+++ b/ts-sdk/tests/sdk/task-scope.test.ts
@@ -0,0 +1,170 @@
+/*!
+ * 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.
+ */
+
+import { describe, expect, it, vi } from "vitest";
+
+import type { TaskClient } from "../../src/sdk/client.js";
+import {
+ getClient,
+ getContext,
+ runInTaskScope,
+ type TaskContext,
+ type TaskScope,
+} from "../../src/sdk/task.js";
+
+function makeScope(taskId: string): TaskScope {
+ const ctx: TaskContext = {
+ dagId: "scope_dag",
+ taskId,
+ runId: "r1",
+ tryNumber: 1,
+ mapIndex: -1,
+ signal: new AbortController().signal,
+ };
+ // Identity is all these tests read, so a cast beats a full fake client.
+ const client = { label: taskId } as unknown as TaskClient;
+ return { ctx, client };
+}
+
+describe("the task scope", () => {
+ it("hands a handler the context and client the runtime installed", async ()
=> {
+ const scope = makeScope("dispatched");
+
+ const seen = await runInTaskScope(scope, async () => ({
+ ctx: getContext(),
+ client: getClient(),
+ }));
+
+ expect(seen.ctx).toBe(scope.ctx);
+ expect(seen.client).toBe(scope.client);
+ });
+
+ it("survives await boundaries and reaches helpers it was never passed",
async () => {
+ // The point of AsyncLocalStorage over a parameter: a helper several frames
+ // and several awaits deep reads the scope without anything being threaded.
+ async function readTaskIdDeep(): Promise<string> {
+ await Promise.resolve();
+ await new Promise((resolve) => setTimeout(resolve, 0));
+ return getContext().taskId;
+ }
+
+ const taskId = await runInTaskScope(makeScope("deep"), async () => {
+ await Promise.resolve();
+ await new Promise((resolve) => setImmediate(resolve));
+ return await readTaskIdDeep();
+ });
+
+ expect(taskId).toBe("deep");
+ });
+
+ it("keeps concurrent handlers from seeing each other's scope", async () => {
+ // A module variable would pass every single-handler test above, so this
+ // is what pins the store as per-call.
+ const run = (taskId: string, delayMs: number) =>
+ runInTaskScope(makeScope(taskId), async () => {
+ await new Promise((resolve) => setTimeout(resolve, delayMs));
+ return getContext().taskId;
+ });
+
+ // The slower one is started first, so a shared store would answer
"second".
+ expect(await Promise.all([run("first", 10), run("second",
0)])).toEqual(["first", "second"]);
+ });
+
+ it("throws outside a handler, naming the accessor", () => {
+ expect(() => getContext()).toThrow(/^getContext\(\) is only available
inside a task handler/);
+ expect(() => getClient()).toThrow(/^getClient\(\) is only available inside
a task handler/);
+ });
+
+ it("throws in work started outside the handler it belongs to", async () => {
+ // A callback handed to a module-level queue or an emitter registered at
+ // import time does not carry the scope with it.
+ let escaped: (() => TaskContext) | null = null;
+ await runInTaskScope(makeScope("escaping"), async () => {
+ escaped = () => getContext();
+ });
+
+ expect(escaped).not.toBeNull();
+ expect(() => escaped!()).toThrow(/only available inside a task handler/);
+ });
+
+ it("still resolves in a floating promise, which is why one must be awaited",
async () => {
+ // By the time a promise the handler never awaited resolves, the runtime
+ // has already reported the task's terminal state. Pinned because the fix
+ // is to await the work, not to expect a throw here.
+ let floating: Promise<string> | null = null;
+ await runInTaskScope(makeScope("floating"), async () => {
+ floating = (async () => {
+ await new Promise((resolve) => setTimeout(resolve, 0));
+ return getContext().taskId;
+ })();
+ });
+
+ await expect(floating!).resolves.toBe("floating");
+ });
+
+ it("restores an outer scope after a nested one returns", async () => {
+ const outer = makeScope("outer");
+
+ const seen = await runInTaskScope(outer, async () => {
+ const inner = await runInTaskScope(makeScope("inner"), async () =>
getContext().taskId);
+ return { inner, afterInner: getContext().taskId };
+ });
+
+ expect(seen).toEqual({ inner: "inner", afterInner: "outer" });
+ });
+
+ it("is keyed globally, so a second resolved copy shares one storage", async
() => {
+ // Two copies of the package each import their own module instance. The
+ // storage lives on globalThis under a well-known symbol so the copy that
+ // dispatches and the copy a handler imported `getClient` from agree.
+ const key = Symbol.for("airflow.ts-sdk.task-scope");
+ const holder = globalThis as unknown as Record<symbol, unknown>;
+
+ await runInTaskScope(makeScope("global"), async () => {
+ expect(holder[key]).toBeDefined();
+ expect(getContext().taskId).toBe("global");
+ });
+ });
+
+ it("does not swallow a handler's own error", async () => {
+ const boom = new Error("handler blew up");
+ await expect(
+ runInTaskScope(makeScope("failing"), async () => {
+ throw boom;
+ }),
+ ).rejects.toBe(boom);
+ // The scope unwinds with the throw rather than leaking to the next call.
+ expect(() => getContext()).toThrow();
+ });
+
+ it("does not hold a reference the handler can reassign", async () => {
+ const scope = makeScope("frozen");
+ const spy = vi.fn();
+
+ await runInTaskScope(scope, async () => {
+ const first = getClient();
+ const second = getClient();
+ // Same object every read: nothing rebuilds a client per access.
+ expect(first).toBe(second);
+ spy();
+ });
+
+ expect(spy).toHaveBeenCalledOnce();
+ });
+});
