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 8533cc40b0a TS SDK: add the task state store to the task client
(#73618)
8533cc40b0a is described below
commit 8533cc40b0af312facfc8f467b786a48e6c8373a
Author: Wilson Josix Wang <[email protected]>
AuthorDate: Fri Oct 9 18:13:24 2026 +0900
TS SDK: add the task state store to the task client (#73618)
* TS SDK: add the task state store to the task client
TypeScript task handlers had no way to reach the per-task-instance state
store that Python exposes as context["task_state_store"], so a retried
TypeScript task could not reattach to an external job it had already
submitted. The supervisor and the generated wire types already carried
the four state-store messages; only the client surface was missing.
The store is reached through client.taskStateStore, with get, set,
delete and clear, so it reads like the Python accessor and the Java and
Go stores rather than four more flat client methods. It is confined to
the caller's own task instance by the Execution API, so the task
instance id is bound when the coordinator client is built rather than
exposed on TaskContext. Absence handling becomes operation-specific
instead of one shared not-found set, because a fourth lookup makes the
shared set answer null for codes an operation can never receive.
Expiry is resolved on the runtime, as Python does: omitting retentionMs
follows AIRFLOW__STATE_STORE__DEFAULT_RETENTION_DAYS, which the
coordinator passes from [state_store] default_retention_days, and the
write fails when that is absent or malformed rather than retaining the
key for a period nobody configured. NEVER_EXPIRE is a symbol, so a
retention that reaches Infinity by accident is rejected instead of
silently storing a key forever.
* TS SDK: cover the task state store in the example bundle and e2e tests
The example bundle is what the ts_sdk e2e mode runs, so the new client
methods need handlers there before the capability can be claimed. The
e2e assertions read the store back through the public REST API rather
than the task's own XCom, so they prove the values reached the database
and that NEVER_EXPIRE arrives as a null expiry. Clearing lives in its
own task because clear removes every key of the task instance that
calls it, which would destroy the rows the round-trip test asserts on.
* TS SDK: document the task state store and mark the capability supported
With the client methods in place and exercised end to end, the
compatibility matrix can stop listing task-state-store as unsupported.
The guide spells out the two places TypeScript deliberately differs from
Python so deployment managers are not surprised: the worker-side
state_store_backend is not used, and a config value of 0 means never
expire while retentionMs: 0 expires the key immediately.
* Fix spelling in TS SDK task state store docs
---------
Co-authored-by: ZHE YOU LIU <[email protected]>
---
.../language-sdks/typescript.rst | 53 ++++-
.../airflow_e2e_tests/e2e_test_utils/clients.py | 14 +-
.../ts_sdk_tests/test_ts_sdk_dag.py | 65 +++++-
ts-sdk/README.md | 15 +-
ts-sdk/api-docs/dag-authoring-api.ts | 11 +-
ts-sdk/api-docs/supporting-types.ts | 8 +-
ts-sdk/capabilities.yaml | 5 +-
ts-sdk/example/dags/typescript_example.py | 10 +
ts-sdk/example/src/main.ts | 36 +++-
ts-sdk/src/coordinator/client.ts | 198 ++++++++++++++---
ts-sdk/src/coordinator/protocol.ts | 5 +
ts-sdk/src/coordinator/runtime.ts | 4 +-
ts-sdk/src/index.ts | 11 +-
ts-sdk/src/sdk/client-types.ts | 23 ++
ts-sdk/src/sdk/client.ts | 74 ++++++-
ts-sdk/tests/coordinator/client.test.ts | 238 ++++++++++++++++++++-
ts-sdk/tests/coordinator/integration.test.ts | 23 ++
ts-sdk/tests/public-api.test.ts | 22 ++
18 files changed, 759 insertions(+), 56 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 d2c6094a105..0cb37f578e2 100644
--- a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst
+++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst
@@ -496,7 +496,7 @@ Two functions give it what it needs while it runs:
that fires when Airflow stops the task. Pass ``signal`` to ``fetch()``,
timers, or any other API that
accepts one, so the task can stop cleanly.
* - ``getClient()``
- - A ``TaskClient`` for Airflow Variables, Connections and XComs.
+ - A ``TaskClient`` for Airflow Variables, Connections, XComs and the task
state store.
Both work anywhere inside the handler's call, including helper functions and
code after an ``await``, and
throw outside it. Await everything the handler starts: work still running
after the handler returns runs
@@ -523,6 +523,8 @@ The ``TaskClient``
``taskId`` and ``mapIndex`` default to the current task; pass ``taskId`` to
read an upstream task's XCom.
See :ref:`typescript-sdk/types` for how values map to JavaScript types.
* ``setXCom({ key, value, ... })`` publishes an XCom value.
+* ``taskStateStore`` is this task instance's key/value store, with
``get<T>(key)``,
+ ``set(key, value, { retentionMs? })``, ``delete(key)`` and ``clear()``.
.. note::
@@ -530,6 +532,55 @@ The ``TaskClient``
stored value when the Variable is read back. ``setVariable`` without a
description clears the existing
description, and ``deleteVariable`` succeeds even when the Variable does
not exist.
+Task state store
+----------------
+
+The task state store is a per-task-instance key/value store, scoped to
``dagId``, ``runId``, ``taskId``,
+and ``mapIndex`` (but not ``tryNumber``), so a value written by one attempt is
still readable by the next.
+It survives worker crashes and task retries within the same Dag run, which
makes it a good place to record
+external job IDs, checkpoints within a task, and progress metadata. See
:doc:`/core-concepts/task-state-store`
+for the full concept, including the equivalent Python API.
+
+.. code-block:: typescript
+
+ import { getClient, NEVER_EXPIRE } from "apache-airflow-ts-sdk";
+
+ export async function runSparkJob() {
+ const store = getClient().taskStateStore;
+
+ let jobId = await store.get<string>("job_id");
+ if (jobId == null) {
+ jobId = await submitSparkJob();
+ await store.set("job_id", jobId, { retentionMs: NEVER_EXPIRE });
+ }
+
+ const result = await waitForSparkJob(jobId);
+ await store.delete("job_id");
+ return result;
+ }
+
+**Retention.** ``retentionMs`` is milliseconds to retain the key, counted from
the time of the write:
+
+* A number retains the key for that many milliseconds from now; ``retentionMs:
0`` expires the key
+ immediately.
+* ``NEVER_EXPIRE`` (imported from ``apache-airflow-ts-sdk``) stores the key
with no expiry, regardless of
+ the deployment default.
+* When ``retentionMs`` is omitted the runtime reads
``AIRFLOW__STATE_STORE__DEFAULT_RETENTION_DAYS``, which
+ the coordinator sets from ``[state_store] default_retention_days``. A value
of ``0`` there means the key
+ never expires. This is not the same as ``retentionMs: 0`` above, which
expires the key immediately.
+
+**Keys** must be non-empty strings of at most 512 characters; they may contain
slashes. **Values** must be
+JSON-compatible (see :ref:`typescript-sdk/types`), and can never be ``null`` —
``set`` rejects a ``null``
+value. ``get`` returns ``null`` only to mean "key not found."
+
+Two behaviors to note:
+
+* ``[workers] state_store_backend`` is not used from TypeScript tasks, so
values are always stored inline
+ through the Execution API. A key written by a Python task through a custom
worker-side backend reads back
+ from TypeScript as that backend's reference marker string, not the original
value.
+* ``[state_store] clear_on_success`` removes all of a task instance's state
store keys when the task
+ succeeds, the same as for Python tasks.
+
Logging
-------
diff --git
a/airflow-e2e-tests/tests/airflow_e2e_tests/e2e_test_utils/clients.py
b/airflow-e2e-tests/tests/airflow_e2e_tests/e2e_test_utils/clients.py
index 87b9b816637..273ab225320 100644
--- a/airflow-e2e-tests/tests/airflow_e2e_tests/e2e_test_utils/clients.py
+++ b/airflow-e2e-tests/tests/airflow_e2e_tests/e2e_test_utils/clients.py
@@ -163,12 +163,14 @@ class AirflowClient:
"""Get an Airflow Variable via API."""
return self._make_request(method="GET", endpoint=f"variables/{key}")
- def get_task_state_store(self, dag_id: str, run_id: str, task_id: str,
key: str):
- """Get a single task state store entry via API."""
- return self._make_request(
- method="GET",
-
endpoint=f"dags/{dag_id}/dagRuns/{run_id}/taskInstances/{task_id}/state-store/{key}",
- )
+ def get_task_state_store(
+ self, dag_id: str, run_id: str, task_id: str, key: str | None = None,
map_index: int = -1
+ ):
+ """Get task state store entries, or a single key, for a task instance
via API."""
+ endpoint =
f"dags/{dag_id}/dagRuns/{run_id}/taskInstances/{task_id}/state-store"
+ if key is not None:
+ endpoint = f"{endpoint}/{key}"
+ return self._make_request(method="GET",
endpoint=f"{endpoint}?map_index={map_index}")
def trigger_dag_and_wait(self, dag_id: str, json=None):
"""Trigger a DAG and wait for it to complete."""
diff --git
a/airflow-e2e-tests/tests/airflow_e2e_tests/ts_sdk_tests/test_ts_sdk_dag.py
b/airflow-e2e-tests/tests/airflow_e2e_tests/ts_sdk_tests/test_ts_sdk_dag.py
index 5424e04dad7..57a20d683fd 100644
--- a/airflow-e2e-tests/tests/airflow_e2e_tests/ts_sdk_tests/test_ts_sdk_dag.py
+++ b/airflow-e2e-tests/tests/airflow_e2e_tests/ts_sdk_tests/test_ts_sdk_dag.py
@@ -25,7 +25,8 @@ Two Dags mix Python tasks with ``@task.stub`` TypeScript
tasks, both served by t
``airflow-ts-pack`` bundle, and each is triggered once via a module-scoped
fixture.
``typescript_example`` covers the runtime: Variable reads and writes,
Connection reads,
-Python <-> TypeScript XCom round-trips, and task logs reaching the log store.
+Python <-> TypeScript XCom round-trips, task state store round-trips and
clearing, and task
+logs reaching the log store.
``typescript_taskflow_example`` covers TaskFlow arguments, including an
upstream output pulled before
the handler runs and a ``withArgNames`` rename on its ``report`` task, and
shares a ``build_message``
@@ -121,6 +122,8 @@ def test_task_states(completed_run: _CompletedRun):
"build_message": "success",
"read_connection": "success",
"write_and_delete_variable": "success",
+ "write_and_read_task_state": "success",
+ "clear_task_state": "success",
}
for task_id, want in expected.items():
assert completed_run.ti_states.get(task_id) == want, (
@@ -173,6 +176,66 @@ def
test_scratch_variable_deleted_by_typescript_task_is_gone(completed_run: _Com
)
+def test_task_state_store_round_trip(completed_run: _CompletedRun):
+ """``write_and_read_task_state`` writes two keys, reads them back, and
deletes one."""
+ value = completed_run.xcom("write_and_read_task_state")
+ assert value["jobId"] == completed_run.run_id, (
+ f"jobId should be this run's id {completed_run.run_id!r}, got
{value!r}"
+ )
+ assert value["scratchBeforeDelete"] == {"attempt": 1}, (
+ f"scratchBeforeDelete should be {{'attempt': 1}}, got {value!r}"
+ )
+ assert value["scratchAfterDelete"] is None, (
+ f"scratchAfterDelete should be None after taskStateStore.delete, got
{value!r}"
+ )
+
+ body = completed_run.client.get_task_state_store(
+ dag_id=completed_run.dag_id, run_id=completed_run.run_id,
task_id="write_and_read_task_state"
+ )
+ assert body["total_entries"] == 1, f"expected a single remaining task
state store entry, got {body!r}"
+ assert {entry["key"] for entry in body["task_state_store"]} ==
{"typescript_example_job_id"}, (
+ f"expected only the job id key to remain, got {body!r}"
+ )
+
+ entry = completed_run.client.get_task_state_store(
+ dag_id=completed_run.dag_id,
+ run_id=completed_run.run_id,
+ task_id="write_and_read_task_state",
+ key="typescript_example_job_id",
+ )
+ assert entry["value"] == completed_run.run_id, (
+ f"job id entry should hold this run's id {completed_run.run_id!r}, got
{entry!r}"
+ )
+ assert entry["expires_at"] is None, f"NEVER_EXPIRE should reach the
database as no expiry, got {entry!r}"
+
+
+def test_task_state_store_clear(completed_run: _CompletedRun):
+ """``clear_task_state`` removes every task state store key it wrote."""
+ value = completed_run.xcom("clear_task_state")
+ assert value["afterClear"] is None, f"afterClear should be None after
taskStateStore.clear, got {value!r}"
+
+ body = completed_run.client.get_task_state_store(
+ dag_id=completed_run.dag_id, run_id=completed_run.run_id,
task_id="clear_task_state"
+ )
+ assert body["total_entries"] == 0, f"taskStateStore.clear should remove
all entries, got {body!r}"
+ assert body["task_state_store"] == [], f"taskStateStore.clear should
remove all entries, got {body!r}"
+
+
+def test_task_state_store_deleted_key_is_gone(completed_run: _CompletedRun):
+ """A key written and then deleted from TypeScript no longer exists via the
REST API."""
+ with pytest.raises(requests.HTTPError) as excinfo:
+ completed_run.client.get_task_state_store(
+ dag_id=completed_run.dag_id,
+ run_id=completed_run.run_id,
+ task_id="write_and_read_task_state",
+ key="typescript_example_scratch",
+ )
+ assert excinfo.value.response.status_code == HTTPStatus.NOT_FOUND, (
+ f"typescript_example_scratch should have been deleted by the
TypeScript task, "
+ f"got HTTP {excinfo.value.response.status_code}"
+ )
+
+
def test_coordinator_logs_reach_task_log_store(completed_run: _CompletedRun):
assert "[ts-sdk.runtime] Coordinator runtime started" in
completed_run.logs("build_message")
diff --git a/ts-sdk/README.md b/ts-sdk/README.md
index 465f0ec5597..90e15609d74 100644
--- a/ts-sdk/README.md
+++ b/ts-sdk/README.md
@@ -179,12 +179,13 @@ both a Dag declared in TypeScript and the tasks of two
Python Dags from one bund
`getClient()` returns a `TaskClient` for task-time Airflow data access, for as
long as a handler is running:
-| Method |
Description |
-| --------------------------------------------------------------- |
----------------------- |
-| `getVariable(key)` / `getVariableOrThrow` | Airflow
Variables |
-| `setVariable(key, value, description?)` / `deleteVariable(key)` | Variable
write / delete |
-| `getXCom(opts)` / `setXCom(opts)` | XCom
read/write |
-| `getConnection(connId)` / `getConnectionOrThrow` | Airflow
Connections |
+| Method |
Description |
+| --------------------------------------------------------------- |
-------------------------------------- |
+| `getVariable(key)` / `getVariableOrThrow` | Airflow
Variables |
+| `setVariable(key, value, description?)` / `deleteVariable(key)` | Variable
write / delete |
+| `getXCom(opts)` / `setXCom(opts)` | XCom
read/write |
+| `getConnection(connId)` / `getConnectionOrThrow` | Airflow
Connections |
+| `taskStateStore.get / set / delete / clear` |
Per-task-instance key/value state store |
Locator fields such as `dagId`, `runId`, and `taskId` default to the
current task context when omitted.
@@ -227,7 +228,7 @@ Do not edit the table by hand. Update the manifest and run
the `update-ts-sdk-re
| capability: `variable-read-write` | MUST | ✓ | 3.4 | getVariable /
setVariable / deleteVariable |
| capability: `self-contained-bundle` | MUST | ✓ | 3.4 | Airflow metadata
embedded in the bundle |
| capability: `retry-policy` | MAY | ✗ | – | no task-facing retry-policy API
yet |
-| capability: `task-state-store` | MAY | ✗ | – | no task-facing state-store
API yet |
+| capability: `task-state-store` | MAY | ✓ | 3.4 | taskStateStore.get / set /
delete / clear |
| capability: `asset-state-store` | MAY | ✗ | – | no task-facing state-store
API yet |
| capability: `asset-event-emit` | MAY | ✗ | – | runtime does not emit asset
events yet |
| capability: `asset-event-read` | MAY | ✗ | – | no task-facing asset-event
API yet |
diff --git a/ts-sdk/api-docs/dag-authoring-api.ts
b/ts-sdk/api-docs/dag-authoring-api.ts
index 79b8ddf5e72..a6147025a30 100644
--- a/ts-sdk/api-docs/dag-authoring-api.ts
+++ b/ts-sdk/api-docs/dag-authoring-api.ts
@@ -19,7 +19,15 @@
/** @module Authoring */
-export { Bundle, Dag, getClient, getContext, TaskHandler, withArgNames } from
"../src/index.js";
+export {
+ Bundle,
+ Dag,
+ getClient,
+ getContext,
+ NEVER_EXPIRE,
+ TaskHandler,
+ withArgNames,
+} from "../src/index.js";
export type {
ArgNameMap,
DagSpec,
@@ -31,4 +39,5 @@ export type {
TaskOptions,
TaskRef,
TaskSpec,
+ TaskStateStore,
} from "../src/index.js";
diff --git a/ts-sdk/api-docs/supporting-types.ts
b/ts-sdk/api-docs/supporting-types.ts
index f928752d91d..3f5dc929e6b 100644
--- a/ts-sdk/api-docs/supporting-types.ts
+++ b/ts-sdk/api-docs/supporting-types.ts
@@ -19,4 +19,10 @@
/** @module Types */
-export type { ConnectionResult, GetXComOpts, JsonValue, SetXComOpts } from
"../src/index.js";
+export type {
+ ConnectionResult,
+ GetXComOpts,
+ JsonValue,
+ SetXComOpts,
+ TaskStateStoreSetOpts,
+} from "../src/index.js";
diff --git a/ts-sdk/capabilities.yaml b/ts-sdk/capabilities.yaml
index 170be9385cb..04f94159f82 100644
--- a/ts-sdk/capabilities.yaml
+++ b/ts-sdk/capabilities.yaml
@@ -86,8 +86,9 @@ capabilities:
supported: false
note: "no task-facing retry-policy API yet"
task-state-store:
- supported: false
- note: "no task-facing state-store API yet"
+ supported: true
+ since: "3.4"
+ note: "taskStateStore.get / set / delete / clear"
asset-state-store:
supported: false
note: "no task-facing state-store API yet"
diff --git a/ts-sdk/example/dags/typescript_example.py
b/ts-sdk/example/dags/typescript_example.py
index 67718ec4bcf..13bc5c34d3d 100644
--- a/ts-sdk/example/dags/typescript_example.py
+++ b/ts-sdk/example/dags/typescript_example.py
@@ -37,12 +37,22 @@ def read_connection(): ...
def write_and_delete_variable(): ...
[email protected](queue="typescript")
+def write_and_read_task_state(): ...
+
+
[email protected](queue="typescript")
+def clear_task_state(): ...
+
+
@dag(dag_id="typescript_example", schedule=None, catchup=False,
tags=["typescript", "example"])
def typescript_example():
start = python_start()
message = build_message()
read_connection()
write_and_delete_variable()
+ write_and_read_task_state()
+ clear_task_state()
start >> message
diff --git a/ts-sdk/example/src/main.ts b/ts-sdk/example/src/main.ts
index fa97d3f9899..d39c10b1c0a 100644
--- a/ts-sdk/example/src/main.ts
+++ b/ts-sdk/example/src/main.ts
@@ -22,7 +22,7 @@
// Both Dags in `dags/` are declared in Python with `@task.stub` tasks routed
to the Node
// coordinator, so this side only supplies the task bodies.
-import { Bundle, getClient, getContext, TaskHandler } from
"apache-airflow-ts-sdk";
+import { Bundle, getClient, getContext, NEVER_EXPIRE, TaskHandler } from
"apache-airflow-ts-sdk";
import { buildSummaryMessage, report, summarize } from "./taskflow.js";
@@ -57,6 +57,38 @@ export async function writeAndDeleteVariable() {
await client.deleteVariable(SCRATCH_VARIABLE);
}
+const JOB_ID_KEY = "typescript_example_job_id";
+const STATE_SCRATCH_KEY = "typescript_example_scratch";
+
+export async function writeAndReadTaskState() {
+ const store = getClient().taskStateStore;
+ const { runId, tryNumber } = getContext();
+
+ await store.set(JOB_ID_KEY, runId, { retentionMs: NEVER_EXPIRE });
+ await store.set(STATE_SCRATCH_KEY, { attempt: tryNumber });
+
+ const jobId = await store.get<string>(JOB_ID_KEY);
+ const scratchBeforeDelete = await store.get(STATE_SCRATCH_KEY);
+
+ await store.delete(STATE_SCRATCH_KEY);
+ const scratchAfterDelete = await store.get(STATE_SCRATCH_KEY);
+
+ return { jobId, scratchBeforeDelete, scratchAfterDelete };
+}
+
+export async function clearTaskState() {
+ const store = getClient().taskStateStore;
+
+ await store.set(JOB_ID_KEY, "placeholder");
+ await store.set(STATE_SCRATCH_KEY, { attempt: 1 });
+
+ await store.clear();
+
+ const afterClear = await store.get(JOB_ID_KEY);
+
+ return { afterClear };
+}
+
export async function readConnection() {
const connection = await
getClient().getConnection("typescript_example_http");
@@ -77,6 +109,8 @@ bundle.register(
new TaskHandler("typescript_example", "build_message", buildMessage),
new TaskHandler("typescript_example", "read_connection", readConnection),
new TaskHandler("typescript_example", "write_and_delete_variable",
writeAndDeleteVariable),
+ new TaskHandler("typescript_example", "write_and_read_task_state",
writeAndReadTaskState),
+ new TaskHandler("typescript_example", "clear_task_state", clearTaskState),
new TaskHandler("typescript_taskflow_example", "summarize", summarize),
new TaskHandler("typescript_taskflow_example", "report", report),
new TaskHandler("typescript_taskflow_example", "build_message",
buildSummaryMessage),
diff --git a/ts-sdk/src/coordinator/client.ts b/ts-sdk/src/coordinator/client.ts
index 8f47955dcdf..eaa49c687b1 100644
--- a/ts-sdk/src/coordinator/client.ts
+++ b/ts-sdk/src/coordinator/client.ts
@@ -20,8 +20,15 @@
import type { CommChannel } from "./comm-channel.js";
import type { LogChannel } from "./log-channel.js";
import type { TaskContext } from "../sdk/task.js";
-import type { TaskClient } from "../sdk/client.js";
-import type { ConnectionResult, GetXComOpts, JsonValue, SetXComOpts } from
"../sdk/client-types.js";
+import type { TaskClient, TaskStateStore } from "../sdk/client.js";
+import type {
+ ConnectionResult,
+ GetXComOpts,
+ JsonValue,
+ SetXComOpts,
+ TaskStateStoreSetOpts,
+} from "../sdk/client-types.js";
+import { NEVER_EXPIRE } from "../sdk/client-types.js";
import { ConnectionNotFoundError, VariableNotFoundError } from
"../sdk/client.js";
import type {
GetVariable,
@@ -34,12 +41,21 @@ import type {
TriggerDagRun,
GetDagRunState,
GetDag,
+ ClearTaskStateStore,
+ DeleteTaskStateStore,
+ GetTaskStateStore,
+ SetTaskStateStore,
ConnectionResult as WireConnectionResult,
} from "./protocol.js";
-/** What a supervisor "row is absent" error means for an operation: only a
lookup
- * can return `null`; for a `void` call a swallowed error would read as
success. */
-type AbsentRowPolicy = "null" | "throw";
+/**
+ * What a supervisor "row is absent" error means for an operation.
+ *
+ * `"throw"` always throws (a swallowed error on a `void` call would read as
+ * success); otherwise, the `ErrorType` codes that mean "absent" plus whether
+ * a wrapped API server 404 also counts (see `isAbsent`).
+ */
+type AbsentRowPolicy = "throw" | { codes: readonly string[]; apiServer404:
boolean };
function resolveWireMapIndex(
requestedMapIndex: number | null | undefined,
@@ -101,9 +117,91 @@ export interface CoordinatorClient extends TaskClient {
isDagPaused(dagId: string): Promise<boolean>;
}
+const VARIABLE_ABSENT_POLICY: AbsentRowPolicy = {
+ codes: ["VARIABLE_NOT_FOUND"],
+ apiServer404: true,
+};
+const XCOM_ABSENT_POLICY: AbsentRowPolicy = { codes: ["XCOM_NOT_FOUND"],
apiServer404: true };
+const CONNECTION_ABSENT_POLICY: AbsentRowPolicy = {
+ codes: ["CONNECTION_NOT_FOUND"],
+ apiServer404: true,
+};
+const TASK_STATE_STORE_ABSENT_POLICY: AbsentRowPolicy = {
+ codes: ["TASK_STORE_NOT_FOUND"],
+ apiServer404: false,
+};
+
+// A language SDK runtime cannot read Airflow config, so the coordinator passes
+// `[state_store] default_retention_days` at launch
+// (task-sdk/src/airflow/sdk/coordinators/_subprocess.py).
+const RETENTION_DAYS_ENV_VAR = "AIRFLOW__STATE_STORE__DEFAULT_RETENTION_DAYS";
+const MS_PER_DAY = 24 * 60 * 60 * 1000;
+
+function assertKey(key: string): void {
+ if (typeof key !== "string") {
+ throw new TypeError(`task state store key must be a string, got ${typeof
key}`);
+ }
+ if (key === "") {
+ throw new RangeError("task state store key must not be empty");
+ }
+}
+
+// Python's `datetime` (and the wire format it parses) cannot represent a year
+// past 9999, and the supervisor silently drops a frame it cannot decode
+// rather than replying with an error, so the task would hang instead of
+// failing fast. Route both retention paths through this check.
+function toExpiresAt(ms: number): string {
+ const at = new Date(Date.now() + ms);
+ if (Number.isNaN(at.getTime()) || at.getUTCFullYear() > 9999) {
+ throw new RangeError(
+ `retention of ${ms}ms overflows the wire timestamp; use NEVER_EXPIRE
instead`,
+ );
+ }
+ return at.toISOString();
+}
+
+/** Resolve `retentionMs` (or the deployment default) to a wire `expires_at`.
*/
+function resolveExpiresAt(retentionMs: number | typeof NEVER_EXPIRE |
undefined): string | null {
+ if (retentionMs === undefined) {
+ return resolveDefaultExpiresAt();
+ }
+ if (retentionMs === NEVER_EXPIRE) {
+ return null;
+ }
+ if (!Number.isFinite(retentionMs) || retentionMs < 0) {
+ throw new RangeError(`retentionMs must be >= 0 or NEVER_EXPIRE, got
${String(retentionMs)}`);
+ }
+ return toExpiresAt(retentionMs);
+}
+
+// The coordinator always passes the setting, so an absent or malformed value
is
+// a misconfiguration and fails the write rather than silently retaining the
key
+// for a period nobody configured.
+function resolveDefaultExpiresAt(): string | null {
+ const raw = process.env[RETENTION_DAYS_ENV_VAR];
+ if (raw === undefined) {
+ throw new Error(
+ `${RETENTION_DAYS_ENV_VAR} is not set, so the default retention is
unknown. It carries ` +
+ "[state_store] default_retention_days to the runtime; pass retentionMs
or NEVER_EXPIRE to " +
+ "set the expiry explicitly.",
+ );
+ }
+ const days = Number(raw);
+ if (raw.trim() === "" || !Number.isInteger(days)) {
+ throw new RangeError(`[state_store] default_retention_days must be a whole
number, got ${raw}`);
+ }
+ if (days < 0) {
+ throw new RangeError(
+ `[state_store] default_retention_days must be >= 0, got ${raw}. Set to 0
to disable expiry.`,
+ );
+ }
+ return days === 0 ? null : toExpiresAt(days * MS_PER_DAY);
+}
+
export function createCoordinatorClient(
comm: CommChannel,
ctx: TaskContext,
+ tiId: string,
logs: LogChannel | null = null,
): CoordinatorClient {
async function rpc<T>(
@@ -122,7 +220,7 @@ export function createCoordinatorClient(
logs?.debug(`${op} answered ${err.code}`);
return null;
}
- if (absent === "null" && isNotFound(err)) {
+ if (absent !== "throw" && isAbsent(err, absent)) {
logs?.debug(`${op} not found`, { error: err.code });
return null;
}
@@ -141,6 +239,57 @@ export function createCoordinatorClient(
return extract(body);
}
+ const taskStateStore: TaskStateStore = {
+ async get<T = unknown>(key: string): Promise<T | null> {
+ assertKey(key);
+ const msg: GetTaskStateStore = { type: "GetTaskStateStore", ti_id: tiId,
key };
+ return rpc(
+ "GetTaskStateStore",
+ "TaskStateStoreResult",
+ msg,
+ (body) => body!.value as T,
+ TASK_STATE_STORE_ABSENT_POLICY,
+ );
+ },
+
+ async set(
+ key: string,
+ value: NonNullable<JsonValue>,
+ opts: TaskStateStoreSetOpts = {},
+ ): Promise<void> {
+ assertKey(key);
+ if (
+ value === null ||
+ value === undefined ||
+ (typeof value === "number" && !Number.isFinite(value))
+ ) {
+ throw new TypeError("task state store value must not be null, NaN, or
Infinity");
+ }
+ // TODO: warn when the serialized value exceeds the deployment's
+ // [state_store] max_value_storage_bytes, as Python's task store setter
+ // does, once the coordinator passes that setting to the runtime.
+ const msg: SetTaskStateStore = {
+ type: "SetTaskStateStore",
+ ti_id: tiId,
+ key,
+ value,
+ expires_at: resolveExpiresAt(opts.retentionMs),
+ };
+ await rpc("SetTaskStateStore", "OKResponse", msg, () => undefined,
"throw");
+ },
+
+ async delete(key: string): Promise<void> {
+ assertKey(key);
+ const msg: DeleteTaskStateStore = { type: "DeleteTaskStateStore", ti_id:
tiId, key };
+ await rpc("DeleteTaskStateStore", "OKResponse", msg, () => undefined,
"throw");
+ },
+
+ async clear(): Promise<void> {
+ const msg: ClearTaskStateStore = { type: "ClearTaskStateStore", ti_id:
tiId };
+ await rpc("ClearTaskStateStore", "OKResponse", msg, () => undefined,
"throw");
+ },
+ };
+
const client: CoordinatorClient = {
// ---- Variables ----
@@ -151,7 +300,7 @@ export function createCoordinatorClient(
"VariableResult",
msg,
(body) => (body!.value as string) ?? null,
- "null",
+ VARIABLE_ABSENT_POLICY,
);
},
@@ -200,7 +349,7 @@ export function createCoordinatorClient(
// result frame decides `found` rather than the value.
value: (body!.value ?? null) as JsonValue,
}),
- "null",
+ XCOM_ABSENT_POLICY,
);
// `rpc` answers null for the supervisor's XCOM_NOT_FOUND.
return entry ?? XCOM_ABSENT;
@@ -271,6 +420,10 @@ export function createCoordinatorClient(
return paused!;
},
+ // ---- Task state store ----
+
+ taskStateStore,
+
// ---- Connections ----
async getConnection(connId: string): Promise<ConnectionResult | null> {
@@ -280,7 +433,7 @@ export function createCoordinatorClient(
"ConnectionResult",
msg,
(body) => fromWireConnection(body as unknown as WireConnectionResult),
- "null",
+ CONNECTION_ABSENT_POLICY,
);
},
@@ -296,8 +449,8 @@ export function createCoordinatorClient(
// -------- Error handling (two functions) --------
//
// parseFrameError: extract a structured error from the frame (once).
-// isNotFound: decide if the error means "absent" (a lookup returns null)
-// or "failed" (throw).
+// isAbsent: decide, per operation policy, if the error means "absent" (a
+// lookup returns null) or "failed" (throw).
interface FrameError {
code: string;
@@ -327,18 +480,13 @@ function parseFrameError(frame: { body: unknown; error?:
unknown }): FrameError
return null;
}
-// Exact supervisor ErrorType codes that mean "absent", not "failed".
-const NOT_FOUND_CODES = new Set(["VARIABLE_NOT_FOUND", "XCOM_NOT_FOUND",
"CONNECTION_NOT_FOUND"]);
-
-/** Is this error a "not found" (caller should get null, not a throw)?
- * Covers both dedicated NOT_FOUND codes and API_SERVER_ERROR with 404. */
-function isNotFound(err: FrameError): boolean {
- if (NOT_FOUND_CODES.has(err.code)) return true;
- // The supervisor wraps API server 404s as API_SERVER_ERROR with
- // detail.status_code=404 (supervisor.py: WatchedSubprocess.handle_requests).
- // Dag / Dag run lookups hit this path.
- // TODO: If the TS client adds lookups beyond variables, XCom, and
connections,
- // make not-found handling operation-specific instead of treating every
- // API_SERVER_ERROR 404 as null.
- return err.code === "API_SERVER_ERROR" && err.statusCode === 404;
+/** Is this error "absent" for `policy` (caller should get null, not a throw)?
+ * The supervisor wraps API server 404s as API_SERVER_ERROR with
+ * detail.status_code=404 (supervisor.py: WatchedSubprocess.handle_requests).
*/
+function isAbsent(
+ err: FrameError,
+ policy: { codes: readonly string[]; apiServer404: boolean },
+): boolean {
+ if (policy.codes.includes(err.code)) return true;
+ return policy.apiServer404 && err.code === "API_SERVER_ERROR" &&
err.statusCode === 404;
}
diff --git a/ts-sdk/src/coordinator/protocol.ts
b/ts-sdk/src/coordinator/protocol.ts
index 9f2f97b949c..3db41259fe7 100644
--- a/ts-sdk/src/coordinator/protocol.ts
+++ b/ts-sdk/src/coordinator/protocol.ts
@@ -72,6 +72,11 @@ export type {
DagRunStateResult,
GetDag,
DagResult,
+ ClearTaskStateStore,
+ DeleteTaskStateStore,
+ GetTaskStateStore,
+ SetTaskStateStore,
+ TaskStateStoreResult,
} from "../generated/supervisor.js";
// -------- Frames from supervisor --------
diff --git a/ts-sdk/src/coordinator/runtime.ts
b/ts-sdk/src/coordinator/runtime.ts
index 8c7227f8e70..9e4a2c3211e 100644
--- a/ts-sdk/src/coordinator/runtime.ts
+++ b/ts-sdk/src/coordinator/runtime.ts
@@ -347,7 +347,7 @@ async function handleTask(
const trigger = getBundleTrigger(bundle, ti.dag_id, ti.task_id);
if (trigger) {
const ctx = buildContext(details, signal);
- const client = createCoordinatorClient(comm, ctx, clientLogs);
+ const client = createCoordinatorClient(comm, ctx, ti.id, clientLogs);
const fail = (message: string) => {
logs.error("Task failed", { task_id: ctx.taskId, error: message });
return buildFailureResponse(details, message);
@@ -377,7 +377,7 @@ async function handleTask(
}
const ctx = buildContext(details, signal);
- const client = createCoordinatorClient(comm, ctx, clientLogs);
+ const client = createCoordinatorClient(comm, ctx, ti.id, clientLogs);
let bound: BoundArgs;
try {
diff --git a/ts-sdk/src/index.ts b/ts-sdk/src/index.ts
index f71bbbf9ac0..4c629bda76a 100644
--- a/ts-sdk/src/index.ts
+++ b/ts-sdk/src/index.ts
@@ -24,6 +24,7 @@ export { TaskHandler } from "./sdk/task-handler.js";
export { withArgNames } from "./sdk/arg-names.js";
export { getClient, getContext } from "./sdk/task.js";
export { ConnectionNotFoundError, VariableNotFoundError } from
"./sdk/client.js";
+export { NEVER_EXPIRE } from "./sdk/client-types.js";
export { SUPERVISOR_API_VERSION } from "./coordinator/index.js";
export type { ArgNameMap } from "./sdk/arg-names.js";
export type { Registerable } from "./sdk/bundle.js";
@@ -44,6 +45,12 @@ export type {
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 { TaskClient, TaskStateStore } from "./sdk/client.js";
+export type {
+ ConnectionResult,
+ GetXComOpts,
+ JsonValue,
+ SetXComOpts,
+ TaskStateStoreSetOpts,
+} from "./sdk/client-types.js";
export type { TaskContext, TaskFunction } from "./sdk/task.js";
diff --git a/ts-sdk/src/sdk/client-types.ts b/ts-sdk/src/sdk/client-types.ts
index 969e432006d..37c3b7b0e42 100644
--- a/ts-sdk/src/sdk/client-types.ts
+++ b/ts-sdk/src/sdk/client-types.ts
@@ -50,6 +50,29 @@ export interface SetXComOpts {
mapIndex?: number | null;
}
+/**
+ * Pass as {@link TaskStateStoreSetOpts.retentionMs} to store a task state
store
+ * key with no expiry, regardless of the deployment default.
+ *
+ * A symbol rather than a number, so a retention that reaches `Infinity` by
+ * accident (a division by zero, for example) is rejected instead of silently
+ * storing a key forever.
+ */
+export const NEVER_EXPIRE: unique symbol =
Symbol("apache-airflow-ts-sdk/NEVER_EXPIRE");
+
+/** Options for a task state store write. */
+export interface TaskStateStoreSetOpts {
+ /**
+ * Milliseconds to retain the key, counted from the write; `0` expires it
+ * immediately.
+ *
+ * Omitting it follows the deployment's `[state_store]
default_retention_days`,
+ * which the coordinator passes to the runtime. Pass {@link NEVER_EXPIRE} to
+ * store the key with no expiry regardless of that default.
+ */
+ retentionMs?: number | typeof NEVER_EXPIRE;
+}
+
/** Airflow Connection details returned by a task client. */
export interface ConnectionResult {
id: string;
diff --git a/ts-sdk/src/sdk/client.ts b/ts-sdk/src/sdk/client.ts
index 0facdd9695c..6927fc47831 100644
--- a/ts-sdk/src/sdk/client.ts
+++ b/ts-sdk/src/sdk/client.ts
@@ -17,7 +17,13 @@
* under the License.
*/
-import type { ConnectionResult, GetXComOpts, SetXComOpts } from
"./client-types.js";
+import type {
+ ConnectionResult,
+ GetXComOpts,
+ JsonValue,
+ SetXComOpts,
+ TaskStateStoreSetOpts,
+} from "./client-types.js";
/**
* Client for reading and writing Airflow task-time data from a task handler.
@@ -88,6 +94,13 @@ export interface TaskClient {
*/
setXCom(opts: SetXComOpts): Promise<void>;
+ /**
+ * Key/value state private to this task instance.
+ *
+ * See {@link TaskStateStore}.
+ */
+ readonly taskStateStore: TaskStateStore;
+
/**
* Look up an Airflow Connection by ID.
*
@@ -109,6 +122,65 @@ export interface TaskClient {
getConnectionOrThrow(connId: string): Promise<ConnectionResult>;
}
+/**
+ * Key/value state private to one task instance, kept across its retries within
+ * the same Dag run.
+ *
+ * Reach it through {@link TaskClient.taskStateStore}:
+ *
+ * ```ts
+ * const store = getClient().taskStateStore;
+ * let jobId = await store.get<string>("job_id");
+ * if (jobId == null) {
+ * jobId = await submit();
+ * await store.set("job_id", jobId, { retentionMs: NEVER_EXPIRE });
+ * }
+ * ```
+ *
+ * Values are stored in the metadata database as-is; the `[workers]
+ * state_store_backend` that Python tasks can configure is not applied here.
+ */
+export interface TaskStateStore {
+ /**
+ * Look up a value.
+ *
+ * Returns `null` when the key is missing. Throws on any other error.
+ *
+ * The generic `T` lets callers narrow the return type when the shape is
+ * known.
+ *
+ * @throws `TypeError` when `key` is not a string.
+ * @throws `RangeError` when `key` is empty.
+ */
+ get<T = unknown>(key: string): Promise<T | null>;
+
+ /**
+ * Store a value, replacing any existing value for the key.
+ *
+ * @throws `TypeError` when `key` is not a string, or `value` is null,
+ * undefined, NaN, or Infinity.
+ * @throws `RangeError` when `key` is empty; when `retentionMs` is
+ * negative, NaN, or Infinity; when the deployment's
+ * `[state_store] default_retention_days` is invalid; or when the resolved
+ * expiry would overflow the wire timestamp (year 9999).
+ */
+ set(key: string, value: NonNullable<JsonValue>, opts?:
TaskStateStoreSetOpts): Promise<void>;
+
+ /**
+ * Delete a key.
+ *
+ * Resolves even when the key does not exist — the Execution API's delete
+ * route is idempotent and does not report a missing key as an error.
+ *
+ * @throws `TypeError` when `key` is not a string.
+ * @throws `RangeError` when `key` is empty.
+ */
+ delete(key: string): Promise<void>;
+
+ /** Delete every key stored for this task instance. */
+ clear(): Promise<void>;
+}
+
/** Error thrown by {@link TaskClient.getVariableOrThrow}. */
export class VariableNotFoundError extends Error {
constructor(public readonly key: string) {
diff --git a/ts-sdk/tests/coordinator/client.test.ts
b/ts-sdk/tests/coordinator/client.test.ts
index cba7154b516..f5e864f6629 100644
--- a/ts-sdk/tests/coordinator/client.test.ts
+++ b/ts-sdk/tests/coordinator/client.test.ts
@@ -17,13 +17,16 @@
* under the License.
*/
-import { describe, it, expect } from "vitest";
+import { afterEach, describe, it, expect, vi } from "vitest";
import { ConnectionNotFoundError } from "../../src/sdk/client.js";
+import { NEVER_EXPIRE } from "../../src/sdk/client-types.js";
import { createCoordinatorClient } from "../../src/coordinator/client.js";
import type { CommChannel } from "../../src/coordinator/comm-channel.js";
import type { TaskClient } from "../../src/sdk/client.js";
import type { TaskContext } from "../../src/sdk/task.js";
+const RETENTION_ENV_VAR = "AIRFLOW__STATE_STORE__DEFAULT_RETENTION_DAYS";
+
function fakeComm(frames: { body: unknown; error?: unknown }[]): CommChannel {
let i = 0;
return {
@@ -40,8 +43,10 @@ const FAKE_CTX: TaskContext = {
signal: new AbortController().signal,
};
+const FAKE_TI_ID = "01890f3e-0000-7000-8000-000000000001";
+
function client(frames: { body: unknown; error?: unknown }[]) {
- return createCoordinatorClient(fakeComm(frames), FAKE_CTX);
+ return createCoordinatorClient(fakeComm(frames), FAKE_CTX, FAKE_TI_ID);
}
/** Captures the request bodies sent, answering each one with `reply`. */
@@ -53,7 +58,7 @@ function recordingClient(reply: unknown = null) {
return { body: reply };
},
} as unknown as CommChannel;
- return { client: createCoordinatorClient(comm, FAKE_CTX), sent };
+ return { client: createCoordinatorClient(comm, FAKE_CTX, FAKE_TI_ID), sent };
}
describe("getVariable not-found contract", () => {
@@ -157,6 +162,17 @@ describe("writes do not read a supervisor 404 as absence",
() => {
["setVariable", "PutVariable", (c: TaskClient) => c.setVariable("k", "v")],
["deleteVariable", "DeleteVariable", (c: TaskClient) =>
c.deleteVariable("k")],
["setXCom", "SetXCom", (c: TaskClient) => c.setXCom({ key: "k", value: 1
})],
+ [
+ "taskStateStore.set",
+ "SetTaskStateStore",
+ (c: TaskClient) => c.taskStateStore.set("k", 1, { retentionMs: 60_000 }),
+ ],
+ [
+ "taskStateStore.delete",
+ "DeleteTaskStateStore",
+ (c: TaskClient) => c.taskStateStore.delete("k"),
+ ],
+ ["taskStateStore.clear", "ClearTaskStateStore", (c: TaskClient) =>
c.taskStateStore.clear()],
])("%s rejects", async (_name, op, call) => {
const c = client([
{ body: null, error: { error: "API_SERVER_ERROR", detail: { status_code:
404 } } },
@@ -208,7 +224,7 @@ describe("client is bound to TaskContext", () => {
return { body: null };
},
} as unknown as CommChannel;
- const c = createCoordinatorClient(recordingComm, FAKE_CTX);
+ const c = createCoordinatorClient(recordingComm, FAKE_CTX, FAKE_TI_ID);
await c.setXCom({ key: "echo", value: 1 });
await c.setXCom({
@@ -243,7 +259,7 @@ describe("client is bound to TaskContext", () => {
return { body: { type: "XComResult", key: b.key, value: null } };
},
} as unknown as CommChannel;
- const c = createCoordinatorClient(recordingComm, FAKE_CTX);
+ const c = createCoordinatorClient(recordingComm, FAKE_CTX, FAKE_TI_ID);
await c.getXCom({ key: "k" });
@@ -272,7 +288,7 @@ describe("client is bound to TaskContext", () => {
// ctx with a real map index, to prove -1 from opts wins over a
// mapped ctx value (caller is explicitly asking "the non-mapped row").
const mappedCtx: TaskContext = { ...FAKE_CTX, mapIndex: 3 };
- const c = createCoordinatorClient(recordingComm, mappedCtx);
+ const c = createCoordinatorClient(recordingComm, mappedCtx, FAKE_TI_ID);
await c.setXCom({ key: "k", value: 1, mapIndex: -1 });
await c.setXCom({ key: "k", value: 2, mapIndex: null });
@@ -377,3 +393,213 @@ describe("getConnectionOrThrow", () => {
await
expect(c.getConnectionOrThrow("warehouse")).rejects.toThrow(/API_SERVER_ERROR/);
});
});
+
+describe("taskStateStore.get", () => {
+ it("sends GetTaskStateStore with the bound ti_id and returns the stored
value", async () => {
+ const { client: c, sent } = recordingClient({ type:
"TaskStateStoreResult", value: { n: 1 } });
+
+ await expect(c.taskStateStore.get("k")).resolves.toEqual({ n: 1 });
+
+ expect(sent[0]).toEqual({ type: "GetTaskStateStore", ti_id: FAKE_TI_ID,
key: "k" });
+ });
+
+ it("returns null for the exact TASK_STORE_NOT_FOUND code", async () => {
+ const c = client([{ body: { type: "ErrorResponse", error:
"TASK_STORE_NOT_FOUND" } }]);
+ expect(await c.taskStateStore.get("k")).toBeNull();
+ });
+
+ it("rejects for an API_SERVER_ERROR with status 404, which is not in its
absence policy", async () => {
+ const c = client([
+ { body: null, error: { error: "API_SERVER_ERROR", detail: { status_code:
404 } } },
+ ]);
+ await
expect(c.taskStateStore.get("k")).rejects.toThrow(/API_SERVER_ERROR/);
+ });
+
+ it("rejects for VARIABLE_NOT_FOUND, which is not in its absence policy",
async () => {
+ const c = client([{ body: { type: "ErrorResponse", error:
"VARIABLE_NOT_FOUND" } }]);
+ await
expect(c.taskStateStore.get("k")).rejects.toThrow(/VARIABLE_NOT_FOUND/);
+ });
+
+ it("rejects on an empty key without sending a request", async () => {
+ const { client: c, sent } = recordingClient();
+ await expect(c.taskStateStore.get("")).rejects.toThrow(RangeError);
+ expect(sent).toHaveLength(0);
+ });
+});
+
+describe("getVariable does not read TASK_STORE_NOT_FOUND as absence", () => {
+ it("rejects for TASK_STORE_NOT_FOUND, which is not in its absence policy",
async () => {
+ const c = client([{ body: { type: "ErrorResponse", error:
"TASK_STORE_NOT_FOUND" } }]);
+ await expect(c.getVariable("k")).rejects.toThrow(/TASK_STORE_NOT_FOUND/);
+ });
+});
+
+describe("taskStateStore.set retention", () => {
+ afterEach(() => {
+ vi.unstubAllEnvs();
+ vi.useRealTimers();
+ vi.restoreAllMocks();
+ });
+
+ it("sends expires_at null for NEVER_EXPIRE", async () => {
+ vi.useFakeTimers();
+ vi.setSystemTime(new Date("2026-01-01T00:00:00.000Z"));
+ const { client: c, sent } = recordingClient({ type: "OKResponse", ok: true
});
+
+ await c.taskStateStore.set("k", 1, { retentionMs: NEVER_EXPIRE });
+
+ expect(sent[0]).toEqual({
+ type: "SetTaskStateStore",
+ ti_id: FAKE_TI_ID,
+ key: "k",
+ value: 1,
+ expires_at: null,
+ });
+ });
+
+ it("sends now + retentionMs for a finite retentionMs", async () => {
+ vi.useFakeTimers();
+ vi.setSystemTime(new Date("2026-01-01T00:00:00.000Z"));
+ const { client: c, sent } = recordingClient({ type: "OKResponse", ok: true
});
+
+ await c.taskStateStore.set("k", 1, { retentionMs: 60_000 });
+
+ expect(sent[0]).toMatchObject({ expires_at: "2026-01-01T00:01:00.000Z" });
+ });
+
+ it("sends now (expire immediately) for retentionMs 0", async () => {
+ vi.useFakeTimers();
+ vi.setSystemTime(new Date("2026-01-01T00:00:00.000Z"));
+ const { client: c, sent } = recordingClient({ type: "OKResponse", ok: true
});
+
+ await c.taskStateStore.set("k", 1, { retentionMs: 0 });
+
+ expect(sent[0]).toMatchObject({ expires_at: "2026-01-01T00:00:00.000Z" });
+ });
+
+ it.each([
+ ["0", "0", null],
+ ["7", "7", "2026-01-08T00:00:00.000Z"],
+ ["7.0", "7.0", "2026-01-08T00:00:00.000Z"],
+ ])("resolves the default_retention_days env var: %s", async (_label, raw,
expected) => {
+ vi.useFakeTimers();
+ vi.setSystemTime(new Date("2026-01-01T00:00:00.000Z"));
+ vi.stubEnv(RETENTION_ENV_VAR, raw);
+ const { client: c, sent } = recordingClient({ type: "OKResponse", ok: true
});
+
+ await c.taskStateStore.set("k", 1);
+
+ expect(sent[0]).toMatchObject({ expires_at: expected });
+ });
+
+ it.each([
+ ["-1", "-1"],
+ ["abc", "abc"],
+ ["fractional", "7.5"],
+ ["empty", ""],
+ ])("rejects an invalid default_retention_days env var: %s", async (_label,
raw) => {
+ vi.stubEnv(RETENTION_ENV_VAR, raw);
+ const { client: c, sent } = recordingClient({ type: "OKResponse", ok: true
});
+
+ await expect(c.taskStateStore.set("k", 1)).rejects.toThrow(RangeError);
+ expect(sent).toHaveLength(0);
+ });
+
+ it("rejects when the coordinator did not pass default_retention_days", async
() => {
+ vi.stubEnv(RETENTION_ENV_VAR, undefined);
+ const { client: c, sent } = recordingClient({ type: "OKResponse", ok: true
});
+
+ await expect(c.taskStateStore.set("k",
1)).rejects.toThrow(RETENTION_ENV_VAR);
+ expect(sent).toHaveLength(0);
+ });
+});
+
+describe("taskStateStore.set guards", () => {
+ afterEach(() => {
+ vi.restoreAllMocks();
+ });
+
+ it.each([
+ ["null", null],
+ ["undefined", undefined],
+ ["NaN", Number.NaN],
+ ["Infinity", Number.POSITIVE_INFINITY],
+ ["-Infinity", Number.NEGATIVE_INFINITY],
+ ])("rejects a %s value without sending a request", async (_label, value) => {
+ const { client: c, sent } = recordingClient();
+ // @ts-expect-error exercising the runtime guard against an untyped caller.
+ await expect(c.taskStateStore.set("k", value)).rejects.toThrow(TypeError);
+ expect(sent).toHaveLength(0);
+ });
+
+ it.each([
+ ["negative", -1],
+ ["NaN", Number.NaN],
+ ])("rejects retentionMs=%s without sending a request", async (_label,
retentionMs) => {
+ const { client: c, sent } = recordingClient();
+ await expect(c.taskStateStore.set("k", 1, { retentionMs
})).rejects.toThrow(RangeError);
+ expect(sent).toHaveLength(0);
+ });
+
+ it("rejects an empty key without sending a request", async () => {
+ const { client: c, sent } = recordingClient();
+ await expect(c.taskStateStore.set("", 1, { retentionMs: 60_000
})).rejects.toThrow(RangeError);
+ expect(sent).toHaveLength(0);
+ });
+});
+
+describe("taskStateStore.delete", () => {
+ it("sends DeleteTaskStateStore and resolves on the supervisor's OKResponse",
async () => {
+ const { client: c, sent } = recordingClient({ type: "OKResponse", ok: true
});
+
+ await expect(c.taskStateStore.delete("k")).resolves.toBeUndefined();
+
+ expect(sent[0]).toEqual({ type: "DeleteTaskStateStore", ti_id: FAKE_TI_ID,
key: "k" });
+ });
+
+ it("rejects on an empty key without sending a request", async () => {
+ const { client: c, sent } = recordingClient();
+ await expect(c.taskStateStore.delete("")).rejects.toThrow(RangeError);
+ expect(sent).toHaveLength(0);
+ });
+});
+
+describe("task state store key type guard", () => {
+ it.each([
+ [
+ "get",
+ (c: TaskClient) => {
+ // @ts-expect-error exercising the runtime guard against an untyped
caller.
+ return c.taskStateStore.get(undefined);
+ },
+ ],
+ [
+ "set",
+ (c: TaskClient) => {
+ // @ts-expect-error exercising the runtime guard against an untyped
caller.
+ return c.taskStateStore.set(undefined, 1, { retentionMs: 60_000 });
+ },
+ ],
+ [
+ "delete",
+ (c: TaskClient) => {
+ // @ts-expect-error exercising the runtime guard against an untyped
caller.
+ return c.taskStateStore.delete(undefined);
+ },
+ ],
+ ])("%s rejects a non-string key without sending a request", async (_label,
call) => {
+ const { client: c, sent } = recordingClient();
+ await expect(call(c)).rejects.toThrow(TypeError);
+ expect(sent).toHaveLength(0);
+ });
+});
+
+describe("taskStateStore.clear", () => {
+ it("sends ClearTaskStateStore with no key and resolves on the supervisor's
OKResponse", async () => {
+ const { client: c, sent } = recordingClient({ type: "OKResponse", ok: true
});
+
+ await expect(c.taskStateStore.clear()).resolves.toBeUndefined();
+
+ expect(sent[0]).toEqual({ type: "ClearTaskStateStore", ti_id: FAKE_TI_ID
});
+ });
+});
diff --git a/ts-sdk/tests/coordinator/integration.test.ts
b/ts-sdk/tests/coordinator/integration.test.ts
index ee945c1c8aa..6ee7b02a0ee 100644
--- a/ts-sdk/tests/coordinator/integration.test.ts
+++ b/ts-sdk/tests/coordinator/integration.test.ts
@@ -759,6 +759,29 @@ describe("coordinator runtime integration", () => {
});
});
+ it("binds GetTaskStateStore's ti_id to the startup TaskInstance id, not the
dag version id", async () => {
+ let observed: unknown = "<unset>";
+ testDag.task("state_store_client", async () => {
+ observed = await getClient().taskStateStore.get("k");
+ });
+
+ const responder: Responder = (msgType) => {
+ if (msgType === "GetTaskStateStore") {
+ return { body: { type: "TaskStateStoreResult", value: 1 } };
+ }
+ return null;
+ };
+
+ const result = await
driveSupervisor(makeStartupDetails("state_store_client"), responder);
+
+ expect(result.firstResponse!.body).toMatchObject({ type: "SucceedTask" });
+ expect(observed).toBe(1);
+
+ const getReq = result.runtimeRequests.find((r) => r.type ===
"GetTaskStateStore")!.body;
+ expect(getReq).toMatchObject({ ti_id: "ti-1", key: "k" });
+ expect(getReq["ti_id"]).not.toBe("dag-version-1");
+ });
+
it("returns null from getVariable when the supervisor signals NOT_FOUND",
async () => {
let observed: string | null = "<unset>";
testDag.task("missing_variable", async () => {
diff --git a/ts-sdk/tests/public-api.test.ts b/ts-sdk/tests/public-api.test.ts
index b6cb9541627..4a9f91bcf12 100644
--- a/ts-sdk/tests/public-api.test.ts
+++ b/ts-sdk/tests/public-api.test.ts
@@ -24,6 +24,7 @@ import type {
ConnectionResult,
DagSpec,
GetXComOpts,
+ JsonValue,
Node,
SetXComOpts,
TaskClient,
@@ -38,6 +39,8 @@ import type {
TaskOptions,
TaskRef,
TaskSpec,
+ TaskStateStore,
+ TaskStateStoreSetOpts,
} from "../src/index.js";
import * as sdk from "../src/index.js";
import {
@@ -46,6 +49,7 @@ import {
Dag,
getClient,
getContext,
+ NEVER_EXPIRE,
SUPERVISOR_API_VERSION,
TaskHandler,
VariableNotFoundError,
@@ -286,6 +290,7 @@ describe("public API", () => {
expectTypeOf<typeof sdk>().not.toHaveProperty("serveDags");
expectTypeOf<typeof sdk>().not.toHaveProperty("DagRegistry");
expectTypeOf(SUPERVISOR_API_VERSION).toMatchTypeOf<string>();
+ expectTypeOf(NEVER_EXPIRE).toMatchTypeOf<symbol>();
});
it("keeps the Dag authoring signatures extensible via trailing specs", () =>
{
@@ -429,11 +434,24 @@ describe("public API", () => {
(key: string, value: string, description?: string | null) =>
Promise<void>
>();
expectTypeOf<TaskClient["deleteVariable"]>().toEqualTypeOf<(key: string)
=> Promise<void>>();
+ expectTypeOf<TaskStateStoreSetOpts>().toEqualTypeOf<{
+ retentionMs?: number | typeof NEVER_EXPIRE;
+ }>();
+
expectTypeOf<TaskClient["taskStateStore"]>().toEqualTypeOf<TaskStateStore>();
+ expectTypeOf<TaskStateStore["get"]>().toEqualTypeOf<
+ <T = unknown>(key: string) => Promise<T | null>
+ >();
+ expectTypeOf<TaskStateStore["set"]>().toEqualTypeOf<
+ (key: string, value: NonNullable<JsonValue>, opts?:
TaskStateStoreSetOpts) => Promise<void>
+ >();
+ expectTypeOf<TaskStateStore["delete"]>().toEqualTypeOf<(key: string) =>
Promise<void>>();
+ expectTypeOf<TaskStateStore["clear"]>().toEqualTypeOf<() =>
Promise<void>>();
});
it("rejects wire-format names and non-JSON XCom values", () => {
function acceptsGetXComOpts(_opts: GetXComOpts): void {}
function acceptsSetXComOpts(_opts: SetXComOpts): void {}
+ function acceptsTaskStateValue(_value: NonNullable<JsonValue>): void {}
acceptsGetXComOpts({
key: "result",
@@ -523,5 +541,9 @@ describe("public API", () => {
expectTypeOf<TaskRef>().toHaveProperty("handler");
// @ts-expect-error XCom values must be JSON-compatible.
acceptsSetXComOpts({ key: "result", value: new Date() });
+ // @ts-expect-error task-state-store values must be JSON-compatible.
+ acceptsTaskStateValue(new Date());
+ // @ts-expect-error task-state-store values must not be null.
+ acceptsTaskStateValue(null);
});
});
