Skip to content

Fix(workflow)/add _process_content_object function in _rehydration_utils file to extract output from event.content object before assigning it to child.output, in _reconstruct_node_states#5909

Open
samarth1224 wants to merge 1 commit into
google:mainfrom
samarth1224:fix/issue-5553

Conversation

@samarth1224
Copy link
Copy Markdown

Link to Issue

1. Link to an existing issue (if applicable):

NOTE

I opened the pull request before but it was closed due to v2 branch being merged into main branch.
The previous test failures in last PR were due to inherent issues in the v2 branche's unittests.

Problem:

During fresh execution of a node with message_as_output == True, process_llm_agent_output extracts text from the Content parts, parses it, validates it against the schema, and assigns the resulting output to event.output.

However, before the event is persisted to the session, _consume_event_queue optimizes for message_as_output nodes by stripping event.output (setting it to None) and only saving the raw event.content

Upon workflow resumption, _reconstruct_node_states rebuilds the node's state. Because event.output is None, it fell back to assigning the raw event.content (a Content object) directly to child.output.

Solution:

Added _process_content_object() to extract the output from the raw event.content during rehydration.
This function mirrors the logic used in the process_llm_agent_output function in _llm_agent_wrapper.py` file

  1. It extracts text from all parts of the Content object.
  2. It correctly filters out Parts.thought parts to prevent Chain-of-Thought reasoning text from leaking into the final output.
  3. It attempts to parse the extracted text as JSON, if fails it returns the raw text.

Testing Plan

I have added new unit tests for the _process_content_object function. These tests cover various scenarios, including plain text extraction, JSON parsing for structured outputs, and the correct filtering of thought parts to prevent internal reasoning from leaking into the final output.

Additionally, I have updated the existing rehydration test _test_scan_message_as_output that previously asserted the raw Content object was returned; It now correctly verify that the output is reconstructed during node state rehydration.

Unit Tests:

  • I have added or updated unit tests for my change.
  • All unit tests pass locally.

Summary Unit Test


tests\unittests\workflow\utils\test_rehydration_utils.py ........................................                [100%]

================================================== warnings summary ===================================================
<frozen abc>:106
<frozen abc>:106
<frozen abc>:106
<frozen abc>:106
  <frozen abc>:106: DeprecationWarning: BaseAgentConfig is deprecated and will be removed in future versions.

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
=========================================== 40 passed, 4 warnings in 10.20s ===========================================

Manual End-to-End (E2E) Tests:

I have already provided the necessary setup and instruction in the description of Issue #5553, along with minimal reproducible code.
The same setup can be used as Manual E2E test.
I am presentig the minimal reprouducible code and successful mitigation of the bug in the screenshot.

from google.adk import Workflow
from google.adk.events import RequestInput
from google.adk import Context
from google.adk.agents import Agent
from google.adk.workflow import node,FunctionNode
from pydantic import BaseModel

from typing import Any

class MySchema(BaseModel):
    greeting: str

analyzer = Agent(
    name="my_agent",
    model="gemini-3.1-flash-lite-preview",
    output_schema=MySchema,
    instruction="Your job is to greet the user. "
                "",
)
@node(rerun_on_resume=False)
async def get_user_approval(ctx: Context, node_input: Any):
    """Yields a RequestInput to pause the workflow and wait for user input."""
    yield RequestInput(message=f'please approve or reject.',response_schema = str)

@node(rerun_on_resume=True)
async def handle_process(ctx: Context, node_input: Any):
    """The orchestrator calling the interactive step."""
    user_response = await ctx.run_node(get_user_approval,node_input)
    if user_response.lower() == "yes":
        yield 'approved'
    yield "Denied"
    return

@node(rerun_on_resume=True)
async def my_workflow(ctx:Context):
    agent_input = "my agent input"
    agent_output = await ctx.run_node(analyzer,agent_input)
    agent_output = MySchema.validate(agent_output)
    result = await ctx.run_node(handle_process)
    yield result

root_agent = Workflow(
    name="greet_user",
    edges=[("START", my_workflow)]
)

Output

image

Additional Context

-I have tested it using the the static Graph-Based workflows. The behavior remains same, as the raw content is assigned to ctx.output during the rehydration.
-Also as specified in the comments in the _consume_event_queue function in runner.py

 async def _consume_event_queue(
      self, ic: InvocationContext, done_sentinel: object
  ) -> AsyncGenerator[Event, None]:
    """Consume events from ic._event_queue until done_sentinel."""
    while True:
      event_or_done, processed_signal = await ic._event_queue.get()
      if event_or_done is done_sentinel:
        break
      event: Event = event_or_done
      # When an LlmAgent node uses ``message_as_output`` (no
      # ``output_schema``), the wrapper sets both ``event.content``
      # (the model's text) AND ``event.output`` (the same text) to
      # signal that the message IS the node's output.  Clear
      # ``event.output`` on a copy here so downstream renderers don't
      # surface the same text twice.  Task-mode agents set
      # ``event.output`` from the ``finish_task`` FC args without
      # ``message_as_output``, so this clearing doesn't affect them.
      if not event.partial:
        if event.node_info.message_as_output and event.content is not None:
          event = event.model_copy()
          event.output = None

It feels like it intend to set message_as_output = True only when the output_schema is not provided, but process_llm_agent_output function in llm_agent_wrapper file, sets the message_as_output =True regardless of whether the output_schema is provided or not.

Checklist

  • I have read the CONTRIBUTING.md document.
  • I have performed a self-review of my own code.
  • I have commented my code, particularly in hard-to-understand areas.
  • I have added tests that prove my fix is effective or that my feature works.
  • New and existing unit tests pass locally with my changes.
  • I have manually tested my changes end-to-end.
  • Any dependent changes have been merged and published in downstream modules.

…tract output from event.content object before assigning it to child.output, in _reconstruct_node_states
@adk-bot adk-bot added the core [Component] This issue is related to the core interface and implementation label May 30, 2026
@samarth1224 samarth1224 changed the title add _process_content_object function in _rehydration_utils file to extract output from event.content object before assigning it to child.output, in _reconstruct_node_states Fix/(workflow)add _process_content_object function in _rehydration_utils file to extract output from event.content object before assigning it to child.output, in _reconstruct_node_states May 30, 2026
@samarth1224 samarth1224 changed the title Fix/(workflow)add _process_content_object function in _rehydration_utils file to extract output from event.content object before assigning it to child.output, in _reconstruct_node_states Fix(workflow)/add _process_content_object function in _rehydration_utils file to extract output from event.content object before assigning it to child.output, in _reconstruct_node_states May 30, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

core [Component] This issue is related to the core interface and implementation

Projects

None yet

2 participants