jason810496 · Oct 13, 2026 · Oct 7, 2026
@@ -20,12 +20,14 @@
 package org.apache.airflow.e2e;
 
 import java.util.List;
+import java.util.Map;
 import org.apache.airflow.sdk.*;
 import org.jetbrains.annotations.NotNull;
 
 /**
  * Bundle for the runner-behaviour E2E tests: deliberately broken task classes that exercise
- * instantiation failures, and a task that round-trips Airflow Variables through the supervisor.
+ * instantiation failures, and tasks that round-trip Airflow Variables and task state through the
+ * supervisor.
  */
 public class TestBundleBuilder implements BundleBuilder {
   public static class MissingNoArgConstructor implements Task {
@@ -60,6 +62,34 @@ public void execute(@NotNull Context context, Client client) {
     }
   }
 
+  /**
+   * Writes task state keys the E2E test reads back through the REST API, and reads them back
+   * here to exercise the get path. No clear: the E2E test reads the keys after the task ends.
+   */
+  public static class RoundtripTaskState implements Task {
+    public void execute(@NotNull Context context, Client client) {
+      var store = client.getTaskStateStore();
+      var runId = context.dagRun.runId;
+      store.set("java_e2e_run_id", runId, TaskStateStore.NEVER_EXPIRE);
+      store.set("java_e2e_counter", Map.of("processed", 3, "cursor", "abc-123"));
+      store.set("java_e2e_retained", "retained");
+      store.set("java_e2e_scratch", "scratch");
+      store.delete("java_e2e_scratch");
+
+      var readBack = store.get("java_e2e_run_id");
+      if (!runId.equals(readBack)) {
+        throw new IllegalStateException("java_e2e_run_id: got " + readBack + ", want " + runId);
+      }
+      var counter = store.get("java_e2e_counter");
+      if (!Map.of("processed", 3L, "cursor", "abc-123").equals(counter)) {
+        throw new IllegalStateException("java_e2e_counter: got " + counter);
+      }
+      if (store.get("java_e2e_scratch") != null) {
+        throw new IllegalStateException("java_e2e_scratch survived its delete");
+      }
+    }
+  }
+
   @NotNull
   @Override
   public Iterable<DagDef> getDags() {
@@ -68,7 +98,9 @@ public Iterable<DagDef> getDags() {
     uninstantiable.addTask("non_static_inner", NonStaticInner.class);
     var variableWrite = new DagDef("java_variable_write");
     variableWrite.addTask("write_and_delete", WriteAndDeleteVariable.class);
-    return List.of(uninstantiable, variableWrite);
+    var taskState = new DagDef("java_task_state");
+    taskState.addTask("roundtrip_task_state", RoundtripTaskState.class);
+    return List.of(uninstantiable, variableWrite, taskState);
   }
 
   public static void main(String[] args) {

@@ -48,3 +48,15 @@ def java_variable_write():
 
 
 java_variable_write()
+
+
+@task.stub(queue="java-test")
+def roundtrip_task_state(): ...
+
+
+@dag(dag_id="java_task_state")
+def java_task_state():
+    roundtrip_task_state()
+
+
+java_task_state()
@@ -45,7 +45,6 @@
     GO_SDK_DAGS_PATH,
     GO_SDK_EXAMPLE_BUNDLE_PKG,
     GO_SDK_ROOT_PATH,
-    GO_SDK_STATE_STORE_RETENTION_DAYS,
     JAVA_COMPOSE_PATH,
     JAVA_DOCKERFILE_PATH,
     JAVA_SDK_EXAMPLE_DAGS_PATH,
@@ -57,6 +56,7 @@
     JAVA_TEST_BUNDLE_ROOT_PATH,
     KAFKA_DIR_PATH,
     LANG_SDK_NATIVE_TOOLCHAIN,
+    LANG_SDK_STATE_STORE_RETENTION_DAYS,
     LOCALSTACK_PATH,
     LOGS_FOLDER,
     NODE_IMAGE,
@@ -310,6 +310,14 @@ def _build_dag_bundle_config(artifact_bundles: dict[str, str]) -> str:
     return json.dumps(bundles)
 
 
+def _write_state_store_retention_config(tmp_dir):
+    # Config file, not an env var: the supervisor copies the worker environment into the
+    # language SDK subprocess, so an env var would reach it even without the propagation under test.
+    (tmp_dir / "config" / "airflow.cfg").write_text(
+        f"[state_store]\ndefault_retention_days = {LANG_SDK_STATE_STORE_RETENTION_DAYS}\n"
+    )
+
+
 def _run_java_sdk_gradle(workdir, *gradle_argv, capture_output=False, native=False):
     """Run the Java SDK Gradle wrapper natively or inside the pinned JDK container.
 
@@ -534,6 +542,7 @@ def _setup_java_sdk_integration(dot_env_file, tmp_dir):
         "AIRFLOW_VAR_MY_VARIABLE=test_value\n"
     )
     os.environ["ENV_FILE_PATH"] = str(dot_env_file)
+    _write_state_store_retention_config(tmp_dir)
 
 
 def _run_go_sdk_pack(output_path, *, capture_output=False, native=False):
@@ -674,11 +683,7 @@ def _setup_go_sdk_integration(dot_env_file, tmp_dir):
     )
     os.environ["ENV_FILE_PATH"] = str(dot_env_file)
 
-    # Config file, not an env var: the supervisor copies the worker environment into the Go
-    # subprocess, so an env var would reach Go even without the propagation under test.
-    (tmp_dir / "config" / "airflow.cfg").write_text(
-        f"[state_store]\ndefault_retention_days = {GO_SDK_STATE_STORE_RETENTION_DAYS}\n"
-    )
+    _write_state_store_retention_config(tmp_dir)
 
 
 def _setup_openlineage_integration(dot_env_file, tmp_dir, compose_file_names):

@@ -62,6 +62,9 @@
 # CI sets this (the same switch the lang-SDK k8s job uses) to build the lang-SDK
 # artifacts with the host toolchain instead of ephemeral toolchain containers.
 LANG_SDK_NATIVE_TOOLCHAIN = os.environ.get("LANG_SDK_NATIVE_TOOLCHAIN", "").lower() == "true"
+# Far from the supervisor's own 30-day config fallback, so the e2e test can tell a
+# propagated value from it.
+LANG_SDK_STATE_STORE_RETENTION_DAYS = 7
 
 # Java SDK E2E test paths
 JAVA_SDK_ROOT_PATH = AIRFLOW_ROOT_PATH / "java-sdk"
@@ -95,9 +98,6 @@
 # Where airflow-go-pack writes the packed bundle inside the repo (go-sdk/bin is gitignored).
 GO_SDK_BIN_PATH = GO_SDK_ROOT_PATH / "bin"
 GO_COMPOSE_PATH = AIRFLOW_ROOT_PATH / "airflow-e2e-tests" / "docker" / "go.yml"
-# Far from the supervisor's own 30-day config fallback, so the e2e test can tell a
-# propagated value from it.
-GO_SDK_STATE_STORE_RETENTION_DAYS = 7
 # Go toolchain image used to build the bundle in the containerized path (i.e. unless
 # LANG_SDK_NATIVE_TOOLCHAIN is set); must satisfy go-sdk/go.mod's toolchain.
 # The Alpine variant is ~7x smaller than the Debian one and is safe here because the

@@ -29,7 +29,7 @@
 import pytest
 import requests
 
-from airflow_e2e_tests.constants import GO_SDK_STATE_STORE_RETENTION_DAYS
+from airflow_e2e_tests.constants import LANG_SDK_STATE_STORE_RETENTION_DAYS
 from airflow_e2e_tests.e2e_test_utils.clients import AirflowClient
 
 _GO_TASK_TIMEOUT = 300
@@ -106,9 +106,9 @@ def test_default_retention_applied(completed_run: _CompletedRun):
     entry = completed_run.state_store("go_e2e_retained")
     assert entry.get("expires_at") is not None, entry
     gap = _parse_timestamp(entry["expires_at"]) - _parse_timestamp(entry["updated_at"])
-    expected = timedelta(days=GO_SDK_STATE_STORE_RETENTION_DAYS)
+    expected = timedelta(days=LANG_SDK_STATE_STORE_RETENTION_DAYS)
     assert abs(gap - expected) <= timedelta(minutes=5), (
-        f"expected ~{GO_SDK_STATE_STORE_RETENTION_DAYS} days, got {gap / timedelta(days=1):.1f} days; "
+        f"expected ~{LANG_SDK_STATE_STORE_RETENTION_DAYS} days, got {gap / timedelta(days=1):.1f} days; "
         "if ~30, the supervisor fell back to its own config default, so it did not see "
         f"[state_store] default_retention_days. entry: {entry!r}"
     )
@@ -69,13 +69,14 @@
 
 import time
 from dataclasses import dataclass, field
-from datetime import UTC, datetime
+from datetime import UTC, datetime, timedelta
 from http import HTTPStatus
 from typing import TYPE_CHECKING
 
 import pytest
 import requests
 
+from airflow_e2e_tests.constants import LANG_SDK_STATE_STORE_RETENTION_DAYS
 from airflow_e2e_tests.e2e_test_utils.clients import AirflowClient
 
 if TYPE_CHECKING:
@@ -95,6 +96,8 @@
 _SCALA_SPARK_DAG_ID = "scala_spark_example"
 _VARIABLE_WRITE_DAG_ID = "java_variable_write"
 _VARIABLE_WRITE_DESCRIPTION = "written by the Java SDK e2e test"
+_TASK_STATE_DAG_ID = "java_task_state"
+_TASK_STATE_TASK_ID = "roundtrip_task_state"
 
 
 @dataclass
@@ -128,6 +131,14 @@ def get_xcom(self, task_id: str, key: str = "return_value"):
             key=key,
         ).get("value")
 
+    def get_task_state_store(self, task_id: str, key: str) -> dict:
+        self._raise_if_failed()
+        if self.run_id is None:
+            raise RuntimeError(f"Dag {self.dag_id!r} did not produce a run ID")
+        return self.client.get_task_state_store(
+            dag_id=self.dag_id, run_id=self.run_id, task_id=task_id, key=key
+        )
+
     def wait_for_log_record(
         self, task_id: str, try_number: int, match: Callable[[dict], bool]
     ) -> tuple[dict | None, list[dict]]:
@@ -175,6 +186,10 @@ def _trigger_and_wait_for_dag(dag_id: str, timeout: int) -> _CompletedRun:
     )
 
 
+def _parse_timestamp(value: str) -> datetime:
+    return datetime.fromisoformat(value.replace("Z", "+00:00"))
+
+
 @pytest.fixture(scope="module")
 def annotation_example_run() -> _CompletedRun:
     """Trigger the annotation example once for all of its assertions."""
@@ -205,6 +220,12 @@ def variable_write_run() -> _CompletedRun:
     return _trigger_and_wait_for_dag(_VARIABLE_WRITE_DAG_ID, _JAVA_TASK_TIMEOUT)
 
 
+@pytest.fixture(scope="module")
+def task_state_run() -> _CompletedRun:
+    """Trigger the task state Dag once for all of its assertions."""
+    return _trigger_and_wait_for_dag(_TASK_STATE_DAG_ID, _JAVA_TASK_TIMEOUT)
+
+
 class TestJavaSDKAnnotationExample:
     """Verify the annotation-based Java SDK example executes correctly."""
 
@@ -480,3 +501,52 @@ def test_spark_etl_pipeline(self, scala_spark_example_run: _CompletedRun):
         assert load_xcom == _SPARK_EXPECTED_TOTAL_REVENUE, (
             f"Expected spark_load to return total revenue {_SPARK_EXPECTED_TOTAL_REVENUE}, got {load_xcom!r}"
         )
+
+
+class TestJavaSDKTaskStateStore:
+    """Verify a Java task can set, get and delete task state store keys.
+
+    The task lives in the java-test-bundle fixture project (served on the
+    dedicated "java-test" queue). It also reads its keys back and fails if a
+    value does not match, so ``success`` covers the get path. Assertions here
+    read the store through the REST API, so the data is proven to have reached
+    the database.
+    """
+
+    def _assert_task_succeeded(self, run: _CompletedRun) -> None:
+        ti = run.get_task_instance(_TASK_STATE_TASK_ID)
+        assert ti.get("state") == "success", (
+            f"Java {_TASK_STATE_TASK_ID!r} task did not succeed.\n"
+            f"  task state : {ti.get('state')!r}\n"
+            f"  dag state  : {run.state!r}\n"
+            f"  all tasks  : {run.ti_states}"
+        )
+
+    def test_never_expire_key_holds_run_id(self, task_state_run: _CompletedRun):
+        self._assert_task_succeeded(task_state_run)
+        entry = task_state_run.get_task_state_store(_TASK_STATE_TASK_ID, "java_e2e_run_id")
+        assert entry.get("value") == task_state_run.run_id, entry
+        assert entry.get("expires_at") is None, entry
+
+    def test_structured_value_roundtrips(self, task_state_run: _CompletedRun):
+        self._assert_task_succeeded(task_state_run)
+        entry = task_state_run.get_task_state_store(_TASK_STATE_TASK_ID, "java_e2e_counter")
+        assert entry.get("value") == {"processed": 3, "cursor": "abc-123"}, entry
+
+    def test_deleted_key_is_gone(self, task_state_run: _CompletedRun):
+        self._assert_task_succeeded(task_state_run)
+        with pytest.raises(requests.HTTPError) as excinfo:
+            task_state_run.get_task_state_store(_TASK_STATE_TASK_ID, "java_e2e_scratch")
+        assert excinfo.value.response.status_code == HTTPStatus.NOT_FOUND
+
+    def test_default_retention_applied(self, task_state_run: _CompletedRun):
+        self._assert_task_succeeded(task_state_run)
+        entry = task_state_run.get_task_state_store(_TASK_STATE_TASK_ID, "java_e2e_retained")
+        assert entry.get("expires_at") is not None, entry
+        gap = _parse_timestamp(entry["expires_at"]) - _parse_timestamp(entry["updated_at"])
+        expected = timedelta(days=LANG_SDK_STATE_STORE_RETENTION_DAYS)
+        assert abs(gap - expected) <= timedelta(minutes=5), (
+            f"expected ~{LANG_SDK_STATE_STORE_RETENTION_DAYS} days, got {gap / timedelta(days=1):.1f} days; "
+            "if ~30, the coordinator passed Airflow's config default, so the JVM did not see "
+            f"[state_store] default_retention_days. entry: {entry!r}"
+        )