Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 17 additions & 6 deletions src/google/adk/cli/utils/state.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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]
Expand Down
48 changes: 48 additions & 0 deletions tests/unittests/cli/utils/test_state.py
Original file line number Diff line number Diff line change
@@ -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': ''
}