From d3428d7f381f8e52c0b67d0795fe41e7ee37bf5d Mon Sep 17 00:00:00 2001
From: Yufeng He <40085740+he-yufeng@users.noreply.github.com>
Date: Sun, 7 Jun 2026 11:07:49 +0800
Subject: [PATCH] fix: collect eval state from workflow nodes

---
 src/google/adk/cli/utils/state.py       | 23 ++++++++----
 tests/unittests/cli/utils/test_state.py | 48 +++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 6 deletions(-)
 create mode 100644 tests/unittests/cli/utils/test_state.py

diff --git a/src/google/adk/cli/utils/state.py b/src/google/adk/cli/utils/state.py
index 36d7b7b4c0..7c28613679 100644
--- a/src/google/adk/cli/utils/state.py
+++ b/src/google/adk/cli/utils/state.py
@@ -18,13 +18,24 @@
 from typing import Any
 from typing import Optional
 
-from ...agents.base_agent import BaseAgent
 from ...agents.llm_agent import LlmAgent
 
 
-def _create_empty_state(agent: BaseAgent, all_state: dict[str, Any]):
-  for sub_agent in agent.sub_agents:
-    _create_empty_state(sub_agent, all_state)
+def _create_empty_state(
+    agent: Any, all_state: dict[str, Any], visited: set[int]
+):
+  agent_id = id(agent)
+  if agent_id in visited:
+    return
+  visited.add(agent_id)
+
+  for sub_agent in getattr(agent, 'sub_agents', []) or []:
+    _create_empty_state(sub_agent, all_state, visited)
+
+  graph = getattr(agent, 'graph', None)
+  if graph is not None:
+    for graph_node in getattr(graph, 'nodes', []) or []:
+      _create_empty_state(graph_node, all_state, visited)
 
   if (
       isinstance(agent, LlmAgent)
@@ -36,11 +47,11 @@ def _create_empty_state(agent: BaseAgent, all_state: dict[str, Any]):
 
 
 def create_empty_state(
-    agent: BaseAgent, initialized_states: Optional[dict[str, Any]] = None
+    agent: Any, initialized_states: Optional[dict[str, Any]] = None
 ) -> dict[str, Any]:
   """Creates empty str for non-initialized states."""
   non_initialized_states = {}
-  _create_empty_state(agent, non_initialized_states)
+  _create_empty_state(agent, non_initialized_states, set())
   for key in initialized_states or {}:
     if key in non_initialized_states:
       del non_initialized_states[key]
diff --git a/tests/unittests/cli/utils/test_state.py b/tests/unittests/cli/utils/test_state.py
new file mode 100644
index 0000000000..48b6f63b17
--- /dev/null
+++ b/tests/unittests/cli/utils/test_state.py
@@ -0,0 +1,48 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from google.adk.agents.llm_agent import LlmAgent
+from google.adk.cli.utils.state import create_empty_state
+from google.adk.workflow import START
+from google.adk.workflow._workflow import Workflow
+
+
+def test_create_empty_state_reads_agent_tree():
+  child = LlmAgent(name='child', instruction='Use {child_key}')
+  root = LlmAgent(
+      name='root',
+      instruction='Use {root_key}',
+      sub_agents=[child],
+  )
+
+  assert create_empty_state(root) == {
+      'child_key': '',
+      'root_key': '',
+  }
+
+
+def test_create_empty_state_reads_workflow_graph_nodes():
+  node = LlmAgent(name='node', instruction='Use {workflow_key}')
+  workflow = Workflow(name='workflow', edges=[(START, node)])
+
+  assert create_empty_state(workflow) == {'workflow_key': ''}
+
+
+def test_create_empty_state_skips_initialized_workflow_state():
+  node = LlmAgent(name='node', instruction='Use {workflow_key} and {fresh_key}')
+  workflow = Workflow(name='workflow', edges=[(START, node)])
+
+  assert create_empty_state(workflow, {'workflow_key': 'set'}) == {
+      'fresh_key': ''
+  }