fix(sdk): tighten async generator chunk normalization

Author: abhinav-galileo

sdks/python/src/agent_control/control_decorators.py

@@ -225,13 +225,12 @@ def output_payload(self) -> JSONValue:
         if not self.normalized_chunks:
             return ""
 
-        if all(isinstance(chunk, str) for chunk in self.normalized_chunks):
-            return "".join(chunk for chunk in self.normalized_chunks if isinstance(chunk, str))
+        text_chunks = [chunk for chunk in self.normalized_chunks if isinstance(chunk, str)]
+        if len(text_chunks) == len(self.normalized_chunks):
+            return "".join(text_chunks)
 
         output: dict[str, JSONValue] = {"chunks": list(self.normalized_chunks)}
-        text_output = "".join(
-            chunk for chunk in self.normalized_chunks if isinstance(chunk, str)
-        )
+        text_output = "".join(text_chunks)
         if text_output:
             output["text"] = text_output
         return output
@@ -276,11 +275,19 @@ def _normalize_json_value(value: Any, *, _seen: set[int] | None = None) -> JSONV
 
     model_dump = getattr(value, "model_dump", None)
     if callable(model_dump):
-        return _normalize_json_value(model_dump(mode="json"), _seen=seen)
+        seen.add(value_id)
+        try:
+            return _normalize_json_value(model_dump(mode="json"), _seen=seen)
+        finally:
+            seen.remove(value_id)
 
     dict_method = getattr(value, "dict", None)
     if callable(dict_method):
-        return _normalize_json_value(dict_method(), _seen=seen)
+        seen.add(value_id)
+        try:
+            return _normalize_json_value(dict_method(), _seen=seen)
+        finally:
+            seen.remove(value_id)
 
     if isinstance(value, Mapping):
         seen.add(value_id)
@@ -369,14 +376,6 @@ def _build_control_context(
     )
 
 
-async def _close_async_generator(stream: AsyncGenerator[Any, None]) -> None:
-    """Close a buffered async generator without masking the original failure."""
-    try:
-        await stream.aclose()
-    except Exception:
-        logger.debug("Failed to close buffered async generator cleanly", exc_info=True)
-
-
 async def _evaluate(
     agent_name: str,
     step: dict[str, Any],
@@ -946,7 +945,10 @@ def decorator(func: F) -> F:
         register(func, policy)
 
         @functools.wraps(func)
-        async def async_gen_wrapper(*args: Any, **kwargs: Any) -> Any:
+        async def async_gen_wrapper(
+            *args: Any,
+            **kwargs: Any,
+        ) -> AsyncGenerator[Any, None]:
             context = _build_control_context(func, args, kwargs, step_name)
             if context is None:
                 logger.warning(
@@ -965,12 +967,8 @@ async def async_gen_wrapper(*args: Any, **kwargs: Any) -> Any:
             try:
                 await _run_control_check(ctx, "pre", ctx.pre_payload(), controls)
 
-                try:
-                    async for chunk in stream:
-                        capture.add(chunk)
-                except BaseException:
-                    await _close_async_generator(stream)
-                    raise
+                async for chunk in stream:
+                    capture.add(chunk)
 
                 await _run_control_check(
                     ctx,

sdks/python/tests/test_control_decorators.py

@@ -800,6 +800,15 @@ class _StructuredChunk:
     value: str
 
 
+class _RecursiveModelDumpChunk:
+    def model_dump(self, mode: str = "json") -> dict[str, object]:
+        _ = mode
+        return {"self": self}
+
+    def __str__(self) -> str:
+        return "recursive-model-dump"
+
+
 class TestAsyncGeneratorControl:
     """Tests for buffered async-generator support in @control()."""
 
@@ -928,6 +937,42 @@ async def stream(message: str):
                 "text": "tail",
             }
 
+    @pytest.mark.asyncio
+    async def test_model_dump_cycle_is_guarded_during_chunk_normalization(
+        self,
+        mock_agent,
+        mock_safe_response,
+    ):
+        """Test that self-referential model_dump payloads do not recurse forever."""
+        post_steps = []
+
+        async def mock_evaluate(
+            agent_name,
+            step,
+            stage,
+            server_url,
+            trace_id=None,
+            span_id=None,
+            controls=None,
+            event_agent_name=None,
+        ):
+            if stage == "post":
+                post_steps.append(step)
+            return mock_safe_response
+
+        with patch("agent_control.control_decorators._get_current_agent", return_value=mock_agent), \
+             patch("agent_control.control_decorators._evaluate", side_effect=mock_evaluate):
+
+            @control()
+            async def stream(message: str):
+                yield _RecursiveModelDumpChunk()
+
+            chunks = [chunk async for chunk in stream("test")]
+            assert len(chunks) == 1
+            assert post_steps[0]["output"] == {
+                "chunks": [{"self": "recursive-model-dump"}]
+            }
+
     @pytest.mark.asyncio
     async def test_post_check_block_happens_before_any_chunk_is_replayed(
         self,