From 4edf844049e0659705d9096baf23e8d5b6f6a176 Mon Sep 17 00:00:00 2001 From: josix Date: Wed, 23 Sep 2026 00:42:31 +0800 Subject: [PATCH 1/4] 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/api-docs/dag-authoring-api.ts | 11 +- ts-sdk/api-docs/supporting-types.ts | 8 +- 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 ++ 11 files changed, 579 insertions(+), 38 deletions(-) diff --git a/ts-sdk/api-docs/dag-authoring-api.ts b/ts-sdk/api-docs/dag-authoring-api.ts index 79b8ddf5e7263..a6147025a3009 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 f928752d91d1b..3f5dc929e6bd9 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/src/coordinator/client.ts b/ts-sdk/src/coordinator/client.ts index 8f47955dcdf86..eaa49c687b1a9 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; } +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( @@ -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(key: string): Promise { + 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, + opts: TaskStateStoreSetOpts = {}, + ): Promise { + 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 { + assertKey(key); + const msg: DeleteTaskStateStore = { type: "DeleteTaskStateStore", ti_id: tiId, key }; + await rpc("DeleteTaskStateStore", "OKResponse", msg, () => undefined, "throw"); + }, + + async clear(): Promise { + 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 { @@ -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 9f2f97b949c3d..3db41259fe74e 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 8c7227f8e70e0..9e4a2c3211e38 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 f71bbbf9ac0a8..4c629bda76a1b 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 969e432006d4c..37c3b7b0e420a 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 0facdd9695c5d..6927fc4783108 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; + /** + * 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; } +/** + * 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("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(key: string): Promise; + + /** + * 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, opts?: TaskStateStoreSetOpts): Promise; + + /** + * 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; + + /** Delete every key stored for this task instance. */ + clear(): Promise; +} + /** 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 cba7154b51675..f5e864f6629cc 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 ee945c1c8aa64..6ee7b02a0ee36 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 = ""; + 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 = ""; testDag.task("missing_variable", async () => { diff --git a/ts-sdk/tests/public-api.test.ts b/ts-sdk/tests/public-api.test.ts index b6cb9541627a9..4a9f91bcf12ff 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().not.toHaveProperty("serveDags"); expectTypeOf().not.toHaveProperty("DagRegistry"); expectTypeOf(SUPERVISOR_API_VERSION).toMatchTypeOf(); + expectTypeOf(NEVER_EXPIRE).toMatchTypeOf(); }); 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 >(); expectTypeOf().toEqualTypeOf<(key: string) => Promise>(); + expectTypeOf().toEqualTypeOf<{ + retentionMs?: number | typeof NEVER_EXPIRE; + }>(); + expectTypeOf().toEqualTypeOf(); + expectTypeOf().toEqualTypeOf< + (key: string) => Promise + >(); + expectTypeOf().toEqualTypeOf< + (key: string, value: NonNullable, opts?: TaskStateStoreSetOpts) => Promise + >(); + expectTypeOf().toEqualTypeOf<(key: string) => Promise>(); + expectTypeOf().toEqualTypeOf<() => Promise>(); }); it("rejects wire-format names and non-JSON XCom values", () => { function acceptsGetXComOpts(_opts: GetXComOpts): void {} function acceptsSetXComOpts(_opts: SetXComOpts): void {} + function acceptsTaskStateValue(_value: NonNullable): void {} acceptsGetXComOpts({ key: "result", @@ -523,5 +541,9 @@ describe("public API", () => { expectTypeOf().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); }); }); From ff2c52ad3ab95b8ee7b5296eb37945a52e6ecb84 Mon Sep 17 00:00:00 2001 From: josix Date: Wed, 23 Sep 2026 00:42:50 +0800 Subject: [PATCH 2/4] 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. --- .../e2e_test_utils/clients.py | 14 ++-- .../ts_sdk_tests/test_ts_sdk_dag.py | 65 ++++++++++++++++++- ts-sdk/example/dags/typescript_example.py | 10 +++ ts-sdk/example/src/main.ts | 36 +++++++++- 4 files changed, 117 insertions(+), 8 deletions(-) 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 87b9b816637ca..273ab22532099 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 @@ 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.""" 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 5424e04dad7b2..57a20d683fde2 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 @@ ``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/example/dags/typescript_example.py b/ts-sdk/example/dags/typescript_example.py index 67718ec4bcf9d..13bc5c34d3df7 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(): ... +@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 diff --git a/ts-sdk/example/src/main.ts b/ts-sdk/example/src/main.ts index fa97d3f98996b..d39c10b1c0a42 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(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), From 481ec489309999696eb54144dfe09cdeaf52529e Mon Sep 17 00:00:00 2001 From: josix Date: Wed, 23 Sep 2026 00:43:29 +0800 Subject: [PATCH 3/4] 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. --- .../language-sdks/typescript.rst | 53 ++++++++++++++++++- ts-sdk/README.md | 15 +++--- ts-sdk/capabilities.yaml | 5 +- 3 files changed, 63 insertions(+), 10 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 d2c6094a105dc..e91149020f9e7 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(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, intra-task checkpoints, 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("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/ts-sdk/README.md b/ts-sdk/README.md index 465f0ec55971f..90e15609d74c6 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/capabilities.yaml b/ts-sdk/capabilities.yaml index 170be9385cb3d..04f94159f82e6 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" From ee216ea089fd8214dc4c2549305a6d6e6cc47d6f Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Fri, 9 Oct 2026 08:21:04 +0000 Subject: [PATCH 4/4] Fix spelling in TS SDK task state store docs Co-Authored-By: Claude --- .../docs/authoring-and-scheduling/language-sdks/typescript.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 e91149020f9e7..0cb37f578e2ef 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst @@ -538,7 +538,7 @@ 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, intra-task checkpoints, and progress metadata. See :doc:`/core-concepts/task-state-store` +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