From 9110115158b13c2eac88dfa04cc1d5623bf7ba94 Mon Sep 17 00:00:00 2001
From: Max Parke <max.parke@scale.com>
Date: Wed, 17 Jun 2026 12:02:37 -0400
Subject: [PATCH] fix(adk): re-send task_id/agent_id in state updates for
 backend compat
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

0.13.0's generated client dropped task_id/agent_id from states.update(),
and the ADK stopped sending them in the body. Backends predating
scale-agentex#278 still require those fields, so every state write 422'd
("Field required") — multi-turn context was lost on every turn (SEV,
~22k failures over ~21h on the Cengage agent).

Re-send both via extra_body. Backends post-#278 accept-and-ignore them,
so this is compatible with old and new servers. Adds a regression test
asserting StateService.update_state forwards them in the request body.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 src/agentex/lib/core/services/adk/state.py |  3 +
 tests/lib/adk/test_state_service.py        | 69 ++++++++++++++++++++++
 2 files changed, 72 insertions(+)
 create mode 100644 tests/lib/adk/test_state_service.py

diff --git a/src/agentex/lib/core/services/adk/state.py b/src/agentex/lib/core/services/adk/state.py
index 88a9b71b3..93012b933 100644
--- a/src/agentex/lib/core/services/adk/state.py
+++ b/src/agentex/lib/core/services/adk/state.py
@@ -98,9 +98,12 @@ async def update_state(
                 "state": state,
             },
         ) as span:
+            # Send task_id/agent_id in the body for backends predating
+            # scale-agentex#278, which still require them (newer ones ignore them).
             state_model = await self._agentex_client.states.update(
                 state_id=state_id,
                 state=state,
+                extra_body={"task_id": task_id, "agent_id": agent_id},
             )
             if span:
                 span.output = state_model.model_dump()
diff --git a/tests/lib/adk/test_state_service.py b/tests/lib/adk/test_state_service.py
new file mode 100644
index 000000000..43b53ff39
--- /dev/null
+++ b/tests/lib/adk/test_state_service.py
@@ -0,0 +1,69 @@
+"""Tests for StateService forwarding task_id/agent_id to the SDK client.
+
+Regression guard for the 0.13.0 incident: the generated client dropped
+task_id/agent_id from states.update(), so the ADK stopped sending them in the
+body and every state write 422'd against backends predating scale-agentex#278.
+"""
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from unittest.mock import Mock, AsyncMock
+
+from agentex.types.state import State
+from agentex.lib.core.services.adk.state import StateService
+
+_TS = datetime(2026, 5, 13, 18, 30, 0, tzinfo=timezone.utc)
+
+
+def _make_state() -> State:
+    return State(
+        id="s1",
+        agent_id="a1",
+        task_id="t1",
+        state={"k": "v"},
+        created_at=_TS,
+    )
+
+
+def _mock_span():
+    span = Mock()
+    span.output = None
+
+    async def __aenter__(_self):
+        return span
+
+    async def __aexit__(_self, *args):
+        return None
+
+    span.__aenter__ = __aenter__
+    span.__aexit__ = __aexit__
+    return span
+
+
+def _make_service() -> tuple[AsyncMock, StateService]:
+    client = AsyncMock()
+    tracer = Mock()
+    trace = Mock()
+    trace.span.return_value = _mock_span()
+    tracer.trace.return_value = trace
+    return client, StateService(agentex_client=client, tracer=tracer)
+
+
+class TestUpdateStateSendsParentIdentifiers:
+    async def test_task_id_and_agent_id_sent_in_body(self) -> None:
+        client, svc = _make_service()
+        client.states.update.return_value = _make_state()
+
+        await svc.update_state(
+            state_id="s1",
+            task_id="t1",
+            agent_id="a1",
+            state={"k": "v"},
+        )
+
+        kwargs = client.states.update.call_args.kwargs
+        assert kwargs["state_id"] == "s1"
+        # task_id/agent_id must ride in extra_body — the generated client dropped
+        # them from the typed signature, but old backends still require them.
+        assert kwargs["extra_body"] == {"task_id": "t1", "agent_id": "a1"}