TS SDK: add the task state store to the task client

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,13 +523,64 @@ 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::
 
    A value from a secrets backend, such as an ``AIRFLOW_VAR_*`` environment variable, still wins over the
    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
 -------
 

airflow-e2e-tests/tests/airflow_e2e_tests/e2e_test_utils/clients.py

@@ -163,12 +163,14 @@ def get_variable(self, key: str):
         """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."""

airflow-e2e-tests/tests/airflow_e2e_tests/ts_sdk_tests/test_ts_sdk_dag.py

@@ -25,7 +25,8 @@
 ``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")
 

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 |

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";

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";

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"

ts-sdk/example/dags/typescript_example.py

@@ -37,12 +37,22 @@ def read_connection(): ...
 def write_and_delete_variable(): ...
 
 
+@task.stub(queue="typescript")
+def write_and_read_task_state(): ...
+
+
+@task.stub(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
 

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),