Warn about oversized task state values from language SDK tasks

go-sdk/pkg/execution/client.go

@@ -477,10 +477,6 @@ func (s taskStateStore) Set(
 		return fmt.Errorf("cannot set task state key %q: %w", key, err)
 	}
 
-	// TODO: warn when the serialized value exceeds the deployment's
-	// [state_store] max_value_storage_bytes, matching Python's
-	// airflow.sdk.execution_time.context task store setter.
-
 	_, err = s.client.comm.Communicate(ctx, genmodels.SetTaskStateStore{
 		TIID:      s.client.tiID,
 		Key:       key,

java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Client.kt

@@ -297,8 +297,6 @@ class TaskStateStore internal constructor(
           )
         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)
   }
 

task-sdk/src/airflow/sdk/coordinators/_subprocess.py

@@ -30,6 +30,7 @@
 import contextlib
 import ipaddress
 import itertools
+import json
 import os
 import selectors
 import signal
@@ -66,6 +67,8 @@
     from airflow.dag_processing.bundles.base import BaseDagBundle  # noqa: SDK002
     from airflow.sdk.api.client import Client
     from airflow.sdk.api.datamodels._generated import TaskInstance, TIRunContext
+    from airflow.sdk.execution_time.comms import SetTaskStateStore
+    from airflow.sdk.execution_time.supervisor import RequestResult
 
     Tracked = TypeVar("Tracked", socket.socket, subprocess.Popen)
 
@@ -445,6 +448,21 @@ def wait(self) -> int:
         self._close_unused_sockets(self._comm_server, self._logs_server)
         return code
 
+    def _handle_set_task_state_store(
+        self, msg: SetTaskStateStore, log: FilteringBoundLogger, req_id: int
+    ) -> RequestResult:
+        # A Python task warns in its own setter; a language SDK runtime cannot compute the size
+        # json.dumps gives, so the supervisor warns for it in the task log.
+        if (limit := conf.getint("state_store", "max_value_storage_bytes")) > 0:
+            if (size := len(json.dumps(msg.value))) > limit:
+                self.process_log.warning(
+                    "Task store value exceeds configured max_value_storage_bytes",
+                    key=msg.key,
+                    size=size,
+                    max_value_storage_bytes=limit,
+                )
+        return super()._handle_set_task_state_store(msg, log, req_id)
+
 
 def _initialize_pinned_bundle(target: BundleInfo, logger: FilteringBoundLogger) -> BaseDagBundle:
     """

task-sdk/tests/task_sdk/coordinators/test_subprocess.py

@@ -55,6 +55,7 @@
     _start_server,
     log,
 )
+from airflow.sdk.execution_time.comms import SetTaskStateStore
 from airflow.sdk.execution_time.coordinator import BaseCoordinator, TaskLaunchError
 from airflow.sdk.execution_time.supervisor import ActivitySubprocess
 from airflow.sdk.importers import DagSourceCode, reset_importer_registry
@@ -938,6 +939,49 @@ def test_register_pipe_readers_called_with_four_sockets(self, mock_client):
         assert mock_register.mock_calls == [call(ANY, ANY, ANY, ANY, data=ANY)]
 
 
+class TestPopenActivitySubprocessSetTaskStateStore:
+    @pytest.mark.parametrize(
+        ("limit", "value", "warns"),
+        [
+            # json.dumps gives '{"a": "\\u00e9"}', 15 bytes; Go's encoding/json gives 11.
+            pytest.param("14", {"a": "é"}, True, id="over-limit"),
+            pytest.param("15", {"a": "é"}, False, id="at-limit"),
+            pytest.param("0", {"a": "é"}, False, id="limit-disabled"),
+        ],
+    )
+    def test_warns_when_value_exceeds_max_value_storage_bytes(self, limit, value, warns):
+        client = MagicMock(spec=Client)
+        process_log = MagicMock()
+        proc = _PopenActivitySubprocess(
+            process_log=process_log,
+            id=uuid7(),
+            pid=12345,
+            stdin=MagicMock(),
+            process=MagicMock(),
+            client=client,
+            comm_server=MagicMock(spec=socket.socket),
+            logs_server=MagicMock(spec=socket.socket),
+        )
+        msg = SetTaskStateStore(ti_id=proc.id, key="k", value=value, expires_at=None)
+
+        with (
+            conf_vars({("state_store", "max_value_storage_bytes"): limit}),
+            patch.object(_PopenActivitySubprocess, "send_msg", autospec=True),
+        ):
+            proc._handle_request(msg, log=MagicMock(), req_id=1)
+
+        client.task_state_store.set.assert_called_once_with(proc.id, "k", value, expires_at=None)
+        if warns:
+            process_log.warning.assert_called_once_with(
+                "Task store value exceeds configured max_value_storage_bytes",
+                key="k",
+                size=15,
+                max_value_storage_bytes=14,
+            )
+        else:
+            process_log.warning.assert_not_called()
+
+
 # This script connects back to the supervisor the way a language SDK runtime does, then waits for
 # StartupDetails. It exits with code 0 as soon as the supervisor sends anything or closes the
 # connection. It exits with code 3 if 10 seconds pass and neither has happened.

ts-sdk/src/coordinator/client.ts

@@ -272,9 +272,6 @@ export function createCoordinatorClient(
       ) {
         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,