jason810496 · Oct 6, 2026 · Sep 21, 2026 · Oct 5, 2026 · Oct 6, 2026 · Oct 6, 2026
diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst
@@ -709,6 +709,38 @@ Durations and date-times are ISO-8601 strings in annotations (``retryDelay = "PT
 ``java.time.OffsetDateTime`` values in ``config`` calls.  An unknown key or a mismatched value type
 fails the build for an annotation, and the ``config`` call itself for an object.
 
+.. _java-sdk/task-state-store:
+
+Task state store
+~~~~~~~~~~~~~~~~
+
+``client.getTaskStateStore()`` gives a task key-value state that is scoped to the task instance and
+survives retry attempts within the same Dag run (see :doc:`/core-concepts/task-state-store`). Use it to
+remember things like an external job ID so a retried task can resume instead of starting over:
+
+.. code-block:: java
+
+    @Builder.Task(id = "submit")
+    public void submit(Client client) throws Exception {
+      var store = client.getTaskStateStore();
+      var jobId = (String) store.get("job_id");
+      if (jobId == null) {
+        jobId = submitJob();
+        store.set("job_id", jobId, Duration.ofHours(6));
+      }
+      waitForJob(jobId);
+      store.delete("job_id");
+    }
+
+``get`` returns ``null`` when the key is not set. ``set`` stores any JSON-serializable value. Pass a positive
+``java.time.Duration`` to expire the key after that long, ``TaskStateStore.NEVER_EXPIRE`` for a key that
+garbage collection skips, or omit the retention to use the deployment's ``[state_store] default_retention_days``
+(0 means never expire). The coordinator passes that setting to the JVM as
+``AIRFLOW__STATE_STORE__DEFAULT_RETENTION_DAYS``. A zero or negative retention is rejected. ``delete`` removes
+one key and ``clear`` removes every key for the task instance. The Java SDK does not use a
+``[workers] state_store_backend``: values always go to the metadata database as-is, so keys written by Python
+tasks through a custom backend are returned to Java as the raw reference marker rather than the stored value.
+
 .. _java-sdk/native-dag-parsing:
 
 Parsing native Java Dags

@@ -630,7 +630,7 @@ prek hook regenerate it.
 | capability: `variable-read-write` | MUST | ✓ | 3.3 |  |
 | capability: `self-contained-bundle` | MUST | ✓ | 3.3 | Airflow metadata embedded in the jar artifact |
 | 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.3 | Client.getTaskStateStore() 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 |

@@ -83,8 +83,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.3"
+    note: "Client.getTaskStateStore() get/set/delete/clear"
   asset-state-store:
     supported: false
     note: "no task-facing state-store API yet"

@@ -51,7 +51,7 @@ meaning of each dimension is defined in the
 | capability: `variable-read-write` | MUST | ✓ | 3.3 |  |
 | capability: `self-contained-bundle` | MUST | ✓ | 3.3 | Airflow metadata embedded in the jar artifact |
 | 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.3 | Client.getTaskStateStore() 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 |

@@ -23,6 +23,11 @@ import org.apache.airflow.sdk.execution.ArgBinding
 import org.apache.airflow.sdk.execution.Client
 import org.apache.airflow.sdk.execution.comm.StartupDetails
 import org.apache.airflow.sdk.execution.decodeArgBindings
+import java.time.Duration
+import java.time.OffsetDateTime
+import java.time.ZoneOffset
+import java.time.temporal.ChronoUnit
+import kotlin.math.floor
 
 /**
  * A connection registered in Airflow's connection store.
@@ -57,6 +62,7 @@ data class Connection(
 class Client internal constructor(
   internal val details: StartupDetails,
   internal val impl: Client,
+  env: (String) -> String? = System::getenv,
 ) {
   internal companion object {
     /**
@@ -65,6 +71,14 @@ class Client internal constructor(
     const val XCOM_RETURN_KEY = "return_value"
   }
 
+  /**
+   * Key-value state scoped to the current task instance.
+   *
+   * Entries survive retries of the task instance within the same Dag run, so
+   * they can carry things like an external job ID across attempts.
+   */
+  val taskStateStore: TaskStateStore = TaskStateStore(details, impl, env)
+
   /**
    * Retrieves a connection from the Airflow connection store.
    *
@@ -215,6 +229,120 @@ class Client internal constructor(
   }
 }
 
+/**
+ * Key-value state scoped to one task instance, shared across its retries
+ * within the same Dag run.
+ *
+ * Values must be JSON-serializable. Every key has an expiry: by default the
+ * deployment's `[state_store] default_retention_days`, which the coordinator
+ * passes to the JVM as `AIRFLOW__STATE_STORE__DEFAULT_RETENTION_DAYS`; [set]
+ * also takes an explicit retention, or [NEVER_EXPIRE] for a key that garbage
+ * collection skips.
+ *
+ * Values are stored in the metadata database as-is; the `[workers]
+ * state_store_backend` used by Python tasks is not applied here.
+ */
+class TaskStateStore internal constructor(
+  private val details: StartupDetails,
+  private val impl: Client,
+  private val env: (String) -> String?,
+) {
+  companion object {
+    /**
+     * Pass as the retention of [set] to store a key that never expires and is
+     * skipped by Airflow's periodic garbage collection.
+     */
+    @JvmField val NEVER_EXPIRE: Duration = ChronoUnit.FOREVER.duration
+
+    internal const val DEFAULT_RETENTION_DAYS_ENV = "AIRFLOW__STATE_STORE__DEFAULT_RETENTION_DAYS"
+  }
+
+  /**
+   * Reads the value stored under [key].
+   *
+   * @return The stored value, or `null` if the key is not set.
+   * @throws ApiError if the API call fails.
+   */
+  fun get(key: String): Any? = impl.getTaskStateStore(details.ti.id, key)?.value
+
+  /**
+   * Stores [value] under [key], replacing any existing value.
+   *
+   * @param key State key.
+   * @param value Value to store. Must be JSON-serializable.
+   * @param retention How long to keep the key. Must be positive, or
+   *   [NEVER_EXPIRE]; `null` uses `[state_store] default_retention_days`.
+   * @throws IllegalArgumentException if [retention] is zero or negative, or the
+   *   default retention from the environment is not a non-negative integer.
+   * @throws IllegalStateException if [retention] is `null` and the coordinator
+   *   did not pass `[state_store] default_retention_days` to the JVM.
+   * @throws ApiError if the API call fails.
+   */
+  @JvmOverloads fun set(
+    key: String,
+    value: Any,
+    retention: Duration? = null,
+  ) {
+    val now = OffsetDateTime.now(ZoneOffset.UTC)
+    val expiresAt =
+      when {
+        retention == null -> resolveDefaultExpiry(now)
+        retention == NEVER_EXPIRE -> null
+        retention.isNegative || retention.isZero ->
+          throw IllegalArgumentException(
+            "Task state retention must be positive or TaskStateStore.NEVER_EXPIRE, got $retention for key '$key'",
+          )
+        else -> now.plus(retention)
+      }
+    // TODO: warn when the serialized value exceeds [state_store] max_value_storage_bytes once the
+    //   coordinator passes it to the JVM, as the Python accessor does.
+    impl.setTaskStateStore(tiId = details.ti.id, key = key, value = value, expiresAt = expiresAt)
+  }
+
+  /**
+   * Deletes the value stored under [key]. Does nothing if the key is not set.
+   *
+   * @throws ApiError if the API call fails.
+   */
+  fun delete(key: String) = impl.deleteTaskStateStore(details.ti.id, key)
+
+  /**
+   * Deletes every key stored for this task instance.
+   *
+   * @throws ApiError if the API call fails.
+   */
+  fun clear() = impl.clearTaskStateStore(details.ti.id)
+
+  private fun resolveDefaultExpiry(now: OffsetDateTime): OffsetDateTime? {
+    val raw =
+      env(DEFAULT_RETENTION_DAYS_ENV)
+        ?: throw IllegalStateException(
+          "$DEFAULT_RETENTION_DAYS_ENV is not set, so the default retention is unknown. The coordinator passes " +
+            "[state_store] default_retention_days to the JVM; pass a retention or TaskStateStore.NEVER_EXPIRE " +
+            "to set the expiry explicitly.",
+        )
+    val days = parseRetentionDays(raw)
+    return if (days == 0) null else now.plusDays(days.toLong())
+  }
+
+  // Accepts "7.0" because Python's conf.getint does.
+  private fun parseRetentionDays(raw: String): Int {
+    val days =
+      raw.trim().toIntOrNull()
+        ?: raw
+          .trim()
+          .toDoubleOrNull()
+          ?.takeIf { it.isFinite() && it == floor(it) }
+          ?.toInt()
+        ?: throw IllegalArgumentException(
+          "Failed to convert value to int. Please check 'default_retention_days' key in 'state_store' section. " +
+            "Current value: '$raw'",
+        )
+    require(days >= 0) { "[state_store] default_retention_days must be >= 0, got $days. Set to 0 to disable expiry." }
+    return days
+  }
+}
+
 /**
  * Thrown when a task's input resolves to nothing where a value is required —
  * a data parameter or a [TaskInput] field with a primitive type.

@@ -20,16 +20,24 @@
 package org.apache.airflow.sdk.execution
 
 import kotlinx.coroutines.runBlocking
+import org.apache.airflow.sdk.execution.comm.ClearTaskStateStore
 import org.apache.airflow.sdk.execution.comm.ConnectionResult
+import org.apache.airflow.sdk.execution.comm.DeleteTaskStateStore
 import org.apache.airflow.sdk.execution.comm.DeleteVariable
+import org.apache.airflow.sdk.execution.comm.ErrorResponse
 import org.apache.airflow.sdk.execution.comm.GetConnection
+import org.apache.airflow.sdk.execution.comm.GetTaskStateStore
 import org.apache.airflow.sdk.execution.comm.GetVariable
 import org.apache.airflow.sdk.execution.comm.GetXCom
 import org.apache.airflow.sdk.execution.comm.OKResponse
 import org.apache.airflow.sdk.execution.comm.PutVariable
+import org.apache.airflow.sdk.execution.comm.SetTaskStateStore
 import org.apache.airflow.sdk.execution.comm.SetXCom
+import org.apache.airflow.sdk.execution.comm.TaskStateStoreResult
 import org.apache.airflow.sdk.execution.comm.VariableResult
 import org.apache.airflow.sdk.execution.comm.XComResult
+import java.time.OffsetDateTime
+import java.util.UUID
 
 /**
  * @suppress
@@ -74,6 +82,26 @@ interface Client {
     runId: String,
     mapIndex: Int,
   )
+
+  /** Returns `null` when the key is not stored for the task instance. */
+  fun getTaskStateStore(
+    tiId: UUID,
+    key: String,
+  ): TaskStateStoreResult?
+
+  fun setTaskStateStore(
+    tiId: UUID,
+    key: String,
+    value: Any,
+    expiresAt: OffsetDateTime?,
+  )
+
+  fun deleteTaskStateStore(
+    tiId: UUID,
+    key: String,
+  )
+
+  fun clearTaskStateStore(tiId: UUID)
 }
 
 /**
@@ -157,4 +185,50 @@ class CoordinatorClient(
       }
     return runBlocking { exec.communicate<XComResult>(message) }
   }
+
+  override fun getTaskStateStore(
+    tiId: UUID,
+    key: String,
+  ): TaskStateStoreResult? {
+    val message =
+      GetTaskStateStore().also {
+        it.tiId = tiId
+        it.key = key
+      }
+    return runBlocking {
+      exec.communicateOrNullIf<TaskStateStoreResult>(message, ErrorResponse.ErrorType.TASK_STORE_NOT_FOUND)
+    }
+  }
+
+  override fun setTaskStateStore(
+    tiId: UUID,
+    key: String,
+    value: Any,
+    expiresAt: OffsetDateTime?,
+  ) {
+    val message =
+      SetTaskStateStore().also {
+        it.tiId = tiId
+        it.key = key
+        it.value = value
+        it.expiresAt = expiresAt
+      }
+    runBlocking { exec.communicate<OKResponse>(message) }
+  }
+
+  override fun deleteTaskStateStore(
+    tiId: UUID,
+    key: String,
+  ) {
+    val message =
+      DeleteTaskStateStore().also {
+        it.tiId = tiId
+        it.key = key
+      }
+    runBlocking { exec.communicate<OKResponse>(message) }
+  }
+
+  override fun clearTaskStateStore(tiId: UUID) {
+    runBlocking { exec.communicate<OKResponse>(ClearTaskStateStore().also { it.tiId = tiId }) }
+  }
 }
@@ -209,6 +209,22 @@ class CoordinatorComm(
     }
   }
 
+  /**
+   * Like [communicate], but maps an [ErrorResponse] of type [absent] to `null`
+   * so callers can treat "not found" as a value instead of an exception.
+   */
+  @Throws(ApiError::class)
+  suspend inline fun <reified T> communicateOrNullIf(
+    request: Any,
+    absent: ErrorResponse.ErrorType,
+  ): T? {
+    when (val response = communicateImpl(request)) {
+      is ErrorResponse -> if (response.error == absent) return null else throw ApiError("[${response.error}] ${response.detail}")
+      is T -> return response
+      else -> throw ApiError("Unexpected response type ${response::class.java}")
+    }
+  }
+
   /**
    * Stop the dispatcher and fail anything still awaiting a response.
    */

@@ -27,6 +27,7 @@ import com.fasterxml.jackson.databind.util.StdDateFormat
 import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule
 import org.apache.airflow.sdk.execution.comm.Discriminator
 import org.apache.airflow.sdk.execution.comm.PutVariable
+import org.apache.airflow.sdk.execution.comm.SetTaskStateStore
 import org.msgpack.core.MessagePack
 import org.msgpack.core.MessageUnpacker
 import org.msgpack.core.buffer.ArrayBufferInput
@@ -45,7 +46,7 @@ data class RawFrame(
  * fields to be present. When one is missing it fails validation, logs the
  * frame and never replies, so the caller blocks forever.
  */
-private val REQUIRED_NULLABLE_REQUESTS = setOf(PutVariable::class.java)
+private val REQUIRED_NULLABLE_REQUESTS = setOf(PutVariable::class.java, SetTaskStateStore::class.java)
 
 @JsonInclude(JsonInclude.Include.ALWAYS)
 private abstract class KeepNullFields