From 9b75e24d8c01878c153fec26ccfea4490417d23b Mon Sep 17 00:00:00 2001
From: Hangfei Lin <hangfei@google.com>
Date: Mon, 30 Jun 2025 09:49:18 -0700
Subject: [PATCH] fix: Avoid pydantic.ValidationError when the model stream
 returns empty final chunk

Bug: When a model emits a stream of tokens, it sometimes emits a final chunk of whitespace or no content. The agent was trying to parse that content into JSON, causing a validation error.

Fix: If a model is expected to return JSON and the last streamed token is empty/whitespace, the agent will no longer try to parse it, and exit gracefully.

New unit tests confirm the scenario and the fix.

PiperOrigin-RevId: 777609415
---
 src/google/adk/agents/llm_agent.py            |  5 ++++
 .../agents/test_llm_agent_output_save.py      | 30 +++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/src/google/adk/agents/llm_agent.py b/src/google/adk/agents/llm_agent.py
index 74091bf2..a5c859e2 100644
--- a/src/google/adk/agents/llm_agent.py
+++ b/src/google/adk/agents/llm_agent.py
@@ -451,6 +451,11 @@ class LlmAgent(BaseAgent):
           [part.text if part.text else '' for part in event.content.parts]
       )
       if self.output_schema:
+        # If the result from the final chunk is just whitespace or empty,
+        # it means this is an empty final chunk of a stream.
+        # Do not attempt to parse it as JSON.
+        if not result.strip():
+          return
         result = self.output_schema.model_validate_json(result).model_dump(
             exclude_none=True
         )
diff --git a/tests/unittests/agents/test_llm_agent_output_save.py b/tests/unittests/agents/test_llm_agent_output_save.py
index a13fde84..0f71f982 100644
--- a/tests/unittests/agents/test_llm_agent_output_save.py
+++ b/tests/unittests/agents/test_llm_agent_output_save.py
@@ -208,3 +208,33 @@ class TestLlmAgentOutputSave:
         "agent1",
         "agent2",
     )
+
+  @pytest.mark.parametrize("empty_content", ["", "  ", "\n"])
+  def test_maybe_save_output_to_state_handles_empty_final_chunk_with_schema(
+      self, empty_content
+  ):
+    """Tests that the agent correctly handles an empty final streaming chunk
+
+    when an output_schema is specified, preventing a crash.
+    """
+    # ARRANGE: Create an agent that expects a JSON output matching a schema.
+    agent = LlmAgent(
+        name="test_agent", output_key="result", output_schema=MockOutputSchema
+    )
+
+    # ARRANGE: Create a final event with empty or whitespace-only content.
+    # This simulates the final, empty chunk from a model's streaming response.
+    event = create_test_event(
+        author="test_agent", content_text=empty_content, is_final=True
+    )
+
+    # ACT: Call the method. The test's primary goal is to ensure this
+    # does NOT raise a pydantic.ValidationError, which it would have before the fix.
+    try:
+      agent._LlmAgent__maybe_save_output_to_state(event)
+    except Exception as e:
+      pytest.fail(f"The method unexpectedly raised an exception: {e}")
+
+    # ASSERT: Because the method should return early, the state_delta
+    # should remain empty.
+    assert len(event.actions.state_delta) == 0