From e360241b8beaaf7112593aa2fe0cf35e2d3a4bfe Mon Sep 17 00:00:00 2001 From: Google Team Member Date: Wed, 1 Jul 2026 09:56:05 -0700 Subject: [PATCH 01/18] fix: prevent crash when parsing default values for typing.Any parameters PiperOrigin-RevId: 941154780 --- .../adk/tools/_function_parameter_parse_util.py | 2 ++ .../tools/test_from_function_with_options.py | 16 ++++++++++++++++ 2 files changed, 18 insertions(+) diff --git a/src/google/adk/tools/_function_parameter_parse_util.py b/src/google/adk/tools/_function_parameter_parse_util.py index 99b6e1991f6..758d8a3893b 100644 --- a/src/google/adk/tools/_function_parameter_parse_util.py +++ b/src/google/adk/tools/_function_parameter_parse_util.py @@ -222,6 +222,8 @@ def _is_default_value_compatible( default_value: Any, annotation: inspect.Parameter.annotation ) -> bool: # None type is expected to be handled external to this function + if annotation is Any: + return True if _is_builtin_primitive_or_compound(annotation): return isinstance(default_value, annotation) diff --git a/tests/unittests/tools/test_from_function_with_options.py b/tests/unittests/tools/test_from_function_with_options.py index d5e1ed54ae5..11e0764e71c 100644 --- a/tests/unittests/tools/test_from_function_with_options.py +++ b/tests/unittests/tools/test_from_function_with_options.py @@ -519,3 +519,19 @@ def generate_image( assert array_schema.items.items.any_of[0].type == types.Type.STRING assert array_schema.items.items.any_of[0].format is None assert array_schema.items.items.any_of[1].type == types.Type.STRING + + +def test_from_function_with_options_any_type_with_default_value(): + """Test that typing.Any with a default value works and doesn't crash.""" + + def my_tool(param: Any = 'default_string') -> str: + return f'ok {param}' + + declaration = _automatic_function_calling_util.from_function_with_options( + my_tool, GoogleLLMVariant.GEMINI_API + ) + + assert declaration.parameters is not None + assert declaration.parameters.properties['param'].default == 'default_string' + # Any type maps to None (no type) in schema + assert declaration.parameters.properties['param'].type is None From d60ac69e044b105be7b492a1aa197b9391f44e88 Mon Sep 17 00:00:00 2001 From: Vivek JM <24496671+vivekjm@users.noreply.github.com> Date: Wed, 1 Jul 2026 10:01:22 -0700 Subject: [PATCH 02/18] chore: Add coverage for pipe union list annotations Merge https://github.com/google/adk-python/pull/6160 Add regression coverage for optional list parameters written with Python pipe union syntax, covering the List[str] | None = None signature shape. Refs #3591. PiperOrigin-RevId: 941157723 --- ...t_function_tool_with_import_annotations.py | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/tests/unittests/tools/test_function_tool_with_import_annotations.py b/tests/unittests/tools/test_function_tool_with_import_annotations.py index ebd21bce95b..0d171628d33 100644 --- a/tests/unittests/tools/test_function_tool_with_import_annotations.py +++ b/tests/unittests/tools/test_function_tool_with_import_annotations.py @@ -159,6 +159,26 @@ def test_function(str_param: str, int_param: int, any_param: Any) -> str: assert declaration.response.type == types.Type.STRING +def test_pipe_union_list_annotation_parameter_vertex(): + """Test function with pipe union list parameter annotation.""" + + def test_function(file_patterns: list[str] | None = None) -> None: + """A test function that accepts optional file patterns.""" + pass + + declaration = _automatic_function_calling_util.from_function_with_options( + test_function, GoogleLLMVariant.VERTEX_AI + ) + + assert declaration.name == 'test_function' + assert declaration.parameters.type == 'OBJECT' + file_patterns_schema = declaration.parameters.properties['file_patterns'] + assert file_patterns_schema.type == types.Type.ARRAY + assert file_patterns_schema.items.type == types.Type.STRING + assert file_patterns_schema.nullable + assert declaration.parameters.required == [] + + def test_string_annotation_no_params_vertex(): """Test function with no parameters but string annotation return.""" From d69deddb1302775f812ed044ac6ba8a4a0685b14 Mon Sep 17 00:00:00 2001 From: Mukunda Rao Katta Date: Wed, 1 Jul 2026 10:23:07 -0700 Subject: [PATCH 03/18] fix(tools): accept dict output_schema in SetModelResponseTool (#5469) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merge https://github.com/google/adk-python/pull/5516 ## Summary Fixes #5469. When `output_schema` is a raw dict (e.g. `{"type": "object", "properties": {...}}`), `SetModelResponseTool.__init__` previously fell through to the generic `else` branch and used the dict **instance** as the parameter annotation. Downstream, `_function_parameter_parse_util._is_builtin_primitive_or_compound` does `annotation in _py_builtin_type_to_schema_type.keys()`, which calls `__hash__` on the annotation and raises `TypeError: unhashable type: 'dict'`. This change adds an explicit `elif isinstance(output_schema, dict)` branch that uses the `dict` type (hashable) as the annotation, so the existing builtin lookup maps it to `Type.OBJECT` cleanly. `run_async` already handles this case via the existing `args.get('response')` pass-through. ## Reproduction (before fix) ```python from google.adk.agents import Agent agent = Agent( name="test", model="gemini-2.5-flash", instruction="You are a helpful agent.", output_schema={"type": "object", "properties": {"result": {"type": "string"}}}, ) # -> TypeError: unhashable type: 'dict' ``` ## Testing plan - Added regression unit tests in `tests/unittests/tools/test_set_model_response_tool.py`: - `test_tool_initialization_raw_dict_schema` — `__init__` with a raw dict does not crash and stores the schema. - `test_function_signature_generation_raw_dict_schema` — generated signature has a single `response: dict` parameter (annotation is the `dict` type, not the dict instance). - `test_get_declaration_raw_dict_schema` — `_get_declaration()` returns a valid declaration without raising `TypeError`. - `test_run_async_raw_dict_schema` — `run_async` returns the response unchanged. - Run: `pytest tests/unittests/tools/test_set_model_response_tool.py -q` - Existing tests for `BaseModel`, `list[BaseModel]`, `list[str]`, `dict[str, int]` schemas remain unchanged and still pass. ## Notes - Schema fidelity (propagating dict-schema constraints into the function declaration) is intentionally out of scope; this PR only fixes the crash, matching the issue's reported scope and the existing handling for `list[str]` / `dict[str, int]`. Co-authored-by: Xuan Yang COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/5516 from MukundaKatta:fix/set-model-response-tool-dict-schema 5e808a978206af2385f093b91de4a915765c3293 PiperOrigin-RevId: 941169611 --- .../adk/tools/set_model_response_tool.py | 20 +++- .../tools/test_set_model_response_tool.py | 100 ++++++++++++++++++ 2 files changed, 119 insertions(+), 1 deletion(-) diff --git a/src/google/adk/tools/set_model_response_tool.py b/src/google/adk/tools/set_model_response_tool.py index 592ec195829..3e80ba64f76 100644 --- a/src/google/adk/tools/set_model_response_tool.py +++ b/src/google/adk/tools/set_model_response_tool.py @@ -52,6 +52,10 @@ def __init__(self, output_schema: SchemaType): - dict: Raw dict schemas - Schema: Google's Schema type """ + # Convert types.Schema instance to raw dict to avoid unhashable crash + if isinstance(output_schema, types.Schema): + output_schema = output_schema.model_dump(exclude_none=True) + self.output_schema = output_schema self._is_basemodel = is_basemodel_schema(output_schema) self._is_list_of_basemodel = is_list_of_basemodel(output_schema) @@ -87,8 +91,22 @@ def set_model_response() -> str: annotation=list[inner_type], ) ] + elif isinstance(output_schema, dict): + # For raw dict schemas (e.g. {"type": "object", "properties": {...}}), + # use the `dict` type itself as the annotation rather than the dict + # instance. Passing the instance would later trigger + # `annotation in _py_builtin_type_to_schema_type.keys()` inside + # `_function_parameter_parse_util`, which calls `__hash__` on the + # annotation and raises `TypeError: unhashable type: 'dict'`. + params = [ + inspect.Parameter( + 'response', + inspect.Parameter.KEYWORD_ONLY, + annotation=dict, + ) + ] else: - # For other schema types (list[str], dict, etc.), + # For other schema types (list[str], dict[str, int], etc.), # create a single parameter with the actual schema type params = [ inspect.Parameter( diff --git a/tests/unittests/tools/test_set_model_response_tool.py b/tests/unittests/tools/test_set_model_response_tool.py index c4374495735..54ff459eeae 100644 --- a/tests/unittests/tools/test_set_model_response_tool.py +++ b/tests/unittests/tools/test_set_model_response_tool.py @@ -25,6 +25,7 @@ from google.adk.sessions.in_memory_session_service import InMemorySessionService from google.adk.tools.set_model_response_tool import SetModelResponseTool from google.adk.tools.tool_context import ToolContext +from google.genai import types from pydantic import BaseModel from pydantic import Field from pydantic import ValidationError @@ -472,6 +473,105 @@ async def test_run_async_dict_schema(): assert result == {'a': 1, 'b': 2, 'c': 3} +def test_tool_initialization_raw_dict_schema(): + """Raw dict output_schema must not crash and must be stored as-is.""" + raw_schema = { + 'type': 'object', + 'properties': {'result': {'type': 'string'}}, + } + + tool = SetModelResponseTool(raw_schema) + + assert tool.output_schema == raw_schema + assert not tool._is_basemodel + assert not tool._is_list_of_basemodel + assert tool.name == 'set_model_response' + assert tool.func is not None + + +def test_function_signature_generation_raw_dict_schema(): + """Raw dict schemas should produce a single `response: dict` parameter. + + The annotation must be the `dict` type (hashable), not the dict instance, + so downstream `_is_builtin_primitive_or_compound` does not raise + `TypeError: unhashable type: 'dict'`. + """ + raw_schema = { + 'type': 'object', + 'properties': {'result': {'type': 'string'}}, + } + + tool = SetModelResponseTool(raw_schema) + + sig = inspect.signature(tool.func) + + assert 'response' in sig.parameters + assert len(sig.parameters) == 1 + assert sig.parameters['response'].kind == inspect.Parameter.KEYWORD_ONLY + # The annotation is the hashable `dict` type, not the dict instance. + assert sig.parameters['response'].annotation is dict + + +def test_get_declaration_raw_dict_schema(): + """`_get_declaration` must not raise when given a raw dict schema.""" + raw_schema = { + 'type': 'object', + 'properties': {'result': {'type': 'string'}}, + } + + tool = SetModelResponseTool(raw_schema) + + declaration = tool._get_declaration() + + assert declaration is not None + assert declaration.name == 'set_model_response' + assert declaration.description is not None + + +@pytest.mark.asyncio +async def test_run_async_raw_dict_schema(): + """Tool execution with a raw dict schema returns the response unchanged.""" + raw_schema = { + 'type': 'object', + 'properties': {'result': {'type': 'string'}}, + } + tool = SetModelResponseTool(raw_schema) + + agent = LlmAgent(name='test_agent', model='gemini-1.5-flash') + invocation_context = await _create_invocation_context(agent) + tool_context = ToolContext(invocation_context) + + result = await tool.run_async( + args={'response': {'result': 'hello'}}, + tool_context=tool_context, + ) + + assert result == {'result': 'hello'} + + +def test_tool_initialization_schema_instance(): + """types.Schema instance output_schema must be converted to dict and not crash.""" + schema_instance = types.Schema( + type=types.Type.OBJECT, + properties={'result': types.Schema(type=types.Type.STRING)}, + ) + + tool = SetModelResponseTool(schema_instance) + + # Check that it converted it to a dictionary + assert isinstance(tool.output_schema, dict) + assert 'result' in tool.output_schema['properties'] + + sig = inspect.signature(tool.func) + assert 'response' in sig.parameters + assert sig.parameters['response'].annotation is dict + + # Check that get_declaration works and doesn't crash with TypeError + declaration = tool._get_declaration() + assert declaration is not None + assert declaration.name == 'set_model_response' + + class SubSchema(BaseModel): field1: str = Field(description='Field 1') From 451c8cc6da037a59fcb58877a15583ad537bfd58 Mon Sep 17 00:00:00 2001 From: Shangjie Chen Date: Wed, 1 Jul 2026 10:35:25 -0700 Subject: [PATCH 04/18] refactor: Split graph validation and parsing from _graph.py Move graph validation and parsing logic from `_graph.py` to dedicated utility modules (`_graph_validation.py` and `_graph_parser.py`) under `workflow/utils/`. This simplifies the core `Graph` and `Edge` model definitions. Co-authored-by: Shangjie Chen PiperOrigin-RevId: 941176423 --- src/google/adk/workflow/_graph.py | 395 +--------- .../adk/workflow/utils/_graph_parser.py | 210 +++++ .../adk/workflow/utils/_graph_validation.py | 211 +++++ tests/unittests/workflow/test_graph.py | 724 +----------------- .../workflow/utils/test_graph_parser.py | 381 +++++++++ .../workflow/utils/test_graph_validation.py | 373 +++++++++ 6 files changed, 1182 insertions(+), 1112 deletions(-) create mode 100644 src/google/adk/workflow/utils/_graph_parser.py create mode 100644 src/google/adk/workflow/utils/_graph_validation.py create mode 100644 tests/unittests/workflow/utils/test_graph_parser.py create mode 100644 tests/unittests/workflow/utils/test_graph_validation.py diff --git a/src/google/adk/workflow/_graph.py b/src/google/adk/workflow/_graph.py index 641e3bda254..56c5459012b 100644 --- a/src/google/adk/workflow/_graph.py +++ b/src/google/adk/workflow/_graph.py @@ -16,15 +16,12 @@ from __future__ import annotations -from collections import Counter from collections.abc import Callable import logging logger = logging.getLogger("google_adk." + __name__) -from collections.abc import Set from typing import Annotated from typing import Any -from typing import get_args from typing import Literal from typing import TypeAlias @@ -36,7 +33,6 @@ from ..tools.base_tool import BaseTool from ._base_node import BaseNode -from ._base_node import START RouteValue: TypeAlias = bool | int | str """Type alias for valid routing values used in conditional graph edges.""" @@ -91,205 +87,9 @@ class Edge(BaseModel): Can be an explicit Edge object, or a tuple representing a chain of elements. """ -from .utils._workflow_graph_utils import build_node -from .utils._workflow_graph_utils import is_node_like - DEFAULT_ROUTE = "__DEFAULT__" - -def _expand_routing_map( - from_element: ChainElement, - routing_map: RoutingMap, -) -> list[tuple[ChainElement, NodeLike | tuple[NodeLike, ...], RouteValue]]: - """Expands a routing map into individual (from, to, route) triples. - - Args: - from_element: The source node(s). Can be a single NodeLike or a - tuple of NodeLike for fan-in. - routing_map: A dict mapping route values to destination nodes. - Values can be a single NodeLike or a tuple of NodeLike for - fan-out. - - Returns: - A list of (from_element, target, route) triples where target can - be a single NodeLike or a tuple for fan-out. - - Raises: - ValueError: If the routing map is empty, has non-RouteValue keys, - or has non-NodeLike values. - """ - if not routing_map: - raise ValueError( - "Routing map must not be empty. Provide at least one route -> node" - " mapping." - ) - - route_value_types = get_args(RouteValue) - expanded: list[ - tuple[ChainElement, NodeLike | tuple[NodeLike, ...], RouteValue] - ] = [] - - for route_key, target in routing_map.items(): - if not isinstance(route_key, route_value_types): - raise ValueError( - f"Invalid routing map key: {route_key!r} (type" - f" {type(route_key).__name__}). Keys must be RouteValue" - " (str, int, or bool)." - ) - if isinstance(target, tuple): - for node in target: - if not is_node_like(node): - raise ValueError( - f"Invalid node in fan-out tuple for route {route_key!r}:" - f" {node!r} (type {type(node).__name__})." - " Values must be NodeLike (BaseNode, BaseAgent, BaseTool," - " callable, or 'START')." - ) - elif not is_node_like(target): - raise ValueError( - f"Invalid routing map value for route {route_key!r}:" - f" {target!r} (type {type(target).__name__})." - " Values must be NodeLike (BaseNode, BaseAgent, BaseTool," - " callable, or 'START')." - ) - expanded.append((from_element, target, route_key)) - - return expanded - - -def _nodes_from_routing_map( - routing_map: RoutingMap, -) -> list[NodeLike]: - """Extracts all target nodes from a routing map, flattening fan-out tuples. - - Args: - routing_map: A dict mapping route values to destination nodes. - - Returns: - A flat list of all NodeLike targets referenced in the map. - """ - nodes: list[NodeLike] = [] - for target in routing_map.values(): - if isinstance(target, tuple): - nodes.extend(target) - else: - nodes.append(target) - return nodes - - -def _flatten_element( - element: NodeLike | tuple[NodeLike, ...] | RoutingMap, -) -> list[NodeLike]: - """Flattens a chain element into a list of individual nodes. - - - A single NodeLike is wrapped in a list. - - A tuple of NodeLike is converted to a list. - - A RoutingMap (dict) has its target nodes extracted and flattened. - """ - if isinstance(element, dict): - return _nodes_from_routing_map(element) - if isinstance(element, tuple): - return list(element) - return [element] - - -def _get_or_build_node( - node_like: NodeLike, node_map: dict[int, BaseNode] -) -> BaseNode: - """Gets a node from the map or builds it if not present.""" - if node_like == "START": - return START - - node_id = id(node_like) - if node_id in node_map: - return node_map[node_id] - - if isinstance(node_like, BaseNode): - wrapped = build_node(node_like) - if wrapped is not node_like: - node_map[node_id] = wrapped - return wrapped - return node_like - - node = build_node(node_like) - node_map[node_id] = node - return node - - -def _process_explicit_edge( - edge: Edge, node_map: dict[int, BaseNode], graph_edges: list[Edge] -) -> None: - """Processes an explicit Edge object.""" - graph_edges.append( - Edge( - from_node=_get_or_build_node(edge.from_node, node_map), - to_node=_get_or_build_node(edge.to_node, node_map), - route=edge.route, - ) - ) - - -def _process_chain( - chain: tuple[Any, ...], - node_map: dict[int, BaseNode], - graph_edges: list[Edge], -) -> None: - """Processes a chain of elements (tuple).""" - for i in range(len(chain) - 1): - from_el = chain[i] - to_el = chain[i + 1] - - if isinstance(to_el, dict): - _process_routing_map_edge(from_el, to_el, node_map, graph_edges) - else: - _process_unconditional_edge(from_el, to_el, node_map, graph_edges) - - -def _process_routing_map_edge( - from_el: Any, - to_el: RoutingMap, - node_map: dict[int, BaseNode], - graph_edges: list[Edge], -) -> None: - """Processes edges where the destination is a routing map.""" - if isinstance(from_el, dict): - raise ValueError( - "Consecutive routing maps are not allowed in a chain." - " Split them into separate edge items." - ) - - # A routing map (dict) in a chain behaves like a fan-out tuple - # but with conditioned incoming edges. - for exp_from, exp_to, route in _expand_routing_map(from_el, to_el): - for from_node in _flatten_element(exp_from): - for to_node in _flatten_element(exp_to): - graph_edges.append( - Edge( - from_node=_get_or_build_node(from_node, node_map), - to_node=_get_or_build_node(to_node, node_map), - route=route, - ) - ) - - -def _process_unconditional_edge( - from_el: Any, - to_el: Any, - node_map: dict[int, BaseNode], - graph_edges: list[Edge], -) -> None: - """Processes unconditional edges between elements.""" - # _flatten_element handles dicts (fan-in from routing map values) - # and tuples (fan-in/out). - for from_node in _flatten_element(from_el): - for to_node in _flatten_element(to_el): - graph_edges.append( - Edge( - from_node=_get_or_build_node(from_node, node_map), - to_node=_get_or_build_node(to_node, node_map), - route=None, - ) - ) +# --- Graph --- class Graph(BaseModel): @@ -311,18 +111,9 @@ class Graph(BaseModel): @classmethod def from_edge_items(cls, edge_items: list[EdgeItem]) -> Graph: """Creates a Graph from a list of edge items.""" - node_map: dict[int, BaseNode] = {} - graph_edges: list[Edge] = [] - - for item in edge_items: - if isinstance(item, Edge): - _process_explicit_edge(item, node_map, graph_edges) - elif isinstance(item, tuple): - _process_chain(item, node_map, graph_edges) - else: - raise ValueError(f"Invalid edge type: {type(item)}") + from .utils._graph_parser import parse_edge_items - return Graph(edges=graph_edges) + return Graph(edges=parse_edge_items(edge_items)) def model_post_init(self, context: Any) -> None: """Populates nodes from edges.""" @@ -391,182 +182,8 @@ def get_next_pending_nodes( return next_pending_nodes - def _detect_unconditional_cycles(self, node_names: Set[str]) -> None: - """Detects unconditional cycles in the graph. - - Edges with route=None are always followed, so a cycle consisting - entirely of such edges would loop forever. Conditional edges - (with a route) are allowed to form cycles (loop patterns). - """ - unconditional_adj: dict[str, list[str]] = {name: [] for name in node_names} - for edge in self.edges: - if edge.route is None: - unconditional_adj[edge.from_node.name].append(edge.to_node.name) - - in_stack: set[str] = set() - done: set[str] = set() - - def _dfs(node: str, path: list[str]) -> None: - in_stack.add(node) - path.append(node) - for neighbor in unconditional_adj[node]: - if neighbor in in_stack: - cycle_start = path.index(neighbor) - cycle = path[cycle_start:] + [neighbor] - raise ValueError( - "Graph validation failed. Unconditional cycle detected:" - f" {' -> '.join(cycle)}. Cycles must include at" - " least one conditional (routed) edge to avoid" - " infinite loops." - ) - if neighbor not in done: - _dfs(neighbor, path) - path.pop() - in_stack.remove(node) - done.add(node) - - for name in node_names: - if name not in done: - _dfs(name, []) - - def _validate_duplicate_node_names(self) -> set[str]: - """Checks for duplicate node names.""" - names = [node.name for node in self.nodes] - duplicates = sorted( - name for name, count in Counter(names).items() if count > 1 - ) - - if duplicates: - raise ValueError( - "Graph validation failed. Duplicate node names found:" - f" {duplicates}. This means multiple distinct node objects" - " have the same name. If you intended to reuse the same node, ensure" - " you pass the exact same object instance. If you intended to have" - " distinct nodes, ensure they have unique names." - ) - return set(names) - - def _validate_start_node(self, node_names: set[str]) -> None: - """Checks for existence of START node.""" - if START.name not in node_names: - raise ValueError( - "Graph validation failed. START node (name: " - f"'{START.name}') not found in graph nodes." - ) - - def _validate_connectivity(self, node_names: set[str]) -> None: - """Checks connectivity and reachability from START.""" - to_nodes: set[str] = set() - adj: dict[str, set[str]] = {name: set() for name in node_names} - for edge in self.edges: - adj[edge.from_node.name].add(edge.to_node.name) - to_nodes.add(edge.to_node.name) - - reachable: set[str] = set() - stack = [START.name] - while stack: - node = stack.pop() - if node in reachable: - continue - reachable.add(node) - stack.extend(adj[node] - reachable) - - unreachable_nodes = node_names - reachable - if unreachable_nodes: - raise ValueError( - "Graph validation failed. The following nodes are unreachable" - f" from START: {sorted(unreachable_nodes)}" - ) - if START.name in to_nodes: - raise ValueError( - "Graph validation failed. START node must not have incoming edges." - ) - - def _validate_duplicate_edges(self) -> None: - """Checks for duplicate edges.""" - seen_edges = set() - for edge in self.edges: - edge_tuple = (edge.from_node.name, edge.to_node.name) - if edge_tuple in seen_edges: - raise ValueError( - "Graph validation failed. Duplicate edge found: from=" - f"{edge.from_node.name}, to={edge.to_node.name}" - ) - seen_edges.add(edge_tuple) - - def _validate_default_routes(self) -> None: - """Checks constraints on DEFAULT_ROUTE.""" - default_route_edges: dict[str, str] = {} - for edge in self.edges: - if isinstance(edge.route, list) and DEFAULT_ROUTE in edge.route: - raise ValueError( - "Graph validation failed. DEFAULT_ROUTE cannot be combined" - " with other routes in a list (edge from=" - f"{edge.from_node.name}, to={edge.to_node.name})." - " Use a separate edge for DEFAULT_ROUTE." - ) - if edge.route == DEFAULT_ROUTE: - from_node_name = edge.from_node.name - if from_node_name in default_route_edges: - raise ValueError( - "Graph validation failed. Multiple DEFAULT_ROUTE edges found" - f" from node {from_node_name} to" - f" {default_route_edges[from_node_name]} and" - f" {edge.to_node.name}" - ) - default_route_edges[from_node_name] = edge.to_node.name - - def _validate_static_schemas(self) -> None: - """Validates static schemas on edges.""" - for edge in self.edges: - from_node = edge.from_node - to_node = edge.to_node - if from_node.output_schema and to_node.input_schema: - if from_node.output_schema != to_node.input_schema: - raise ValueError( - "Graph validation failed. Schema mismatch on edge" - f" {from_node.name} -> {to_node.name}." - f" Output schema {from_node.output_schema} does not match" - f" input schema {to_node.input_schema}." - ) - - def _validate_chat_agent_wiring(self) -> None: - """Validates that chat-mode agents do not have incoming edges from non-START nodes.""" - from ..agents.llm_agent import LlmAgent - - for edge in self.edges: - to_node = edge.to_node - if ( - isinstance(to_node, LlmAgent) - and getattr(to_node, "mode", None) == "chat" - ): - if edge.from_node.name != START.name: - raise ValueError( - f"The agent '{to_node.name}' has been added to the workflow with" - f" mode='chat' following node '{edge.from_node.name}'. This is" - " not supported because chat-mode agents rely on conversational" - " history (session events) and cannot consume direct node inputs" - " from preceding nodes. Please change the agent's mode to" - " 'single_turn'" - ) - - def _compute_terminal_nodes(self) -> None: - """Computes terminal nodes (no outgoing edges).""" - from_names = {edge.from_node.name for edge in self.edges} - self._terminal_node_names = { - n.name - for n in self.nodes - if n.name != START.name and n.name not in from_names - } - def validate_graph(self) -> None: """Validates the workflow graph.""" - node_names = self._validate_duplicate_node_names() - self._validate_start_node(node_names) - self._validate_connectivity(node_names) - self._validate_duplicate_edges() - self._validate_default_routes() - self._detect_unconditional_cycles(node_names) - self._validate_static_schemas() - self._validate_chat_agent_wiring() - self._compute_terminal_nodes() + from .utils._graph_validation import validate_graph + + self._terminal_node_names = validate_graph(self.nodes, self.edges) diff --git a/src/google/adk/workflow/utils/_graph_parser.py b/src/google/adk/workflow/utils/_graph_parser.py new file mode 100644 index 00000000000..45dddb66b8c --- /dev/null +++ b/src/google/adk/workflow/utils/_graph_parser.py @@ -0,0 +1,210 @@ +# 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. + +"""Utility functions for parsing workflow edges and chains.""" + +from __future__ import annotations + +from typing import Any +from typing import get_args + +from .._base_node import BaseNode +from .._base_node import START +from .._graph import ChainElement +from .._graph import Edge +from .._graph import EdgeItem +from .._graph import NodeLike +from .._graph import RouteValue +from .._graph import RoutingMap +from ._workflow_graph_utils import build_node +from ._workflow_graph_utils import is_node_like + + +def _expand_routing_map( + from_element: ChainElement, + routing_map: RoutingMap, +) -> list[tuple[ChainElement, NodeLike | tuple[NodeLike, ...], RouteValue]]: + """Expands a routing map into individual (from, to, route) triples.""" + if not routing_map: + raise ValueError( + "Routing map must not be empty. Provide at least one route -> node" + " mapping." + ) + + route_value_types = get_args(RouteValue) + expanded: list[ + tuple[ChainElement, NodeLike | tuple[NodeLike, ...], RouteValue] + ] = [] + + for route_key, target in routing_map.items(): + if not isinstance(route_key, route_value_types): + raise ValueError( + f"Invalid routing map key: {route_key!r} (type" + f" {type(route_key).__name__}). Keys must be RouteValue" + " (str, int, or bool)." + ) + if isinstance(target, tuple): + for node in target: + if not is_node_like(node): + raise ValueError( + f"Invalid node in fan-out tuple for route {route_key!r}:" + f" {node!r} (type {type(node).__name__})." + " Values must be NodeLike (BaseNode, BaseAgent, BaseTool," + " callable, or 'START')." + ) + elif not is_node_like(target): + raise ValueError( + f"Invalid routing map value for route {route_key!r}:" + f" {target!r} (type {type(target).__name__})." + " Values must be NodeLike (BaseNode, BaseAgent, BaseTool," + " callable, or 'START')." + ) + expanded.append((from_element, target, route_key)) + + return expanded + + +def _nodes_from_routing_map( + routing_map: RoutingMap, +) -> list[NodeLike]: + """Extracts all target nodes from a routing map, flattening fan-out tuples.""" + nodes: list[NodeLike] = [] + for target in routing_map.values(): + if isinstance(target, tuple): + nodes.extend(target) + else: + nodes.append(target) + return nodes + + +def _flatten_element( + element: NodeLike | tuple[NodeLike, ...] | RoutingMap, +) -> list[NodeLike]: + """Flattens a chain element into a list of individual nodes.""" + if isinstance(element, dict): + return _nodes_from_routing_map(element) + if isinstance(element, tuple): + return list(element) + return [element] + + +def _get_or_build_node( + node_like: NodeLike, node_map: dict[int, BaseNode] +) -> BaseNode: + """Gets a node from the map or builds it if not present.""" + if node_like == "START": + return START + + node_id = id(node_like) + if node_id in node_map: + return node_map[node_id] + + if isinstance(node_like, BaseNode): + wrapped = build_node(node_like) + if wrapped is not node_like: + node_map[node_id] = wrapped + return wrapped + return node_like + + node = build_node(node_like) + node_map[node_id] = node + return node + + +def _process_explicit_edge( + edge: Edge, node_map: dict[int, BaseNode], graph_edges: list[Edge] +) -> None: + """Processes an explicit Edge object.""" + graph_edges.append( + Edge( + from_node=_get_or_build_node(edge.from_node, node_map), + to_node=_get_or_build_node(edge.to_node, node_map), + route=edge.route, + ) + ) + + +def _process_routing_map_edge( + from_el: Any, + to_el: RoutingMap, + node_map: dict[int, BaseNode], + graph_edges: list[Edge], +) -> None: + """Processes edges where the destination is a routing map.""" + if isinstance(from_el, dict): + raise ValueError( + "Consecutive routing maps are not allowed in a chain." + " Split them into separate edge items." + ) + + for exp_from, exp_to, route in _expand_routing_map(from_el, to_el): + for from_node in _flatten_element(exp_from): + for to_node in _flatten_element(exp_to): + graph_edges.append( + Edge( + from_node=_get_or_build_node(from_node, node_map), + to_node=_get_or_build_node(to_node, node_map), + route=route, + ) + ) + + +def _process_unconditional_edge( + from_el: Any, + to_el: Any, + node_map: dict[int, BaseNode], + graph_edges: list[Edge], +) -> None: + """Processes unconditional edges between elements.""" + for from_node in _flatten_element(from_el): + for to_node in _flatten_element(to_el): + graph_edges.append( + Edge( + from_node=_get_or_build_node(from_node, node_map), + to_node=_get_or_build_node(to_node, node_map), + route=None, + ) + ) + + +def _process_chain( + chain: tuple[Any, ...], + node_map: dict[int, BaseNode], + graph_edges: list[Edge], +) -> None: + """Processes a chain of elements (tuple).""" + for i in range(len(chain) - 1): + from_el = chain[i] + to_el = chain[i + 1] + + if isinstance(to_el, dict): + _process_routing_map_edge(from_el, to_el, node_map, graph_edges) + else: + _process_unconditional_edge(from_el, to_el, node_map, graph_edges) + + +def parse_edge_items(edge_items: list[EdgeItem]) -> list[Edge]: + """Parses a list of edge items into a flat list of Edge objects.""" + node_map: dict[int, BaseNode] = {} + graph_edges: list[Edge] = [] + + for item in edge_items: + if isinstance(item, Edge): + _process_explicit_edge(item, node_map, graph_edges) + elif isinstance(item, tuple): + _process_chain(item, node_map, graph_edges) + else: + raise ValueError(f"Invalid edge type: {type(item)}") + + return graph_edges diff --git a/src/google/adk/workflow/utils/_graph_validation.py b/src/google/adk/workflow/utils/_graph_validation.py new file mode 100644 index 00000000000..03962f07c4c --- /dev/null +++ b/src/google/adk/workflow/utils/_graph_validation.py @@ -0,0 +1,211 @@ +# 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. + +"""Utility functions for validating workflow graphs.""" + +from __future__ import annotations + +from collections import Counter + +from .._base_node import BaseNode +from .._base_node import START +from .._graph import DEFAULT_ROUTE +from .._graph import Edge + + +def _detect_unconditional_cycles( + edges: list[Edge], node_names: Set[str] +) -> None: + """Detects unconditional cycles in the graph.""" + unconditional_adj: dict[str, list[str]] = {name: [] for name in node_names} + for edge in edges: + if edge.route is None: + unconditional_adj[edge.from_node.name].append(edge.to_node.name) + + in_stack: set[str] = set() + done: set[str] = set() + + def _dfs(node: str, path: list[str]) -> None: + in_stack.add(node) + path.append(node) + for neighbor in unconditional_adj[node]: + if neighbor in in_stack: + cycle_start = path.index(neighbor) + cycle = path[cycle_start:] + [neighbor] + raise ValueError( + "Graph validation failed. Unconditional cycle detected:" + f" {' -> '.join(cycle)}. Cycles must include at" + " least one conditional (routed) edge to avoid" + " infinite loops." + ) + if neighbor not in done: + _dfs(neighbor, path) + path.pop() + in_stack.remove(node) + done.add(node) + + for name in node_names: + if name not in done: + _dfs(name, []) + + +def _validate_duplicate_node_names(nodes: list[BaseNode]) -> set[str]: + """Checks for duplicate node names.""" + names = [node.name for node in nodes] + duplicates = sorted( + name for name, count in Counter(names).items() if count > 1 + ) + + if duplicates: + raise ValueError( + "Graph validation failed. Duplicate node names found:" + f" {duplicates}. This means multiple distinct node objects" + " have the same name. If you intended to reuse the same node, ensure" + " you pass the exact same object instance. If you intended to have" + " distinct nodes, ensure they have unique names." + ) + return set(names) + + +def _validate_start_node(node_names: set[str]) -> None: + """Checks for existence of START node.""" + if START.name not in node_names: + raise ValueError( + "Graph validation failed. START node (name: " + f"'{START.name}') not found in graph nodes." + ) + + +def _validate_connectivity(edges: list[Edge], node_names: set[str]) -> None: + """Checks connectivity and reachability from START.""" + to_nodes: set[str] = set() + adj: dict[str, set[str]] = {name: set() for name in node_names} + for edge in edges: + adj[edge.from_node.name].add(edge.to_node.name) + to_nodes.add(edge.to_node.name) + + reachable: set[str] = set() + stack = [START.name] + while stack: + node = stack.pop() + if node in reachable: + continue + reachable.add(node) + stack.extend(adj[node] - reachable) + + unreachable_nodes = node_names - reachable + if unreachable_nodes: + raise ValueError( + "Graph validation failed. The following nodes are unreachable" + f" from START: {sorted(unreachable_nodes)}" + ) + if START.name in to_nodes: + raise ValueError( + "Graph validation failed. START node must not have incoming edges." + ) + + +def _validate_duplicate_edges(edges: list[Edge]) -> None: + """Checks for duplicate edges.""" + seen_edges = set() + for edge in edges: + edge_tuple = (edge.from_node.name, edge.to_node.name) + if edge_tuple in seen_edges: + raise ValueError( + "Graph validation failed. Duplicate edge found: from=" + f"{edge.from_node.name}, to={edge.to_node.name}" + ) + seen_edges.add(edge_tuple) + + +def _validate_default_routes(edges: list[Edge]) -> None: + """Checks constraints on DEFAULT_ROUTE.""" + default_route_edges: dict[str, str] = {} + for edge in edges: + if isinstance(edge.route, list) and DEFAULT_ROUTE in edge.route: + raise ValueError( + "Graph validation failed. DEFAULT_ROUTE cannot be combined" + " with other routes in a list (edge from=" + f"{edge.from_node.name}, to={edge.to_node.name})." + " Use a separate edge for DEFAULT_ROUTE." + ) + if edge.route == DEFAULT_ROUTE: + from_node_name = edge.from_node.name + if from_node_name in default_route_edges: + raise ValueError( + "Graph validation failed. Multiple DEFAULT_ROUTE edges found" + f" from node {from_node_name} to" + f" {default_route_edges[from_node_name]} and" + f" {edge.to_node.name}" + ) + default_route_edges[from_node_name] = edge.to_node.name + + +def _validate_static_schemas(edges: list[Edge]) -> None: + """Validates static schemas on edges.""" + for edge in edges: + from_node = edge.from_node + to_node = edge.to_node + if from_node.output_schema and to_node.input_schema: + if from_node.output_schema != to_node.input_schema: + raise ValueError( + "Graph validation failed. Schema mismatch on edge" + f" {from_node.name} -> {to_node.name}." + f" Output schema {from_node.output_schema} does not match" + f" input schema {to_node.input_schema}." + ) + + +def _validate_chat_agent_wiring(edges: list[Edge]) -> None: + """Validates that chat-mode agents do not have incoming edges from non-START nodes.""" + from ...agents.llm_agent import LlmAgent + + for edge in edges: + to_node = edge.to_node + if ( + isinstance(to_node, LlmAgent) + and getattr(to_node, "mode", None) == "chat" + ): + if edge.from_node.name != START.name: + raise ValueError( + f"The agent '{to_node.name}' has been added to the workflow with" + f" mode='chat' following node '{edge.from_node.name}'. This is" + " not supported because chat-mode agents rely on conversational" + " history (session events) and cannot consume direct node inputs" + " from preceding nodes. Please change the agent's mode to" + " 'single_turn'" + ) + + +def _compute_terminal_nodes( + nodes: list[BaseNode], edges: list[Edge] +) -> set[str]: + """Computes terminal nodes (no outgoing edges).""" + from_names = {edge.from_node.name for edge in edges} + return { + n.name for n in nodes if n.name != START.name and n.name not in from_names + } + + +def validate_graph(nodes: list[BaseNode], edges: list[Edge]) -> set[str]: + """Validates the workflow graph and returns terminal node names.""" + node_names = _validate_duplicate_node_names(nodes) + _validate_start_node(node_names) + _validate_connectivity(edges, node_names) + _validate_duplicate_edges(edges) + _validate_default_routes(edges) + _detect_unconditional_cycles(edges, node_names) + _validate_static_schemas(edges) + _validate_chat_agent_wiring(edges) + return _compute_terminal_nodes(nodes, edges) diff --git a/tests/unittests/workflow/test_graph.py b/tests/unittests/workflow/test_graph.py index 31b21aa5c40..9eb7355175d 100644 --- a/tests/unittests/workflow/test_graph.py +++ b/tests/unittests/workflow/test_graph.py @@ -12,17 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Tests for Graph validation.""" +"""Tests for Graph validation and routing.""" import logging from google.adk.workflow import Edge -from google.adk.workflow import FunctionNode from google.adk.workflow import START from google.adk.workflow._graph import DEFAULT_ROUTE from google.adk.workflow._graph import Graph -from pydantic import BaseModel -import pytest from .workflow_testing_utils import TestingNode @@ -38,709 +35,6 @@ def test_valid_graph() -> None: graph.validate_graph() # Should not raise -def test_missing_start_node() -> None: - """Tests that a graph missing the START node fails validation.""" - node_a = TestingNode(name='NodeA') - node_b = TestingNode(name='NodeB') - graph = Graph( - edges=[Edge(from_node=node_a, to_node=node_b)], - ) - with pytest.raises( - ValueError, - match=( - r"Graph validation failed\. START node \(name: '__START__'\) not" - r' found in graph nodes\.' - ), - ): - graph.validate_graph() - - -def test_unreachable_node() -> None: - """Tests that a graph with an unreachable node fails validation.""" - node_a = TestingNode(name='NodeA') - node_b = TestingNode(name='NodeB') # Unreachable - graph = Graph( - edges=[ - Edge(from_node=START, to_node=node_a), - Edge(from_node=node_b, to_node=node_a), - ], - ) - with pytest.raises( - ValueError, - match=( - r'Graph validation failed\. The following nodes are unreachable' - r" from START: \['NodeB'\]" - ), - ): - graph.validate_graph() - - -def test_disconnected_routed_subgraph_is_unreachable() -> None: - """Tests that a disconnected subgraph with routed edges fails validation. - - Even though B and C each appear as a to_node in some edge, neither is - reachable from START. The old "has incoming edge" heuristic would let - this pass; true reachability from START catches it. - """ - node_a = TestingNode(name='NodeA') - node_b = TestingNode(name='NodeB') - node_c = TestingNode(name='NodeC') - graph = Graph( - edges=[ - Edge(from_node=START, to_node=node_a), - Edge(from_node=node_b, to_node=node_c, route='x'), - Edge(from_node=node_c, to_node=node_b, route='y'), - ], - ) - with pytest.raises( - ValueError, - match=( - r'Graph validation failed\. The following nodes are unreachable' - r" from START: \['NodeB', 'NodeC'\]" - ), - ): - graph.validate_graph() - - -@pytest.mark.parametrize( - 'routes', - [ - (None, None), - ('route1', 'route1'), - ('route1', 'route2'), - ('route1', None), - ], -) -def test_duplicate_edges_fail_validation( - routes: tuple[str | None, str | None], -) -> None: - """Tests that duplicate edges fail validation, regardless of routes.""" - node_a = TestingNode(name='NodeA') - graph = Graph( - edges=[ - Edge( - from_node=START, - to_node=node_a, - route=routes[0], - ), - Edge( - from_node=START, - to_node=node_a, - route=routes[1], - ), - ], - ) - with pytest.raises( - ValueError, - match=( - r'Graph validation failed\. Duplicate edge found: from=__START__,' - r' to=NodeA' - ), - ): - graph.validate_graph() - - -def test_start_node_with_incoming_edge() -> None: - """Tests graph with incoming edge to START node fails validation.""" - node_a = TestingNode(name='NodeA') - graph = Graph( - edges=[ - Edge(from_node=node_a, to_node=START), - Edge(from_node=START, to_node=node_a), - ], - ) - with pytest.raises( - ValueError, - match=( - r'Graph validation failed\. START node must not have incoming edges\.' - ), - ): - graph.validate_graph() - - -def test_multiple_default_routes_fail_validation() -> None: - """Tests that multiple DEFAULT_ROUTE edges from a node fail validation.""" - node_a = TestingNode(name='NodeA') - node_b = TestingNode(name='NodeB') - node_c = TestingNode(name='NodeC') - graph = Graph( - edges=[ - Edge(from_node=START, to_node=node_a), - Edge(from_node=node_a, to_node=node_b, route=DEFAULT_ROUTE), - Edge(from_node=node_a, to_node=node_c, route=DEFAULT_ROUTE), - ], - ) - with pytest.raises( - ValueError, - match=( - r'Graph validation failed\. Multiple DEFAULT_ROUTE edges found from' - r' node NodeA to NodeB and NodeC' - ), - ): - graph.validate_graph() - - -def test_single_default_route_passes_validation() -> None: - """Tests that a single DEFAULT_ROUTE edge from a node passes validation.""" - node_a = TestingNode(name='NodeA') - node_b = TestingNode(name='NodeB') - node_c = TestingNode(name='NodeC') - graph = Graph( - edges=[ - Edge(from_node=START, to_node=node_a), - Edge(from_node=node_a, to_node=node_b, route=DEFAULT_ROUTE), - Edge(from_node=node_a, to_node=node_c, route='another_route'), - ], - ) - graph.validate_graph() # Should not raise - - -def test_duplicate_node_names_fail_validation() -> None: - """Tests that duplicate nodes raise error.""" - - node_a1 = TestingNode(name='NodeA') - node_a2 = TestingNode(name='NodeA') - graph = Graph( - edges=[ - Edge(from_node=START, to_node=node_a1), - Edge(from_node=node_a1, to_node=node_a2), - ], - ) - with pytest.raises( - ValueError, - match=( - r"Graph validation failed\. Duplicate node names found: \['NodeA'\]\." - r' This means multiple distinct node objects have the same name\. If' - r' you intended to reuse the same node, ensure you pass the exact' - r' same object instance\. If you intended to have distinct nodes,' - r' ensure they have unique names\.' - ), - ): - graph.validate_graph() - - -def test_from_edge_items_with_node_reuse_passes_validation() -> None: - """Tests that node reuse with from_edge_items passes validation. - - The same my_node_func instance is used in the graph multiple times, and - the workflow graph should recognize it as the same instance and not throw - an error during validation. - """ - - def my_node_func() -> None: - pass - - node_b = TestingNode(name='NodeB') - graph = Graph.from_edge_items([ - (START, my_node_func), - (my_node_func, node_b), - ]) - graph.validate_graph() # Should not raise duplicate name error - - node_names = {n.name for n in graph.nodes} - assert node_names == {'__START__', 'my_node_func', 'NodeB'} - assert len(graph.nodes) == 3 - # Check that my_node_func was wrapped and deduplicated. - func_node = next(n for n in graph.nodes if n.name == 'my_node_func') - assert isinstance(func_node, FunctionNode) - - -def test_unconditional_cycle_fails_validation() -> None: - """Tests that a cycle of unconditional edges (route=None) fails.""" - node_a = TestingNode(name='NodeA') - node_b = TestingNode(name='NodeB') - graph = Graph( - edges=[ - Edge(from_node=START, to_node=node_a), - Edge(from_node=node_a, to_node=node_b), - Edge(from_node=node_b, to_node=node_a), - ], - ) - with pytest.raises( - ValueError, - match=r'Graph validation failed\. Unconditional cycle detected:', - ): - graph.validate_graph() - - -def test_unconditional_self_loop_fails_validation() -> None: - """Tests that an unconditional self-loop (A -> A) fails.""" - node_a = TestingNode(name='NodeA') - graph = Graph( - edges=[ - Edge(from_node=START, to_node=node_a), - Edge(from_node=node_a, to_node=node_a), - ], - ) - with pytest.raises( - ValueError, - match=r'Graph validation failed\. Unconditional cycle detected:', - ): - graph.validate_graph() - - -def test_longer_unconditional_cycle_fails_validation() -> None: - """Tests that a longer unconditional cycle (A -> B -> C -> A) fails.""" - node_a = TestingNode(name='NodeA') - node_b = TestingNode(name='NodeB') - node_c = TestingNode(name='NodeC') - graph = Graph( - edges=[ - Edge(from_node=START, to_node=node_a), - Edge(from_node=node_a, to_node=node_b), - Edge(from_node=node_b, to_node=node_c), - Edge(from_node=node_c, to_node=node_a), - ], - ) - with pytest.raises( - ValueError, - match=r'Graph validation failed\. Unconditional cycle detected:', - ): - graph.validate_graph() - - -def test_conditional_cycle_passes_validation() -> None: - """Tests that a cycle with a routed edge (loop pattern) passes.""" - node_a = TestingNode(name='NodeA') - node_b = TestingNode(name='NodeB') - graph = Graph( - edges=[ - Edge(from_node=START, to_node=node_a), - Edge(from_node=node_a, to_node=node_b), - Edge(from_node=node_b, to_node=node_a, route='retry'), - ], - ) - graph.validate_graph() # Should not raise — routed back-edge - - -def test_conditional_self_loop_passes_validation() -> None: - """Tests that a self-loop with a route passes validation.""" - node_a = TestingNode(name='NodeA') - node_b = TestingNode(name='NodeB') - graph = Graph( - edges=[ - Edge(from_node=START, to_node=node_a), - Edge(from_node=node_a, to_node=node_a, route='continue'), - Edge(from_node=node_a, to_node=node_b, route='done'), - ], - ) - graph.validate_graph() # Should not raise — routed self-loop - - -def test_dag_with_diamond_passes_validation() -> None: - """Tests that a DAG with a diamond shape passes validation.""" - node_a = TestingNode(name='NodeA') - node_b = TestingNode(name='NodeB') - node_c = TestingNode(name='NodeC') - graph = Graph( - edges=[ - Edge(from_node=START, to_node=node_a), - Edge(from_node=START, to_node=node_b), - Edge(from_node=node_a, to_node=node_c), - Edge(from_node=node_b, to_node=node_c), - ], - ) - graph.validate_graph() # Should not raise - - -# --- Routing map tests --- - - -def test_routing_map_basic() -> None: - """Tests that a string-keyed routing map expands to correct edges.""" - node_a = TestingNode(name='NodeA') - node_b = TestingNode(name='NodeB') - node_c = TestingNode(name='NodeC') - graph = Graph.from_edge_items([ - (START, node_a), - (node_a, {'route_b': node_b, 'route_c': node_c}), - ]) - graph.validate_graph() - - assert len(graph.edges) == 3 # START->A, A->B(route_b), A->C(route_c) - - routed_edges = [e for e in graph.edges if e.route is not None] - assert len(routed_edges) == 2 - - routes_and_targets = {(e.route, e.to_node.name) for e in routed_edges} - assert routes_and_targets == {('route_b', 'NodeB'), ('route_c', 'NodeC')} - - for e in routed_edges: - assert e.from_node.name == 'NodeA' - - -def test_routing_map_int_keys() -> None: - """Tests that integer route keys work in routing maps.""" - node_a = TestingNode(name='NodeA') - node_b = TestingNode(name='NodeB') - node_c = TestingNode(name='NodeC') - graph = Graph.from_edge_items([ - (START, node_a), - (node_a, {1: node_b, 2: node_c}), - ]) - graph.validate_graph() - - routed_edges = [e for e in graph.edges if e.route is not None] - assert len(routed_edges) == 2 - routes = [e.route for e in routed_edges] - assert 1 in routes - assert 2 in routes - - -def test_routing_map_bool_keys() -> None: - """Tests that boolean route keys work in routing maps.""" - node_a = TestingNode(name='NodeA') - node_b = TestingNode(name='NodeB') - node_c = TestingNode(name='NodeC') - graph = Graph.from_edge_items([ - (START, node_a), - (node_a, {True: node_b, False: node_c}), - ]) - graph.validate_graph() - - routed_edges = [e for e in graph.edges if e.route is not None] - assert len(routed_edges) == 2 - routes = [e.route for e in routed_edges] - assert True in routes - assert False in routes - - -def test_routing_map_with_fan_in_source() -> None: - """Tests that fan-in on the source side works with routing maps.""" - node_a = TestingNode(name='NodeA') - node_b = TestingNode(name='NodeB') - node_c = TestingNode(name='NodeC') - node_d = TestingNode(name='NodeD') - graph = Graph.from_edge_items([ - (START, node_a), - (START, node_b), - ((node_a, node_b), {'route_x': node_c, 'route_y': node_d}), - ]) - graph.validate_graph() - - # 2 from START + 4 from fan-in (A->C, A->D, B->C, B->D) - assert len(graph.edges) == 6 - - fan_in_edges = [ - e for e in graph.edges if e.from_node.name in ('NodeA', 'NodeB') - ] - assert len(fan_in_edges) == 4 - - combos = {(e.from_node.name, e.to_node.name, e.route) for e in fan_in_edges} - assert combos == { - ('NodeA', 'NodeC', 'route_x'), - ('NodeA', 'NodeD', 'route_y'), - ('NodeB', 'NodeC', 'route_x'), - ('NodeB', 'NodeD', 'route_y'), - } - - -def test_routing_map_with_callable_target() -> None: - """Tests that callable values in routing maps get wrapped via build_node.""" - node_a = TestingNode(name='NodeA') - - def my_target_func() -> None: - pass - - graph = Graph.from_edge_items([ - (START, node_a), - (node_a, {'route_x': my_target_func}), - ]) - graph.validate_graph() - - target_edge = next(e for e in graph.edges if e.route == 'route_x') - assert isinstance(target_edge.to_node, FunctionNode) - assert target_edge.to_node.name == 'my_target_func' - - -def test_routing_map_node_reuse() -> None: - """Tests that the same callable used in a map and elsewhere is deduplicated.""" - - def my_func() -> None: - pass - - node_b = TestingNode(name='NodeB') - graph = Graph.from_edge_items([ - (START, my_func), - (my_func, {'route_x': node_b}), - ]) - graph.validate_graph() - - # my_func should be wrapped once and reused. - func_nodes = [n for n in graph.nodes if n.name == 'my_func'] - assert len(func_nodes) == 1 - assert isinstance(func_nodes[0], FunctionNode) - - -def test_routing_map_empty_dict_raises() -> None: - """Tests that an empty routing map raises ValueError.""" - node_a = TestingNode(name='NodeA') - with pytest.raises( - ValueError, - match=r'Routing map must not be empty', - ): - Graph.from_edge_items([ - (START, node_a), - (node_a, {}), - ]) - - -def test_routing_map_invalid_key_raises() -> None: - """Tests that a non-RouteValue key raises ValueError.""" - node_a = TestingNode(name='NodeA') - node_b = TestingNode(name='NodeB') - with pytest.raises( - ValueError, - match=r'Invalid routing map key', - ): - Graph.from_edge_items([ - (START, node_a), - (node_a, {1.5: node_b}), - ]) - - -def test_routing_map_invalid_value_raises() -> None: - """Tests that a non-NodeLike value raises ValueError.""" - node_a = TestingNode(name='NodeA') - with pytest.raises( - ValueError, - match=r'Invalid routing map value', - ): - Graph.from_edge_items([ - (START, node_a), - (node_a, {'route_x': 42}), - ]) - - -def test_routing_map_fan_out_target() -> None: - """Tests that a tuple value in a routing map creates fan-out edges.""" - node_a = TestingNode(name='NodeA') - node_b = TestingNode(name='NodeB') - node_c = TestingNode(name='NodeC') - graph = Graph.from_edge_items([ - (START, node_a), - (node_a, {'route_x': (node_b, node_c)}), - ]) - graph.validate_graph() - - # START->A, A->B(route_x), A->C(route_x) - assert len(graph.edges) == 3 - - routed_edges = [e for e in graph.edges if e.route is not None] - assert len(routed_edges) == 2 - - # Both fan-out edges share the same route and source. - for e in routed_edges: - assert e.from_node.name == 'NodeA' - assert e.route == 'route_x' - - targets = {e.to_node.name for e in routed_edges} - assert targets == {'NodeB', 'NodeC'} - - -def test_routing_map_fan_out_invalid_element_raises() -> None: - """Tests that a non-NodeLike element inside a fan-out tuple raises.""" - node_a = TestingNode(name='NodeA') - node_b = TestingNode(name='NodeB') - with pytest.raises( - ValueError, - match=r'Invalid node in fan-out tuple', - ): - Graph.from_edge_items([ - (START, node_a), - (node_a, {'route_x': (node_b, 42)}), - ]) - - -# --- Routing map as chain element tests --- - - -def test_routing_map_chain_ending_with_dict() -> None: - """Tests a chain ending with a routing map creates correct edges.""" - node_a = TestingNode(name='NodeA') - node_b = TestingNode(name='NodeB') - node_c = TestingNode(name='NodeC') - graph = Graph.from_edge_items([ - (START, node_a, {'r1': node_b, 'r2': node_c}), - ]) - graph.validate_graph() - - # START->A (None), A->B (r1), A->C (r2) - assert len(graph.edges) == 3 - - start_edge = next(e for e in graph.edges if e.from_node.name == '__START__') - assert start_edge.to_node.name == 'NodeA' - assert start_edge.route is None - - routed_edges = [e for e in graph.edges if e.route is not None] - assert len(routed_edges) == 2 - routes_and_targets = {(e.route, e.to_node.name) for e in routed_edges} - assert routes_and_targets == {('r1', 'NodeB'), ('r2', 'NodeC')} - for e in routed_edges: - assert e.from_node.name == 'NodeA' - - -def test_routing_map_mid_chain_with_fan_in() -> None: - """Tests routing map mid-chain with fan-in to the next element.""" - node_a = TestingNode(name='NodeA') - node_b = TestingNode(name='NodeB') - node_c = TestingNode(name='NodeC') - node_d = TestingNode(name='NodeD') - graph = Graph.from_edge_items([ - (START, node_a, {'r1': node_b, 'r2': node_c}, node_d), - ]) - graph.validate_graph() - - # START->A (None), A->B (r1), A->C (r2), B->D (None), C->D (None) - assert len(graph.edges) == 5 - - routed_edges = sorted( - [e for e in graph.edges if e.route is not None], - key=lambda e: e.to_node.name, - ) - assert len(routed_edges) == 2 - assert routed_edges[0].from_node.name == 'NodeA' - assert routed_edges[0].to_node.name == 'NodeB' - assert routed_edges[0].route == 'r1' - assert routed_edges[1].from_node.name == 'NodeA' - assert routed_edges[1].to_node.name == 'NodeC' - assert routed_edges[1].route == 'r2' - - fan_in_edges = sorted( - [e for e in graph.edges if e.to_node.name == 'NodeD'], - key=lambda e: e.from_node.name, - ) - assert len(fan_in_edges) == 2 - assert fan_in_edges[0].from_node.name == 'NodeB' - assert fan_in_edges[0].route is None - assert fan_in_edges[1].from_node.name == 'NodeC' - assert fan_in_edges[1].route is None - - -def test_routing_map_mid_chain_fan_out_values() -> None: - """Tests routing map with fan-out tuple values, followed by fan-in.""" - node_a = TestingNode(name='NodeA') - node_b = TestingNode(name='NodeB') - node_c = TestingNode(name='NodeC') - node_d = TestingNode(name='NodeD') - graph = Graph.from_edge_items([ - (START, node_a, {'r1': (node_b, node_c)}, node_d), - ]) - graph.validate_graph() - - # START->A (None), A->B (r1), A->C (r1), B->D (None), C->D (None) - assert len(graph.edges) == 5 - - routed_edges = [e for e in graph.edges if e.route is not None] - assert len(routed_edges) == 2 - for e in routed_edges: - assert e.from_node.name == 'NodeA' - assert e.route == 'r1' - - fan_in_edges = [e for e in graph.edges if e.to_node.name == 'NodeD'] - assert len(fan_in_edges) == 2 - fan_in_sources = {e.from_node.name for e in fan_in_edges} - assert fan_in_sources == {'NodeB', 'NodeC'} - - -def test_routing_map_consecutive_dicts_raises() -> None: - """Tests that consecutive routing maps in a chain are rejected.""" - node_a = TestingNode(name='NodeA') - node_b = TestingNode(name='NodeB') - node_c = TestingNode(name='NodeC') - node_d = TestingNode(name='NodeD') - with pytest.raises( - ValueError, match=r'Consecutive routing maps are not allowed' - ): - Graph.from_edge_items([ - (START, node_a, {'r1': node_b, 'r2': node_c}, {'r3': node_d}), - ]) - - -def test_routing_map_empty_dict_in_chain_raises() -> None: - """Tests that an empty routing map in a chain raises ValueError.""" - node_a = TestingNode(name='NodeA') - node_b = TestingNode(name='NodeB') - with pytest.raises(ValueError, match=r'Routing map must not be empty'): - Graph.from_edge_items([ - (START, node_a, {}, node_b), - ]) - - -def test_routing_map_invalid_key_in_chain_raises() -> None: - """Tests that invalid routing map keys in a chain raise ValueError.""" - node_a = TestingNode(name='NodeA') - node_b = TestingNode(name='NodeB') - with pytest.raises(ValueError, match=r'Invalid routing map key'): - Graph.from_edge_items([ - (START, node_a, {1.5: node_b}), - ]) - - -def test_routing_map_2_tuple_backward_compat() -> None: - """Ensures existing 2-tuple routing map syntax still works.""" - node_a = TestingNode(name='NodeA') - node_b = TestingNode(name='NodeB') - node_c = TestingNode(name='NodeC') - graph = Graph.from_edge_items([ - (START, node_a), - (node_a, {'r1': node_b, 'r2': node_c}), - ]) - graph.validate_graph() - assert len(graph.edges) == 3 - - -class ModelA(BaseModel): - x: int - - -class ModelB(BaseModel): - x: int - - -def test_schema_match_passes() -> None: - """Tests that edges with matching schemas pass validation.""" - node_a = TestingNode(name='NodeA', output_schema=ModelA) - node_b = TestingNode(name='NodeB', input_schema=ModelA) - graph = Graph( - edges=[ - Edge(from_node=START, to_node=node_a), - Edge(from_node=node_a, to_node=node_b), - ], - ) - graph.validate_graph() # Should not raise - - -def test_schema_mismatch_raises() -> None: - """Tests that edges with mismatching schemas fail validation.""" - node_a = TestingNode(name='NodeA', output_schema=ModelA) - node_b = TestingNode(name='NodeB', input_schema=ModelB) - graph = Graph( - edges=[ - Edge(from_node=START, to_node=node_a), - Edge(from_node=node_a, to_node=node_b), - ], - ) - with pytest.raises( - ValueError, - match=r'Graph validation failed\. Schema mismatch on edge', - ): - graph.validate_graph() - - -def test_schema_missing_passes() -> None: - """Tests that edges with missing schemas pass validation.""" - node_a = TestingNode(name='NodeA', output_schema=ModelA) - node_b = TestingNode(name='NodeB') # No input schema - graph = Graph( - edges=[ - Edge(from_node=START, to_node=node_a), - Edge(from_node=node_a, to_node=node_b), - ], - ) - graph.validate_graph() # Should not raise - - def test_get_next_pending_nodes() -> None: """Tests that get_next_pending_nodes returns correct nodes based on routes.""" node_a = TestingNode(name='NodeA') @@ -779,22 +73,6 @@ def test_get_next_pending_nodes() -> None: assert set(next_nodes) == {'NodeB', 'NodeC'} -def test_chat_agent_wiring_validation_only_runs_on_llm_agent() -> None: - """Tests that _validate_chat_agent_wiring checks non-LlmAgent nodes safely.""" - node_a = TestingNode(name='NodeA') - node_b = TestingNode(name='NodeB') - # Set mode='chat' on a non-LlmAgent node - object.__setattr__(node_b, 'mode', 'chat') - - graph = Graph( - edges=[ - Edge(from_node=START, to_node=node_a), - Edge(from_node=node_a, to_node=node_b), - ], - ) - graph.validate_graph() # Should not raise because node_b is a TestingNode, not LlmAgent - - def test_get_next_pending_nodes_unmatched_route_warning(caplog) -> None: """Tests that a warning is logged when a route is unmatched and there's no DEFAULT_ROUTE.""" node_a = TestingNode(name='NodeA') diff --git a/tests/unittests/workflow/utils/test_graph_parser.py b/tests/unittests/workflow/utils/test_graph_parser.py new file mode 100644 index 00000000000..82f5fc55396 --- /dev/null +++ b/tests/unittests/workflow/utils/test_graph_parser.py @@ -0,0 +1,381 @@ +# 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. + +"""Tests for Graph parser utility.""" + +from google.adk.workflow import Edge +from google.adk.workflow import FunctionNode +from google.adk.workflow import START +from google.adk.workflow._graph import DEFAULT_ROUTE +from google.adk.workflow.utils._graph_parser import parse_edge_items +import pytest + +from ..workflow_testing_utils import TestingNode + + +def test_parse_edge_items_with_node_reuse() -> None: + """Tests that node reuse during parsing works and deduplicates wrapped nodes.""" + + def my_node_func() -> None: + pass + + node_b = TestingNode(name='NodeB') + edges = parse_edge_items([ + (START, my_node_func), + (my_node_func, node_b), + ]) + + assert len(edges) == 2 + edge1, edge2 = edges + assert edge1.from_node == START + assert isinstance(edge1.to_node, FunctionNode) + assert edge1.to_node.name == 'my_node_func' + + assert isinstance(edge2.from_node, FunctionNode) + assert edge2.from_node.name == 'my_node_func' + assert edge2.to_node == node_b + + # Verify exact same object instance was reused + assert edge1.to_node is edge2.from_node + + +def test_routing_map_basic() -> None: + """Tests that a string-keyed routing map expands to correct edges.""" + node_a = TestingNode(name='NodeA') + node_b = TestingNode(name='NodeB') + node_c = TestingNode(name='NodeC') + edges = parse_edge_items([ + (START, node_a), + (node_a, {'route_b': node_b, 'route_c': node_c}), + ]) + + assert len(edges) == 3 # START->A, A->B(route_b), A->C(route_c) + + routed_edges = [e for e in edges if e.route is not None] + assert len(routed_edges) == 2 + + routes_and_targets = {(e.route, e.to_node.name) for e in routed_edges} + assert routes_and_targets == {('route_b', 'NodeB'), ('route_c', 'NodeC')} + + for e in routed_edges: + assert e.from_node.name == 'NodeA' + + +def test_routing_map_int_keys() -> None: + """Tests that integer route keys work in routing maps.""" + node_a = TestingNode(name='NodeA') + node_b = TestingNode(name='NodeB') + node_c = TestingNode(name='NodeC') + edges = parse_edge_items([ + (START, node_a), + (node_a, {1: node_b, 2: node_c}), + ]) + + routed_edges = [e for e in edges if e.route is not None] + assert len(routed_edges) == 2 + routes = [e.route for e in routed_edges] + assert 1 in routes + assert 2 in routes + + +def test_routing_map_bool_keys() -> None: + """Tests that boolean route keys work in routing maps.""" + node_a = TestingNode(name='NodeA') + node_b = TestingNode(name='NodeB') + node_c = TestingNode(name='NodeC') + edges = parse_edge_items([ + (START, node_a), + (node_a, {True: node_b, False: node_c}), + ]) + + routed_edges = [e for e in edges if e.route is not None] + assert len(routed_edges) == 2 + routes = [e.route for e in routed_edges] + assert True in routes + assert False in routes + + +def test_routing_map_with_fan_in_source() -> None: + """Tests that fan-in on the source side works with routing maps.""" + node_a = TestingNode(name='NodeA') + node_b = TestingNode(name='NodeB') + node_c = TestingNode(name='NodeC') + node_d = TestingNode(name='NodeD') + edges = parse_edge_items([ + (START, node_a), + (START, node_b), + ((node_a, node_b), {'route_x': node_c, 'route_y': node_d}), + ]) + + # 2 from START + 4 from fan-in (A->C, A->D, B->C, B->D) + assert len(edges) == 6 + + fan_in_edges = [e for e in edges if e.from_node.name in ('NodeA', 'NodeB')] + assert len(fan_in_edges) == 4 + + combos = {(e.from_node.name, e.to_node.name, e.route) for e in fan_in_edges} + assert combos == { + ('NodeA', 'NodeC', 'route_x'), + ('NodeA', 'NodeD', 'route_y'), + ('NodeB', 'NodeC', 'route_x'), + ('NodeB', 'NodeD', 'route_y'), + } + + +def test_routing_map_with_callable_target() -> None: + """Tests that callable values in routing maps get wrapped via build_node.""" + node_a = TestingNode(name='NodeA') + + def my_target_func() -> None: + pass + + edges = parse_edge_items([ + (START, node_a), + (node_a, {'route_x': my_target_func}), + ]) + + target_edge = next(e for e in edges if e.route == 'route_x') + assert isinstance(target_edge.to_node, FunctionNode) + assert target_edge.to_node.name == 'my_target_func' + + +def test_routing_map_node_reuse() -> None: + """Tests that the same callable used in a map and elsewhere is deduplicated.""" + + def my_func() -> None: + pass + + node_b = TestingNode(name='NodeB') + edges = parse_edge_items([ + (START, my_func), + (my_func, {'route_x': node_b}), + ]) + + # my_func should be wrapped once and reused. + assert len(edges) == 2 + assert edges[0].to_node is edges[1].from_node + assert isinstance(edges[0].to_node, FunctionNode) + + +def test_routing_map_empty_dict_raises() -> None: + """Tests that an empty routing map raises ValueError.""" + node_a = TestingNode(name='NodeA') + with pytest.raises( + ValueError, + match=r'Routing map must not be empty', + ): + parse_edge_items([ + (START, node_a), + (node_a, {}), + ]) + + +def test_routing_map_invalid_key_raises() -> None: + """Tests that a non-RouteValue key raises ValueError.""" + node_a = TestingNode(name='NodeA') + node_b = TestingNode(name='NodeB') + with pytest.raises( + ValueError, + match=r'Invalid routing map key', + ): + parse_edge_items([ + (START, node_a), + (node_a, {1.5: node_b}), + ]) + + +def test_routing_map_invalid_value_raises() -> None: + """Tests that a non-NodeLike value raises ValueError.""" + node_a = TestingNode(name='NodeA') + with pytest.raises( + ValueError, + match=r'Invalid routing map value', + ): + parse_edge_items([ + (START, node_a), + (node_a, {'route_x': 42}), + ]) + + +def test_routing_map_fan_out_target() -> None: + """Tests that a tuple value in a routing map creates fan-out edges.""" + node_a = TestingNode(name='NodeA') + node_b = TestingNode(name='NodeB') + node_c = TestingNode(name='NodeC') + edges = parse_edge_items([ + (START, node_a), + (node_a, {'route_x': (node_b, node_c)}), + ]) + + # START->A, A->B(route_x), A->C(route_x) + assert len(edges) == 3 + + routed_edges = [e for e in edges if e.route is not None] + assert len(routed_edges) == 2 + + # Both fan-out edges share the same route and source. + for e in routed_edges: + assert e.from_node.name == 'NodeA' + assert e.route == 'route_x' + + targets = {e.to_node.name for e in routed_edges} + assert targets == {'NodeB', 'NodeC'} + + +def test_routing_map_fan_out_invalid_element_raises() -> None: + """Tests that a non-NodeLike element inside a fan-out tuple raises.""" + node_a = TestingNode(name='NodeA') + node_b = TestingNode(name='NodeB') + with pytest.raises( + ValueError, + match=r'Invalid node in fan-out tuple', + ): + parse_edge_items([ + (START, node_a), + (node_a, {'route_x': (node_b, 42)}), + ]) + + +# --- Routing map as chain element tests --- + + +def test_routing_map_chain_ending_with_dict() -> None: + """Tests a chain ending with a routing map creates correct edges.""" + node_a = TestingNode(name='NodeA') + node_b = TestingNode(name='NodeB') + node_c = TestingNode(name='NodeC') + edges = parse_edge_items([ + (START, node_a, {'r1': node_b, 'r2': node_c}), + ]) + + # START->A (None), A->B (r1), A->C (r2) + assert len(edges) == 3 + + start_edge = next(e for e in edges if e.from_node.name == '__START__') + assert start_edge.to_node.name == 'NodeA' + assert start_edge.route is None + + routed_edges = [e for e in edges if e.route is not None] + assert len(routed_edges) == 2 + routes_and_targets = {(e.route, e.to_node.name) for e in routed_edges} + assert routes_and_targets == {('r1', 'NodeB'), ('r2', 'NodeC')} + for e in routed_edges: + assert e.from_node.name == 'NodeA' + + +def test_routing_map_mid_chain_with_fan_in() -> None: + """Tests routing map mid-chain with fan-in to the next element.""" + node_a = TestingNode(name='NodeA') + node_b = TestingNode(name='NodeB') + node_c = TestingNode(name='NodeC') + node_d = TestingNode(name='NodeD') + edges = parse_edge_items([ + (START, node_a, {'r1': node_b, 'r2': node_c}, node_d), + ]) + + # START->A (None), A->B (r1), A->C (r2), B->D (None), C->D (None) + assert len(edges) == 5 + + routed_edges = sorted( + [e for e in edges if e.route is not None], + key=lambda e: e.to_node.name, + ) + assert len(routed_edges) == 2 + assert routed_edges[0].from_node.name == 'NodeA' + assert routed_edges[0].to_node.name == 'NodeB' + assert routed_edges[0].route == 'r1' + assert routed_edges[1].from_node.name == 'NodeA' + assert routed_edges[1].to_node.name == 'NodeC' + assert routed_edges[1].route == 'r2' + + fan_in_edges = sorted( + [e for e in edges if e.to_node.name == 'NodeD'], + key=lambda e: e.from_node.name, + ) + assert len(fan_in_edges) == 2 + assert fan_in_edges[0].from_node.name == 'NodeB' + assert fan_in_edges[0].route is None + assert fan_in_edges[1].from_node.name == 'NodeC' + assert fan_in_edges[1].route is None + + +def test_routing_map_mid_chain_fan_out_values() -> None: + """Tests routing map with fan-out tuple values, followed by fan-in.""" + node_a = TestingNode(name='NodeA') + node_b = TestingNode(name='NodeB') + node_c = TestingNode(name='NodeC') + node_d = TestingNode(name='NodeD') + edges = parse_edge_items([ + (START, node_a, {'r1': (node_b, node_c)}, node_d), + ]) + + # START->A (None), A->B (r1), A->C (r1), B->D (None), C->D (None) + assert len(edges) == 5 + + routed_edges = [e for e in edges if e.route is not None] + assert len(routed_edges) == 2 + for e in routed_edges: + assert e.from_node.name == 'NodeA' + assert e.route == 'r1' + + fan_in_edges = [e for e in edges if e.to_node.name == 'NodeD'] + assert len(fan_in_edges) == 2 + fan_in_sources = {e.from_node.name for e in fan_in_edges} + assert fan_in_sources == {'NodeB', 'NodeC'} + + +def test_routing_map_consecutive_dicts_raises() -> None: + """Tests that consecutive routing maps in a chain are rejected.""" + node_a = TestingNode(name='NodeA') + node_b = TestingNode(name='NodeB') + node_c = TestingNode(name='NodeC') + node_d = TestingNode(name='NodeD') + with pytest.raises( + ValueError, match=r'Consecutive routing maps are not allowed' + ): + parse_edge_items([ + (START, node_a, {'r1': node_b, 'r2': node_c}, {'r3': node_d}), + ]) + + +def test_routing_map_empty_dict_in_chain_raises() -> None: + """Tests that an empty routing map in a chain raises ValueError.""" + node_a = TestingNode(name='NodeA') + node_b = TestingNode(name='NodeB') + with pytest.raises(ValueError, match=r'Routing map must not be empty'): + parse_edge_items([ + (START, node_a, {}, node_b), + ]) + + +def test_routing_map_invalid_key_in_chain_raises() -> None: + """Tests that invalid routing map keys in a chain raise ValueError.""" + node_a = TestingNode(name='NodeA') + node_b = TestingNode(name='NodeB') + with pytest.raises(ValueError, match=r'Invalid routing map key'): + parse_edge_items([ + (START, node_a, {1.5: node_b}), + ]) + + +def test_routing_map_2_tuple_backward_compat() -> None: + """Ensures existing 2-tuple routing map syntax still works.""" + node_a = TestingNode(name='NodeA') + node_b = TestingNode(name='NodeB') + node_c = TestingNode(name='NodeC') + edges = parse_edge_items([ + (START, node_a), + (node_a, {'r1': node_b, 'r2': node_c}), + ]) + assert len(edges) == 3 diff --git a/tests/unittests/workflow/utils/test_graph_validation.py b/tests/unittests/workflow/utils/test_graph_validation.py new file mode 100644 index 00000000000..10ef747e0e9 --- /dev/null +++ b/tests/unittests/workflow/utils/test_graph_validation.py @@ -0,0 +1,373 @@ +# 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. + +"""Tests for Graph validation utility.""" + +import logging + +from google.adk.workflow import Edge +from google.adk.workflow import START +from google.adk.workflow._graph import DEFAULT_ROUTE +from google.adk.workflow._graph import Graph +from google.adk.workflow.utils._graph_validation import validate_graph +from pydantic import BaseModel +import pytest + +from ..workflow_testing_utils import TestingNode + + +def test_missing_start_node() -> None: + """Tests that a graph missing the START node fails validation.""" + node_a = TestingNode(name='NodeA') + node_b = TestingNode(name='NodeB') + graph = Graph( + edges=[Edge(from_node=node_a, to_node=node_b)], + ) + with pytest.raises( + ValueError, + match=( + r"Graph validation failed\. START node \(name: '__START__'\) not" + r' found in graph nodes\.' + ), + ): + validate_graph(graph.nodes, graph.edges) + + +def test_unreachable_node() -> None: + """Tests that a graph with an unreachable node fails validation.""" + node_a = TestingNode(name='NodeA') + node_b = TestingNode(name='NodeB') # Unreachable + graph = Graph( + edges=[ + Edge(from_node=START, to_node=node_a), + Edge(from_node=node_b, to_node=node_a), + ], + ) + with pytest.raises( + ValueError, + match=( + r'Graph validation failed\. The following nodes are unreachable' + r" from START: \['NodeB'\]" + ), + ): + validate_graph(graph.nodes, graph.edges) + + +def test_disconnected_routed_subgraph_is_unreachable() -> None: + """Tests that a disconnected subgraph with routed edges fails validation.""" + node_a = TestingNode(name='NodeA') + node_b = TestingNode(name='NodeB') + node_c = TestingNode(name='NodeC') + graph = Graph( + edges=[ + Edge(from_node=START, to_node=node_a), + Edge(from_node=node_b, to_node=node_c, route='x'), + Edge(from_node=node_c, to_node=node_b, route='y'), + ], + ) + with pytest.raises( + ValueError, + match=( + r'Graph validation failed\. The following nodes are unreachable' + r" from START: \['NodeB', 'NodeC'\]" + ), + ): + validate_graph(graph.nodes, graph.edges) + + +@pytest.mark.parametrize( + 'routes', + [ + (None, None), + ('route1', 'route1'), + ('route1', 'route2'), + ('route1', None), + ], +) +def test_duplicate_edges_fail_validation( + routes: tuple[str | None, str | None], +) -> None: + """Tests that duplicate edges fail validation, regardless of routes.""" + node_a = TestingNode(name='NodeA') + graph = Graph( + edges=[ + Edge( + from_node=START, + to_node=node_a, + route=routes[0], + ), + Edge( + from_node=START, + to_node=node_a, + route=routes[1], + ), + ], + ) + with pytest.raises( + ValueError, + match=( + r'Graph validation failed\. Duplicate edge found: from=__START__,' + r' to=NodeA' + ), + ): + validate_graph(graph.nodes, graph.edges) + + +def test_start_node_with_incoming_edge() -> None: + """Tests graph with incoming edge to START node fails validation.""" + node_a = TestingNode(name='NodeA') + graph = Graph( + edges=[ + Edge(from_node=node_a, to_node=START), + Edge(from_node=START, to_node=node_a), + ], + ) + with pytest.raises( + ValueError, + match=( + r'Graph validation failed\. START node must not have incoming edges\.' + ), + ): + validate_graph(graph.nodes, graph.edges) + + +def test_multiple_default_routes_fail_validation() -> None: + """Tests that multiple DEFAULT_ROUTE edges from a node fail validation.""" + node_a = TestingNode(name='NodeA') + node_b = TestingNode(name='NodeB') + node_c = TestingNode(name='NodeC') + graph = Graph( + edges=[ + Edge(from_node=START, to_node=node_a), + Edge(from_node=node_a, to_node=node_b, route=DEFAULT_ROUTE), + Edge(from_node=node_a, to_node=node_c, route=DEFAULT_ROUTE), + ], + ) + with pytest.raises( + ValueError, + match=( + r'Graph validation failed\. Multiple DEFAULT_ROUTE edges found from' + r' node NodeA to NodeB and NodeC' + ), + ): + validate_graph(graph.nodes, graph.edges) + + +def test_single_default_route_passes_validation() -> None: + """Tests that a single DEFAULT_ROUTE edge from a node passes validation.""" + node_a = TestingNode(name='NodeA') + node_b = TestingNode(name='NodeB') + node_c = TestingNode(name='NodeC') + graph = Graph( + edges=[ + Edge(from_node=START, to_node=node_a), + Edge(from_node=node_a, to_node=node_b, route=DEFAULT_ROUTE), + Edge(from_node=node_a, to_node=node_c, route='another_route'), + ], + ) + validate_graph(graph.nodes, graph.edges) # Should not raise + + +def test_duplicate_node_names_fail_validation() -> None: + """Tests that duplicate nodes raise error.""" + node_a1 = TestingNode(name='NodeA') + node_a2 = TestingNode(name='NodeA') + graph = Graph( + edges=[ + Edge(from_node=START, to_node=node_a1), + Edge(from_node=node_a1, to_node=node_a2), + ], + ) + with pytest.raises( + ValueError, + match=( + r"Graph validation failed\. Duplicate node names found: \['NodeA'\]\." + r' This means multiple distinct node objects have the same name\. If' + r' you intended to reuse the same node, ensure you pass the exact' + r' same object instance\. If you intended to have distinct nodes,' + r' ensure they have unique names\.' + ), + ): + validate_graph(graph.nodes, graph.edges) + + +def test_unconditional_cycle_fails_validation() -> None: + """Tests that a cycle of unconditional edges (route=None) fails.""" + node_a = TestingNode(name='NodeA') + node_b = TestingNode(name='NodeB') + graph = Graph( + edges=[ + Edge(from_node=START, to_node=node_a), + Edge(from_node=node_a, to_node=node_b), + Edge(from_node=node_b, to_node=node_a), + ], + ) + with pytest.raises( + ValueError, + match=r'Graph validation failed\. Unconditional cycle detected:', + ): + validate_graph(graph.nodes, graph.edges) + + +def test_unconditional_self_loop_fails_validation() -> None: + """Tests that an unconditional self-loop (A -> A) fails.""" + node_a = TestingNode(name='NodeA') + graph = Graph( + edges=[ + Edge(from_node=START, to_node=node_a), + Edge(from_node=node_a, to_node=node_a), + ], + ) + with pytest.raises( + ValueError, + match=r'Graph validation failed\. Unconditional cycle detected:', + ): + validate_graph(graph.nodes, graph.edges) + + +def test_longer_unconditional_cycle_fails_validation() -> None: + """Tests that a longer unconditional cycle (A -> B -> C -> A) fails.""" + node_a = TestingNode(name='NodeA') + node_b = TestingNode(name='NodeB') + node_c = TestingNode(name='NodeC') + graph = Graph( + edges=[ + Edge(from_node=START, to_node=node_a), + Edge(from_node=node_a, to_node=node_b), + Edge(from_node=node_b, to_node=node_c), + Edge(from_node=node_c, to_node=node_a), + ], + ) + with pytest.raises( + ValueError, + match=r'Graph validation failed\. Unconditional cycle detected:', + ): + validate_graph(graph.nodes, graph.edges) + + +def test_conditional_cycle_passes_validation() -> None: + """Tests that a cycle with a routed edge (loop pattern) passes.""" + node_a = TestingNode(name='NodeA') + node_b = TestingNode(name='NodeB') + graph = Graph( + edges=[ + Edge(from_node=START, to_node=node_a), + Edge(from_node=node_a, to_node=node_b), + Edge(from_node=node_b, to_node=node_a, route='retry'), + ], + ) + validate_graph( + graph.nodes, graph.edges + ) # Should not raise — routed back-edge + + +def test_conditional_self_loop_passes_validation() -> None: + """Tests that a self-loop with a route passes validation.""" + node_a = TestingNode(name='NodeA') + node_b = TestingNode(name='NodeB') + graph = Graph( + edges=[ + Edge(from_node=START, to_node=node_a), + Edge(from_node=node_a, to_node=node_a, route='continue'), + Edge(from_node=node_a, to_node=node_b, route='done'), + ], + ) + validate_graph( + graph.nodes, graph.edges + ) # Should not raise — routed self-loop + + +def test_dag_with_diamond_passes_validation() -> None: + """Tests that a DAG with a diamond shape passes validation.""" + node_a = TestingNode(name='NodeA') + node_b = TestingNode(name='NodeB') + node_c = TestingNode(name='NodeC') + graph = Graph( + edges=[ + Edge(from_node=START, to_node=node_a), + Edge(from_node=START, to_node=node_b), + Edge(from_node=node_a, to_node=node_c), + Edge(from_node=node_b, to_node=node_c), + ], + ) + validate_graph(graph.nodes, graph.edges) # Should not raise + + +class ModelA(BaseModel): + x: int + + +class ModelB(BaseModel): + x: int + + +def test_schema_match_passes() -> None: + """Tests that edges with matching schemas pass validation.""" + node_a = TestingNode(name='NodeA', output_schema=ModelA) + node_b = TestingNode(name='NodeB', input_schema=ModelA) + graph = Graph( + edges=[ + Edge(from_node=START, to_node=node_a), + Edge(from_node=node_a, to_node=node_b), + ], + ) + validate_graph(graph.nodes, graph.edges) # Should not raise + + +def test_schema_mismatch_raises() -> None: + """Tests that edges with mismatching schemas fail validation.""" + node_a = TestingNode(name='NodeA', output_schema=ModelA) + node_b = TestingNode(name='NodeB', input_schema=ModelB) + graph = Graph( + edges=[ + Edge(from_node=START, to_node=node_a), + Edge(from_node=node_a, to_node=node_b), + ], + ) + with pytest.raises( + ValueError, + match=r'Graph validation failed\. Schema mismatch on edge', + ): + validate_graph(graph.nodes, graph.edges) + + +def test_schema_missing_passes() -> None: + """Tests that edges with missing schemas pass validation.""" + node_a = TestingNode(name='NodeA', output_schema=ModelA) + node_b = TestingNode(name='NodeB') # No input schema + graph = Graph( + edges=[ + Edge(from_node=START, to_node=node_a), + Edge(from_node=node_a, to_node=node_b), + ], + ) + validate_graph(graph.nodes, graph.edges) # Should not raise + + +def test_chat_agent_wiring_validation_only_runs_on_llm_agent() -> None: + """Tests that _validate_chat_agent_wiring checks non-LlmAgent nodes safely.""" + node_a = TestingNode(name='NodeA') + node_b = TestingNode(name='NodeB') + # Set mode='chat' on a non-LlmAgent node + object.__setattr__(node_b, 'mode', 'chat') + + graph = Graph( + edges=[ + Edge(from_node=START, to_node=node_a), + Edge(from_node=node_a, to_node=node_b), + ], + ) + validate_graph( + graph.nodes, graph.edges + ) # Should not raise because node_b is a TestingNode, not LlmAgent From aceb7be08144748a9591984a7fb06463f4689204 Mon Sep 17 00:00:00 2001 From: Shangjie Chen Date: Wed, 1 Jul 2026 11:15:41 -0700 Subject: [PATCH 05/18] refactor: Clean up unused variables and arguments in adk/workflow - Remove unused local variable `target_state` in `_workflow.py`. - Remove unused `ctx` argument from `NodeRunner._attempt_retry`. - Remove unused `node_path` and `curr_parent_ctx` arguments from `check_interception` utility function. Co-authored-by: Shangjie Chen PiperOrigin-RevId: 941198065 --- .../adk/workflow/_dynamic_node_scheduler.py | 2 -- src/google/adk/workflow/_node_runner.py | 6 ++-- src/google/adk/workflow/_workflow.py | 3 -- .../adk/workflow/utils/_replay_interceptor.py | 2 -- .../workflow/utils/test_replay_interceptor.py | 34 ------------------- 5 files changed, 2 insertions(+), 45 deletions(-) diff --git a/src/google/adk/workflow/_dynamic_node_scheduler.py b/src/google/adk/workflow/_dynamic_node_scheduler.py index 6cdd45e6c14..e59e1ac624c 100644 --- a/src/google/adk/workflow/_dynamic_node_scheduler.py +++ b/src/google/adk/workflow/_dynamic_node_scheduler.py @@ -266,11 +266,9 @@ async def _check_existing_run( # Delegate replay and same-turn interception check to ReplayInterceptor. result = check_interception( - node_path=node_path, node=curr_node, recovered=run.recovered_state, current_run=run, - curr_parent_ctx=curr_parent_ctx, ) if not result.should_run: diff --git a/src/google/adk/workflow/_node_runner.py b/src/google/adk/workflow/_node_runner.py index 11158773346..e2d3a3dbe70 100644 --- a/src/google/adk/workflow/_node_runner.py +++ b/src/google/adk/workflow/_node_runner.py @@ -157,7 +157,7 @@ async def run( ) await self._enqueue_event(error_event, ctx) - if not await self._attempt_retry(e, ctx, attempt_count): + if not await self._attempt_retry(e, attempt_count): ctx._error = e ctx._error_node_path = ctx.node_path logger.debug("node %s end.", ctx.node_path) @@ -169,9 +169,7 @@ async def run( ) attempt_count += 1 - async def _attempt_retry( - self, e: Exception, ctx: Context, attempt_count: int - ) -> bool: + async def _attempt_retry(self, e: Exception, attempt_count: int) -> bool: """Checks if node should retry and sleeps if so.""" from ._node_state import NodeState from .utils._retry_utils import _get_retry_delay diff --git a/src/google/adk/workflow/_workflow.py b/src/google/adk/workflow/_workflow.py index d022dcd8528..b6c92ab673c 100644 --- a/src/google/adk/workflow/_workflow.py +++ b/src/google/adk/workflow/_workflow.py @@ -555,10 +555,8 @@ def _start_node_task( recovered = loop_state.recovered_executions[key] result = check_interception( - node_path=f'{ctx.node_path}/{node_name}@{run_id}', node=node, recovered=recovered, - curr_parent_ctx=ctx, ) if not result.should_run: @@ -688,7 +686,6 @@ def _buffer_downstream_triggers( use_sub_branch = len(next_nodes) > 1 for target_name in next_nodes: target_node = self._get_static_node_by_name(target_name) - target_state = loop_state.nodes.get(target_name) if target_node._requires_all_predecessors: # Wait for all predecessors diff --git a/src/google/adk/workflow/utils/_replay_interceptor.py b/src/google/adk/workflow/utils/_replay_interceptor.py index f72ca62cb62..f39dc1774fa 100644 --- a/src/google/adk/workflow/utils/_replay_interceptor.py +++ b/src/google/adk/workflow/utils/_replay_interceptor.py @@ -56,11 +56,9 @@ class InterceptionResult: def check_interception( *, - node_path: str, node: BaseNode, recovered: _ChildScanState | None = None, current_run: DynamicNodeRun | None = None, - curr_parent_ctx: Context, ) -> InterceptionResult: """Determine if a node execution should be intercepted based on history.""" from .._workflow import Workflow # pylint: disable=g-import-not-at-top diff --git a/tests/unittests/workflow/utils/test_replay_interceptor.py b/tests/unittests/workflow/utils/test_replay_interceptor.py index edd9a737b08..1c1bd74ca7b 100644 --- a/tests/unittests/workflow/utils/test_replay_interceptor.py +++ b/tests/unittests/workflow/utils/test_replay_interceptor.py @@ -18,9 +18,6 @@ replay interception. """ -from unittest.mock import MagicMock - -from google.adk.agents.context import Context from google.adk.workflow._base_node import BaseNode from google.adk.workflow._dynamic_node_scheduler import DynamicNodeRun from google.adk.workflow._node_state import NodeState @@ -30,13 +27,6 @@ import pytest -def _make_parent_ctx(): - ctx = MagicMock(spec=Context) - ctx._invocation_context = MagicMock() - ctx.resume_inputs = {} - return ctx - - def test_same_turn_completed(): """Same-turn completed run intercepts and returns cached output.""" # Given a same-turn completed run @@ -45,14 +35,11 @@ def test_same_turn_completed(): output='cached-out', transfer_to_agent='target-agent', ) - ctx = _make_parent_ctx() # When checked result = check_interception( - node_path='wf/node@1', node=BaseNode(name='node'), current_run=run, - curr_parent_ctx=ctx, ) # Then it intercepts with cached results @@ -67,14 +54,11 @@ def test_same_turn_waiting(): run = DynamicNodeRun( state=NodeState(status=NodeStatus.WAITING, interrupts=['fc-1']), ) - ctx = _make_parent_ctx() # When checked result = check_interception( - node_path='wf/node@1', node=BaseNode(name='node'), current_run=run, - curr_parent_ctx=ctx, ) # Then it intercepts and keeps waiting @@ -91,14 +75,11 @@ def test_cross_turn_unresolved_interrupts_no_rerun(): resolved_ids={'fc-1'}, ) node = BaseNode(name='node', rerun_on_resume=False) - ctx = _make_parent_ctx() # When checked result = check_interception( - node_path='wf/node@1', node=node, recovered=recovered, - curr_parent_ctx=ctx, ) # Then it stays waiting on unresolved interrupts @@ -116,14 +97,11 @@ def test_cross_turn_unresolved_interrupts_rerun(): resolved_responses={'fc-1': 'ans'}, ) node = BaseNode(name='node', rerun_on_resume=True) - ctx = _make_parent_ctx() # When checked result = check_interception( - node_path='wf/node@1', node=node, recovered=recovered, - curr_parent_ctx=ctx, ) # Then it reruns with partial resolved inputs @@ -140,14 +118,11 @@ def test_cross_turn_completed(): route='route-a', ) node = BaseNode(name='node') - ctx = _make_parent_ctx() # When checked result = check_interception( - node_path='wf/node@1', node=node, recovered=recovered, - curr_parent_ctx=ctx, ) # Then it fast-forwards with cached output and route @@ -166,17 +141,11 @@ def test_cross_turn_all_resolved_no_rerun(): resolved_responses={'fc-1': 'ans'}, ) node = BaseNode(name='node', rerun_on_resume=False) - ctx = _make_parent_ctx() - ctx.resume_inputs = { - 'fc-1': {'result': 'ans'} - } # Simulate FunctionResponse dict # When checked result = check_interception( - node_path='wf/node@1', node=node, recovered=recovered, - curr_parent_ctx=ctx, ) # Then it auto-completes @@ -194,14 +163,11 @@ def test_cross_turn_all_resolved_rerun(): resolved_responses={'fc-1': 'ans'}, ) node = BaseNode(name='node', rerun_on_resume=True) - ctx = _make_parent_ctx() # When checked result = check_interception( - node_path='wf/node@1', node=node, recovered=recovered, - curr_parent_ctx=ctx, ) # Then it reruns From 14a24f2beeb8583d98c5b1b1933e82652751e0aa Mon Sep 17 00:00:00 2001 From: Anwesha Das Date: Wed, 1 Jul 2026 11:33:15 -0700 Subject: [PATCH 06/18] =?UTF-8?q?feat(bigtable):=20Support=20parameterized?= =?UTF-8?q?=20views=20with=20secure=20parameter=20inj=E2=80=A6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merge https://github.com/google/adk-python/pull/6128 …ection Expose a parameterized query tool execute_sql_parameterized that automatically maps and injects secure parameters (like user_id) from the tool context to Bigtable's view_parameters. **Please ensure you have read the [contribution guide](https://github.com/google/adk-python/blob/main/CONTRIBUTING.md) before creating a pull request.** ### Link to Issue or Description of Change **1. Link to an existing issue (if applicable):** - Closes: #_issue_number_ - Related: #_issue_number_ **2. Or, if no issue exists, describe the change:** _If applicable, please follow the issue templates to provide as much detail as possible._ **Problem:** _A clear and concise description of what the problem is._ **Solution:** _A clear and concise description of what you want to happen and why you choose this solution._ ### Testing Plan _Please describe the tests that you ran to verify your changes. This is required for all PRs that are not small documentation or typo fixes._ **Unit Tests:** - [ ] I have added or updated unit tests for my change. - [ ] All unit tests pass locally. _Please include a summary of passed `pytest` results._ **Manual End-to-End (E2E) Tests:** _Please provide instructions on how to manually test your changes, including any necessary setup or configuration. Please provide logs or screenshots to help reviewers better understand the fix._ ### Checklist - [ ] I have read the [CONTRIBUTING.md](https://github.com/google/adk-python/blob/main/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. ### Additional context _Add any other context or screenshots about the feature request here._ COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/6128 from ad548:feat/bigtable-parameterized-views f5902fd7ef834f326a43c7c01990d0648e7ec06c PiperOrigin-RevId: 941207338 --- pyproject.toml | 6 +- .../adk/tools/bigtable/bigtable_toolset.py | 95 ++++++++ src/google/adk/tools/bigtable/query_tool.py | 3 + .../bigtable/test_bigtable_query_tool.py | 79 +++++++ .../tools/bigtable/test_bigtable_toolset.py | 205 ++++++++++++++++++ 5 files changed, 385 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index c746c5aa709..3daff0eba28 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -70,7 +70,7 @@ optional-dependencies.all = [ "google-cloud-aiplatform[agent-engines]>=1.148.1,<2", "google-cloud-bigquery>=2.2", "google-cloud-bigquery-storage>=2", - "google-cloud-bigtable>=2.32", + "google-cloud-bigtable>=2.39.1", "google-cloud-dataplex>=1.7,<3", "google-cloud-discoveryengine>=0.13.12,<0.14", "google-cloud-parametermanager>=0.4,<1", @@ -165,7 +165,7 @@ optional-dependencies.gcp = [ "google-cloud-aiplatform[agent-engines]>=1.148.1,<2", "google-cloud-bigquery>=2.2", "google-cloud-bigquery-storage>=2", - "google-cloud-bigtable>=2.32", + "google-cloud-bigtable>=2.39.1", "google-cloud-dataplex>=1.7,<3", "google-cloud-discoveryengine>=0.13.12,<0.14", "google-cloud-parametermanager>=0.4,<1", @@ -208,7 +208,7 @@ optional-dependencies.test = [ "google-cloud-aiplatform[agent-engines,evaluation]>=1.148.1,<2", "google-cloud-bigquery>=2.2", "google-cloud-bigquery-storage>=2", - "google-cloud-bigtable>=2.32", + "google-cloud-bigtable>=2.39.1", "google-cloud-dataplex>=1.7,<3", "google-cloud-discoveryengine>=0.13.12,<0.14", "google-cloud-firestore>=2.11,<3", diff --git a/src/google/adk/tools/bigtable/bigtable_toolset.py b/src/google/adk/tools/bigtable/bigtable_toolset.py index 97fc2eb0b63..704908b8e87 100644 --- a/src/google/adk/tools/bigtable/bigtable_toolset.py +++ b/src/google/adk/tools/bigtable/bigtable_toolset.py @@ -14,11 +14,16 @@ from __future__ import annotations +import inspect +from typing import Any +from typing import Callable from typing import List from typing import Optional from typing import Union from google.adk.agents.readonly_context import ReadonlyContext +from google.auth.credentials import Credentials +from pydantic import BaseModel from typing_extensions import override from . import metadata_tool @@ -29,9 +34,88 @@ from ...tools.base_toolset import BaseToolset from ...tools.base_toolset import ToolPredicate from ...tools.google_tool import GoogleTool +from ..tool_context import ToolContext from .bigtable_credentials import BigtableCredentialsConfig from .settings import BigtableToolSettings + +class BigtableParameterizedViewTool(GoogleTool): + """Wrapper FunctionTool for Bigtable execute_sql query tool that passes view parameters. + + This tool wraps the Bigtable query tool to automatically resolve and inject + parameters from the ToolContext (e.g. user_id) into the query's + view_parameters. The parameter names to resolve are configured via + view_parameter_names. + + Example: + If a parameterized view `purchase_history_pv` was created with the query: + `SELECT * FROM purchases WHERE user_id = VIEW_PARAMETERS('user_id')` + + By configuring `view_parameter_names=["user_id"]`, the wrapper will + resolve the `user_id` value from the `tool_context.user_id` at runtime and + pass it as `view_parameters={"user_id": user_id}`. + This securely restricts query execution to the logged-in user's data + without exposing the `user_id` parameter to the LLM. + """ + + def __init__( + self, + func: Callable[..., Any], + *, + credentials_config: Optional[BigtableCredentialsConfig] = None, + tool_settings: Optional[BigtableToolSettings] = None, + view_parameter_names: Optional[List[str]] = None, + ): + """Initializes the BigtableParameterizedViewTool. + + Args: + func: The Bigtable query function to wrap. + credentials_config: The credentials configuration. + tool_settings: The tool settings. + view_parameter_names: A list of parameter names to resolve from + tool_context and pass into view_parameters. This is configured on the + toolset (BigtableToolset) and forwarded here. + """ + super().__init__( + func=func, + credentials_config=credentials_config, + tool_settings=tool_settings, + ) + self.name = "execute_sql_parameterized" + self.description = ( + "Execute a GoogleSQL query from a Bigtable table using parameterized" + " views to securely check permissions." + ) + self.view_parameter_names = view_parameter_names + # Exclude from being parsed and exposed to the LLM when generating tool schemas + self._ignore_params.append("_view_parameters") + + @override + async def _run_async_with_credential( + self, + credentials: Credentials, + tool_settings: BaseModel, + args: dict[str, Any], + tool_context: ToolContext, + ) -> Any: + args_to_call = args.copy() + signature = inspect.signature(self.func) + if "_view_parameters" in signature.parameters and self.view_parameter_names: + view_params = {} + for param_name in self.view_parameter_names: + # 1. Check if it's a strongly-typed top-level property (like 'user_id') + if (val := getattr(tool_context, param_name, None)) is not None: + view_params[param_name] = val + # 2. Fallback to checking application-level session state + elif tool_context.state and param_name in tool_context.state: + view_params[param_name] = tool_context.state[param_name] + + args_to_call["_view_parameters"] = view_params + return await super()._run_async_with_credential( + credentials, tool_settings, args_to_call, tool_context + ) + + DEFAULT_BIGTABLE_TOOL_NAME_PREFIX = "bigtable" @@ -55,6 +139,7 @@ def __init__( tool_filter: Optional[Union[ToolPredicate, List[str]]] = None, credentials_config: Optional[BigtableCredentialsConfig] = None, bigtable_tool_settings: Optional[BigtableToolSettings] = None, + view_parameter_names: Optional[List[str]] = None, ): super().__init__( tool_filter=tool_filter, @@ -66,6 +151,7 @@ def __init__( if bigtable_tool_settings else BigtableToolSettings() ) + self.view_parameter_names = view_parameter_names def _is_tool_selected( self, tool: BaseTool, readonly_context: ReadonlyContext @@ -102,6 +188,15 @@ async def get_tools( query_tool.execute_sql, ] ] + if self.view_parameter_names: + all_tools.append( + BigtableParameterizedViewTool( + func=query_tool.execute_sql, + credentials_config=self._credentials_config, + tool_settings=self._tool_settings, + view_parameter_names=self.view_parameter_names, + ) + ) return [ tool for tool in all_tools diff --git a/src/google/adk/tools/bigtable/query_tool.py b/src/google/adk/tools/bigtable/query_tool.py index bf64b282a12..195d6a08314 100644 --- a/src/google/adk/tools/bigtable/query_tool.py +++ b/src/google/adk/tools/bigtable/query_tool.py @@ -42,6 +42,7 @@ async def execute_sql( tool_context: ToolContext, parameters: Dict[str, Any] | None = None, parameter_types: Dict[str, Any] | None = None, + _view_parameters: Dict[str, Any] | None = None, ) -> dict: """Execute a GoogleSQL query from a Bigtable table. @@ -56,6 +57,7 @@ async def execute_sql( parameters (dict): properties for parameter replacement. Keys must match the names used in ``query``. parameter_types (dict): maps explicit types for one or more param values. + _view_parameters (dict): maps properties for parameterized views. Returns: dict: Dictionary containing the status and the rows read. @@ -91,6 +93,7 @@ def _execute_sql(): instance_id=instance_id, parameters=parameters, parameter_types=parameter_types, + view_parameters=_view_parameters, ) rows: List[Dict[str, Any]] = [] diff --git a/tests/unittests/tools/bigtable/test_bigtable_query_tool.py b/tests/unittests/tools/bigtable/test_bigtable_query_tool.py index 46b65a3a07c..3873fedde4a 100644 --- a/tests/unittests/tools/bigtable/test_bigtable_query_tool.py +++ b/tests/unittests/tools/bigtable/test_bigtable_query_tool.py @@ -193,6 +193,7 @@ def raise_error(): instance_id=instance_id, parameters=parameters, parameter_types=parameter_types, + view_parameters=None, ) mock_iterator.close.assert_called_once() @@ -228,3 +229,81 @@ async def test_execute_sql_row_value_circular_reference_fallback(): assert result["status"] == "SUCCESS" assert result["rows"][0]["col1"] == str(circular_value) + + +@pytest.mark.asyncio +async def test_execute_sql_with_view_parameters(): + """Test execute_sql with _view_parameters passed.""" + project = "my_project" + instance_id = "my_instance" + query = "SELECT * FROM my_table" + credentials = mock.create_autospec(Credentials, instance=True) + tool_context = mock.create_autospec(ToolContext, instance=True) + view_parameters = {"user_id": "test-user-123"} + + with mock.patch.object(client, "get_bigtable_data_client") as mock_get_client: + mock_client = mock.MagicMock() + mock_get_client.return_value = mock_client + mock_iterator = mock.create_autospec(ExecuteQueryIterator, instance=True) + mock_client.execute_query.return_value = mock_iterator + mock_iterator.__iter__.return_value = [] + + result = await execute_sql( + project_id=project, + instance_id=instance_id, + credentials=credentials, + query=query, + settings=BigtableToolSettings(), + tool_context=tool_context, + _view_parameters=view_parameters, + ) + + assert result["status"] == "SUCCESS" + mock_client.execute_query.assert_called_once_with( + query=query, + instance_id=instance_id, + parameters=None, + parameter_types=None, + view_parameters=view_parameters, + ) + + +@pytest.mark.asyncio +async def test_execute_sql_with_multiple_view_parameters(): + """Test execute_sql with multiple view_parameters of different names.""" + project = "my_project" + instance_id = "my_instance" + query = "SELECT * FROM my_table" + credentials = mock.create_autospec(Credentials, instance=True) + tool_context = mock.create_autospec(ToolContext, instance=True) + view_parameters = { + "user_id": "test-user-123", + "tenant_id": "tenant-xyz", + "role": "admin", + } + + with mock.patch.object(client, "get_bigtable_data_client") as mock_get_client: + mock_client = mock.MagicMock() + mock_get_client.return_value = mock_client + mock_iterator = mock.create_autospec(ExecuteQueryIterator, instance=True) + mock_client.execute_query.return_value = mock_iterator + mock_iterator.__iter__.return_value = [] + + result = await execute_sql( + project_id=project, + instance_id=instance_id, + credentials=credentials, + query=query, + settings=BigtableToolSettings(), + tool_context=tool_context, + _view_parameters=view_parameters, + ) + + assert result["status"] == "SUCCESS" + mock_client.execute_query.assert_called_once_with( + query=query, + instance_id=instance_id, + parameters=None, + parameter_types=None, + view_parameters=view_parameters, + ) diff --git a/tests/unittests/tools/bigtable/test_bigtable_toolset.py b/tests/unittests/tools/bigtable/test_bigtable_toolset.py index b5698cfc075..aa5fd440be2 100644 --- a/tests/unittests/tools/bigtable/test_bigtable_toolset.py +++ b/tests/unittests/tools/bigtable/test_bigtable_toolset.py @@ -16,12 +16,19 @@ from unittest import mock +from google.adk.agents.context import Context +from google.adk.agents.invocation_context import InvocationContext +from google.adk.sessions.session import Session from google.adk.tools.bigtable import BigtableCredentialsConfig from google.adk.tools.bigtable import metadata_tool from google.adk.tools.bigtable import query_tool +from google.adk.tools.bigtable.bigtable_toolset import BigtableParameterizedViewTool from google.adk.tools.bigtable.bigtable_toolset import BigtableToolset from google.adk.tools.bigtable.bigtable_toolset import DEFAULT_BIGTABLE_TOOL_NAME_PREFIX +from google.adk.tools.bigtable.settings import BigtableToolSettings from google.adk.tools.google_tool import GoogleTool +from google.adk.tools.tool_context import ToolContext +from google.auth.credentials import Credentials import pytest @@ -133,3 +140,201 @@ async def test_bigtable_toolset_unknown_tool(selected_tools, returned_tools): expected_tool_names = set(returned_tools) actual_tool_names = set([tool.name for tool in tools]) assert actual_tool_names == expected_tool_names + + +@pytest.mark.asyncio +async def test_bigtable_toolset_query_tool_wrapped(): + """Test that execute_sql is wrapped in BigtableQueryTool.""" + credentials_config = BigtableCredentialsConfig( + client_id="abc", client_secret="def" + ) + toolset = BigtableToolset(credentials_config=credentials_config) + + tools = await toolset.get_tools() + query_tools = [tool for tool in tools if tool.name == "execute_sql"] + assert len(query_tools) == 1 + + parameterized_tools = [ + tool for tool in tools if tool.name == "execute_sql_parameterized" + ] + assert len(parameterized_tools) == 0 + + +@pytest.mark.asyncio +async def test_bigtable_toolset_query_tool_wrapped_custom_mapping(): + """Test that BigtableParameterizedViewTool accepts custom mapping.""" + credentials_config = BigtableCredentialsConfig( + client_id="abc", client_secret="def" + ) + toolset = BigtableToolset( + credentials_config=credentials_config, + view_parameter_names=["user_id"], + ) + + tools = await toolset.get_tools() + parameterized_tools = [ + tool for tool in tools if tool.name == "execute_sql_parameterized" + ] + assert len(parameterized_tools) == 1 + tool = parameterized_tools[0] + assert isinstance(tool, BigtableParameterizedViewTool) + assert tool.view_parameter_names == ["user_id"] + + +@pytest.mark.asyncio +async def test_bigtable_parameterized_view_tool_execution(): + """Test that BigtableParameterizedViewTool maps attributes from tool_context to view_parameters.""" + + # Define a dummy function to wrap that has '_view_parameters' parameter + def mock_execute_sql(_view_parameters=None): + return {"status": "SUCCESS", "_view_parameters": _view_parameters} + + credentials = mock.create_autospec(Credentials, instance=True) + tool_settings = BigtableToolSettings() + tool_context = mock.create_autospec(ToolContext, instance=True) + tool_context.user_id = "test-user-123" + + # Create tool with custom mapping + tool = BigtableParameterizedViewTool( + func=mock_execute_sql, + view_parameter_names=["user_id"], + ) + + # Run the tool + res = await tool._run_async_with_credential( + credentials=credentials, + tool_settings=tool_settings, + args={}, + tool_context=tool_context, + ) + + assert res == { + "status": "SUCCESS", + "_view_parameters": {"user_id": "test-user-123"}, + } + + +@pytest.mark.asyncio +async def test_bigtable_parameterized_view_tool_login_flow(): + """Test showing how user_id is updated in the session/context after login.""" + + def mock_execute_sql(_view_parameters=None): + return {"status": "SUCCESS", "_view_parameters": _view_parameters} + + session = Session(id="session-1", app_name="test-app", user_id="anonymous") + + invocation_context = mock.create_autospec(InvocationContext, instance=True) + invocation_context.session = session + type(invocation_context).user_id = property(lambda self: self.session.user_id) + + tool_context = Context(invocation_context=invocation_context) + assert tool_context.user_id == "anonymous" + + # Simulate login by updating user_id on session + session.user_id = "authenticated-user-999" + + credentials = mock.create_autospec(Credentials, instance=True) + tool = BigtableParameterizedViewTool( + func=mock_execute_sql, + view_parameter_names=["user_id"], + ) + + res = await tool._run_async_with_credential( + credentials=credentials, + tool_settings=BigtableToolSettings(), + args={}, + tool_context=tool_context, + ) + + assert res == { + "status": "SUCCESS", + "_view_parameters": {"user_id": "authenticated-user-999"}, + } + + +@pytest.mark.asyncio +async def test_bigtable_parameterized_view_tool_execution_session_state_fallback(): + """Test that BigtableParameterizedViewTool falls back to tool_context.state for custom parameters.""" + + def mock_execute_sql(_view_parameters=None): + return {"status": "SUCCESS", "_view_parameters": _view_parameters} + + # Create session with application-level state + session = Session( + id="session-1", + app_name="test-app", + user_id="user-123", + state={"tenant_id": "tenant-xyz"}, + ) + + invocation_context = mock.create_autospec(InvocationContext, instance=True) + invocation_context.session = session + type(invocation_context).user_id = property(lambda self: self.session.user_id) + + tool_context = Context(invocation_context=invocation_context) + + # Ensure 'tenant_id' is NOT a top-level property or attribute on tool_context + assert not hasattr(tool_context, "tenant_id") + assert "tenant_id" in tool_context.state + + credentials = mock.create_autospec(Credentials, instance=True) + tool = BigtableParameterizedViewTool( + func=mock_execute_sql, + view_parameter_names=["tenant_id"], + ) + + res = await tool._run_async_with_credential( + credentials=credentials, + tool_settings=BigtableToolSettings(), + args={}, + tool_context=tool_context, + ) + + assert res == { + "status": "SUCCESS", + "_view_parameters": {"tenant_id": "tenant-xyz"}, + } + + +@pytest.mark.asyncio +async def test_bigtable_parameterized_view_tool_execution_multiple_parameters(): + """Test that BigtableParameterizedViewTool maps multiple attributes to view_parameters.""" + + def mock_execute_sql(_view_parameters=None): + return {"status": "SUCCESS", "_view_parameters": _view_parameters} + + session = Session( + id="session-1", + app_name="test-app", + user_id="user-123", + state={"tenant_id": "tenant-xyz", "agent_id": "agent-123"}, + ) + + invocation_context = mock.create_autospec(InvocationContext, instance=True) + invocation_context.session = session + type(invocation_context).user_id = property(lambda self: self.session.user_id) + + tool_context = Context(invocation_context=invocation_context) + + credentials = mock.create_autospec(Credentials, instance=True) + # Pass list of parameter names to be matched + tool = BigtableParameterizedViewTool( + func=mock_execute_sql, + view_parameter_names=["user_id", "tenant_id", "agent_id"], + ) + + res = await tool._run_async_with_credential( + credentials=credentials, + tool_settings=BigtableToolSettings(), + args={}, + tool_context=tool_context, + ) + + assert res == { + "status": "SUCCESS", + "_view_parameters": { + "user_id": "user-123", + "tenant_id": "tenant-xyz", + "agent_id": "agent-123", + }, + } From ffa184395136c366bf498324191d43d7afb86df5 Mon Sep 17 00:00:00 2001 From: doughayden <110487462+doughayden@users.noreply.github.com> Date: Wed, 1 Jul 2026 11:51:30 -0700 Subject: [PATCH 07/18] fix(auth): strip redirect_uri from credential_key MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merge https://github.com/google/adk-python/pull/5692 ### Link to Issue or Description of Change **1. Link to an existing issue (if applicable):** - Closes: #5691 This change adds `auth_credential.oauth2.redirect_uri = None` to the OAuth2 strip block at the three call sites where credential hashing happens. `redirect_uri` is deployment configuration (which callback URL the auth server should redirect to), not part of the credential identity, so it should be excluded from the hash just as access_token, refresh_token, expires_at, and the other transient OAuth2 fields already are. Without this change, a credential minted under one deployment URL cannot be retrieved when the deployment moves to another. ### Testing Plan **Unit Tests:** - [x] I have added or updated unit tests for my change. - [x] All unit tests pass locally. Three new tests, one per affected method. Each constructs two `AuthCredential` instances that differ only in `redirect_uri` and asserts the computed key is identical: - `tests/unittests/tools/openapi_tool/openapi_spec_parser/test_tool_auth_handler.py::test_credential_key_is_stable_across_redirect_uri` - `tests/unittests/tools/openapi_tool/openapi_spec_parser/test_tool_auth_handler.py::test_legacy_credential_key_is_stable_across_redirect_uri` - `tests/unittests/auth/test_auth_config.py::test_credential_key_is_stable_across_redirect_uri` ``` $ pytest tests/unittests/tools/openapi_tool/openapi_spec_parser/test_tool_auth_handler.py tests/unittests/auth/test_auth_config.py ======================== 14 passed, 8 warnings in 2.90s ======================== $ pytest tests/unittests/ ================ 5740 passed, 2340 warnings in 86.64s (0:01:26) ================ ``` **Manual End-to-End (E2E) Tests:** A self-contained Runner-based reproduction is at https://github.com/doughayden/adk-issue-examples/tree/main/05-redirect_uri_in_credential_hash. The agent definition (`agent.py`) wires up an `OpenAPIToolset` against a local OAuth2 test server, configured with one redirect_uri value. `main.py` constructs an `InMemoryRunner`, seeds a real (non-expired) OAuth2 credential into session state under a different redirect_uri's hash, and runs the agent. The `--apply-fix` flag monkey-patches the proposed fix to demonstrate the resolution end-to-end. Without the fix: ``` 🌤️ WeatherAssistant Agent — redirect_uri-in-hash repro ============================================================ Proposed fix applied: False STORED_REDIRECT_URI: http://localhost:8080/callback CURRENT_REDIRECT_URI: http://localhost:8081/callback Hash keys produced by ToolContextCredentialStore.get_credential_key: STORED → oauth2_55f666541ad22e39_oauth2_8ba0457897522d9d_existing_exchanged_credential CURRENT → oauth2_55f666541ad22e39_oauth2_ae16199243c358df_existing_exchanged_credential ❌ Keys differ — credentials minted under STORED are not retrievable. 👤 User: What's the weather in San Francisco? 🌤️ Weather Assistant event stream: [function_call] get_weather by WeatherAssistant [auth_event] adk_request_credential by WeatherAssistant [function_response] get_weather by WeatherAssistant [text] WeatherAssistant: 'It seems I need your authorization to access weather data. Could you please g...' Event counts: function_calls: 1 auth_events: 1 function_responses: 1 text_events: 1 ✅ Bug reproduced: agent emitted 1 adk_request_credential event(s) despite a valid seeded credential being present in state. ``` With the fix: ``` 🌤️ WeatherAssistant Agent — redirect_uri-in-hash repro ============================================================ Proposed fix applied: True STORED_REDIRECT_URI: http://localhost:8080/callback CURRENT_REDIRECT_URI: http://localhost:8081/callback Hash keys produced by ToolContextCredentialStore.get_credential_key: STORED → oauth2_55f666541ad22e39_oauth2_c2ad46dffd26cd87_existing_exchanged_credential CURRENT → oauth2_55f666541ad22e39_oauth2_c2ad46dffd26cd87_existing_exchanged_credential ✅ Keys match — fix is taking effect at the hash level. 👤 User: What's the weather in San Francisco? 🌤️ Weather Assistant event stream: [function_call] get_weather by WeatherAssistant [function_response] get_weather by WeatherAssistant [text] WeatherAssistant: 'The weather in San Francisco is Clear with a temperature of 30 degrees Celsiu...' Event counts: function_calls: 1 auth_events: 0 function_responses: 1 text_events: 1 ✅ Fix verified: tool call succeeded against the seeded credential without an adk_request_credential prompt. ``` ### Checklist - [x] I have read the [CONTRIBUTING.md](https://github.com/google/adk-python/blob/main/CONTRIBUTING.md) document. - [x] I have performed a self-review of my own code. - [x] I have commented my code, particularly in hard-to-understand areas. - [x] I have added tests that prove my fix is effective or that my feature works. - [x] New and existing unit tests pass locally with my changes. - [x] I have manually tested my changes end-to-end. - [ ] Any dependent changes have been merged and published in downstream modules. ### Additional context **Scope:** The same strip block exists at three call sites and has the same gap at all three. This PR patches all three. Patching only the tool-level pair (`tool_auth_handler.py`) and leaving the framework-level path (`auth_tool.py:AuthConfig.get_credential_key`) would leave the bug reachable for any consumer that does not work around #5327 with `get_auth_config = lambda: None`. Patching only `AuthConfig.get_credential_key` and leaving the tool-level pair would leave the bug reachable on the standard tool-level credential lookup path. **Upgrade note:** The credential_key shape changes with this fix: `redirect_uri` is no longer included in the hash. OAuth credentials cached in existing session state under the pre-fix key shape become unreachable under the new key. Users should expect a one-time re-auth prompt on the first run after upgrading. Subsequent runs use the new key normally. **Related:** - #5327 (preemptive toolset auth) - #5328 (refresh request scope) - #5329 (refreshed credential persistence, fixed in 218ea76e) - #5637 (tool-level auth termination) Co-authored-by: George Weale COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/5692 from doughayden:fix/credential-key-strip-redirect-uri 3be2a85c75ebca865467499de8bf29af3e735c31 PiperOrigin-RevId: 941217122 --- src/google/adk/auth/auth_tool.py | 18 +++--- .../openapi_spec_parser/tool_auth_handler.py | 36 ++++++----- tests/unittests/auth/test_auth_config.py | 38 ++++++++++++ .../test_tool_auth_handler.py | 62 +++++++++++++++++++ 4 files changed, 130 insertions(+), 24 deletions(-) diff --git a/src/google/adk/auth/auth_tool.py b/src/google/adk/auth/auth_tool.py index 820540ef128..dbf85870445 100644 --- a/src/google/adk/auth/auth_tool.py +++ b/src/google/adk/auth/auth_tool.py @@ -121,14 +121,16 @@ def get_credential_key(self): auth_credential.model_extra.clear() if auth_credential and auth_credential.oauth2: auth_credential = auth_credential.model_copy(deep=True) - auth_credential.oauth2.auth_uri = None - auth_credential.oauth2.state = None - auth_credential.oauth2.auth_response_uri = None - auth_credential.oauth2.auth_code = None - auth_credential.oauth2.access_token = None - auth_credential.oauth2.refresh_token = None - auth_credential.oauth2.expires_at = None - auth_credential.oauth2.expires_in = None + if auth_credential.oauth2: + auth_credential.oauth2.auth_uri = None + auth_credential.oauth2.state = None + auth_credential.oauth2.auth_response_uri = None + auth_credential.oauth2.auth_code = None + auth_credential.oauth2.access_token = None + auth_credential.oauth2.refresh_token = None + auth_credential.oauth2.expires_at = None + auth_credential.oauth2.expires_in = None + auth_credential.oauth2.redirect_uri = None credential_name = ( f"{auth_credential.auth_type.value}_{_stable_model_digest(auth_credential)}" if auth_credential diff --git a/src/google/adk/tools/openapi_tool/openapi_spec_parser/tool_auth_handler.py b/src/google/adk/tools/openapi_tool/openapi_spec_parser/tool_auth_handler.py index 0d78a5759b6..ce47b9dc46a 100644 --- a/src/google/adk/tools/openapi_tool/openapi_spec_parser/tool_auth_handler.py +++ b/src/google/adk/tools/openapi_tool/openapi_spec_parser/tool_auth_handler.py @@ -62,14 +62,16 @@ def _get_legacy_credential_key( ) -> str: if auth_credential and auth_credential.oauth2: auth_credential = auth_credential.model_copy(deep=True) - auth_credential.oauth2.auth_uri = None - auth_credential.oauth2.state = None - auth_credential.oauth2.auth_response_uri = None - auth_credential.oauth2.auth_code = None - auth_credential.oauth2.access_token = None - auth_credential.oauth2.refresh_token = None - auth_credential.oauth2.expires_at = None - auth_credential.oauth2.expires_in = None + if auth_credential.oauth2: + auth_credential.oauth2.auth_uri = None + auth_credential.oauth2.state = None + auth_credential.oauth2.auth_response_uri = None + auth_credential.oauth2.auth_code = None + auth_credential.oauth2.access_token = None + auth_credential.oauth2.refresh_token = None + auth_credential.oauth2.expires_at = None + auth_credential.oauth2.expires_in = None + auth_credential.oauth2.redirect_uri = None scheme_name = ( f"{auth_scheme.type_.name}_{self._legacy_stable_digest(auth_scheme.model_dump_json())}" if auth_scheme @@ -91,14 +93,16 @@ def get_credential_key( if auth_credential and auth_credential.oauth2: auth_credential = auth_credential.model_copy(deep=True) - auth_credential.oauth2.auth_uri = None - auth_credential.oauth2.state = None - auth_credential.oauth2.auth_response_uri = None - auth_credential.oauth2.auth_code = None - auth_credential.oauth2.access_token = None - auth_credential.oauth2.refresh_token = None - auth_credential.oauth2.expires_at = None - auth_credential.oauth2.expires_in = None + if auth_credential.oauth2: + auth_credential.oauth2.auth_uri = None + auth_credential.oauth2.state = None + auth_credential.oauth2.auth_response_uri = None + auth_credential.oauth2.auth_code = None + auth_credential.oauth2.access_token = None + auth_credential.oauth2.refresh_token = None + auth_credential.oauth2.expires_at = None + auth_credential.oauth2.expires_in = None + auth_credential.oauth2.redirect_uri = None scheme_name = ( f"{auth_scheme.type_.name}_{_stable_model_digest(auth_scheme)}" if auth_scheme diff --git a/tests/unittests/auth/test_auth_config.py b/tests/unittests/auth/test_auth_config.py index ab5f6b584c3..a4cd1c27bf7 100644 --- a/tests/unittests/auth/test_auth_config.py +++ b/tests/unittests/auth/test_auth_config.py @@ -176,3 +176,41 @@ def test_credential_key_with_custom_auth_scheme(): key = custom_config.credential_key assert key.startswith("adk_mock_custom_type_") assert len(key) > len("adk_mock_custom_type_") + + +def test_credential_key_is_stable_across_redirect_uri(oauth2_auth_scheme): + """AuthConfig.credential_key should be invariant under redirect_uri changes. + + redirect_uri is deployment configuration (which callback URL the auth + server should redirect to), not part of the credential identity. Two + AuthConfig instances built from credentials that share the same client_id, + client_secret, and scopes but differ only in redirect_uri should produce + the same credential_key. + """ + credential_local = AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth( + client_id="client", + client_secret="secret", + redirect_uri="http://localhost:8001/oauth2callback", + ), + ) + credential_deployed = AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth( + client_id="client", + client_secret="secret", + redirect_uri="https://deployed.example.com/oauth2callback", + ), + ) + + config_local = AuthConfig( + auth_scheme=oauth2_auth_scheme, + raw_auth_credential=credential_local, + ) + config_deployed = AuthConfig( + auth_scheme=oauth2_auth_scheme, + raw_auth_credential=credential_deployed, + ) + + assert config_local.credential_key == config_deployed.credential_key diff --git a/tests/unittests/tools/openapi_tool/openapi_spec_parser/test_tool_auth_handler.py b/tests/unittests/tools/openapi_tool/openapi_spec_parser/test_tool_auth_handler.py index d32fc132da7..bd561951266 100644 --- a/tests/unittests/tools/openapi_tool/openapi_spec_parser/test_tool_auth_handler.py +++ b/tests/unittests/tools/openapi_tool/openapi_spec_parser/test_tool_auth_handler.py @@ -353,3 +353,65 @@ async def test_refreshed_credential_is_persisted_to_store( assert persisted is not None assert persisted.oauth2.access_token == 'new_access_token' assert persisted.oauth2.refresh_token == 'new_refresh_token' + + +def test_credential_key_is_stable_across_redirect_uri(): + """get_credential_key should be invariant under redirect_uri changes. + + redirect_uri is deployment configuration (which callback URL the auth + server should redirect to), not part of the credential identity. Two + AuthCredential instances that share the same client_id, client_secret, + and scopes but differ only in redirect_uri should produce the same key. + """ + scheme, _ = get_mock_openid_scheme_credential() + credential_local = AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth( + client_id='client', + client_secret='secret', + redirect_uri='http://localhost:8001/oauth2callback', + ), + ) + credential_deployed = AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth( + client_id='client', + client_secret='secret', + redirect_uri='https://deployed.example.com/oauth2callback', + ), + ) + store = ToolContextCredentialStore(tool_context=create_mock_tool_context()) + + assert store.get_credential_key( + scheme, credential_local + ) == store.get_credential_key(scheme, credential_deployed) + + +def test_legacy_credential_key_is_stable_across_redirect_uri(): + """_get_legacy_credential_key should be invariant under redirect_uri changes. + + The same redirect_uri-strip behavior must apply to the legacy key path so + that already-stored credentials remain findable after the fix. + """ + scheme, _ = get_mock_openid_scheme_credential() + credential_local = AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth( + client_id='client', + client_secret='secret', + redirect_uri='http://localhost:8001/oauth2callback', + ), + ) + credential_deployed = AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth( + client_id='client', + client_secret='secret', + redirect_uri='https://deployed.example.com/oauth2callback', + ), + ) + store = ToolContextCredentialStore(tool_context=create_mock_tool_context()) + + assert store._get_legacy_credential_key( + scheme, credential_local + ) == store._get_legacy_credential_key(scheme, credential_deployed) From 48e6e504eb17d7c684e386573021f5635c620650 Mon Sep 17 00:00:00 2001 From: Xuan Yang Date: Wed, 1 Jul 2026 13:21:58 -0700 Subject: [PATCH 08/18] chore: Update check_new_py_files.sh to run in multiple VCS environments Co-authored-by: Xuan Yang PiperOrigin-RevId: 941262556 --- scripts/check_new_py_files.sh | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/scripts/check_new_py_files.sh b/scripts/check_new_py_files.sh index 1a2cf9a1509..c595dd75047 100755 --- a/scripts/check_new_py_files.sh +++ b/scripts/check_new_py_files.sh @@ -16,12 +16,24 @@ exit_code=0 -# Get list of newly added files using diff-filter=A -# Using process substitution to avoid subshell and handle spaces in filenames +get_added_files() { + if git rev-parse --is-inside-work-tree >/dev/null 2>&1; then + git diff --cached --name-only --diff-filter=A + elif jj root >/dev/null 2>&1; then + jj diff --summary 2>/dev/null | awk '/^A / {print $2}' + elif hg root >/dev/null 2>&1; then + hg status --added --no-status 2>/dev/null + elif g4 info >/dev/null 2>&1; then + g4 opened 2>/dev/null | awk '/ - add / {print $1}' | sed 's/#.*//' + elif p4 info >/dev/null 2>&1; then + p4 opened 2>/dev/null | awk '/ - add / {print $1}' | sed 's/#.*//' + fi +} + while read -r file; do # Check if file is not empty (happens if no new files) if [[ -n "$file" ]]; then - if [[ "$file" == src/google/adk/*.py ]]; then + if [[ "$file" == */google/adk/*.py ]] || [[ "$file" == google/adk/*.py ]]; then filename=$(basename "$file") if [[ ! "$filename" == _* ]]; then echo "Error: New Python file '$file' must have a '_' prefix." @@ -32,6 +44,6 @@ while read -r file; do fi fi fi -done < <(git diff --cached --name-only --diff-filter=A) +done < <(get_added_files) exit $exit_code From 4d88a5227f7624a2a07bebd7fe23ad86b1d57c10 Mon Sep 17 00:00:00 2001 From: George Weale Date: Wed, 1 Jul 2026 13:32:06 -0700 Subject: [PATCH 09/18] fix: avoid yielding a None function-response event in live mode Co-authored-by: George Weale PiperOrigin-RevId: 941267314 --- .../adk/flows/llm_flows/base_llm_flow.py | 33 ++++++++------- .../flows/llm_flows/test_base_llm_flow.py | 42 +++++++++++++++++++ 2 files changed, 61 insertions(+), 14 deletions(-) diff --git a/src/google/adk/flows/llm_flows/base_llm_flow.py b/src/google/adk/flows/llm_flows/base_llm_flow.py index 8fb5fcd2441..1a48b66c41a 100644 --- a/src/google/adk/flows/llm_flows/base_llm_flow.py +++ b/src/google/adk/flows/llm_flows/base_llm_flow.py @@ -1221,23 +1221,28 @@ async def _postprocess_live( # Handles function calls. if model_response_event.get_function_calls(): - function_response_event = await functions.handle_function_calls_live( + # handle_function_calls_live returns None when every call is deferred + # (e.g. all long-running), so guard before yielding to avoid emitting a + # None event into the live stream. + if function_response_event := await functions.handle_function_calls_live( invocation_context, model_response_event, llm_request.tools_dict - ) - # Always yield the function response event first - yield function_response_event - - # Check if this is a set_model_response function response - if json_response := _output_schema_processor.get_structured_model_response( - function_response_event ): - # Create and yield a final model response event - final_event = ( - _output_schema_processor.create_final_model_response_event( - invocation_context, json_response + # Always yield the function response event first + yield function_response_event + + # Check if this is a set_model_response function response + if json_response := ( + _output_schema_processor.get_structured_model_response( + function_response_event ) - ) - yield final_event + ): + # Create and yield a final model response event + final_event = ( + _output_schema_processor.create_final_model_response_event( + invocation_context, json_response + ) + ) + yield final_event async def _postprocess_run_processors_async( self, invocation_context: InvocationContext, llm_response: LlmResponse diff --git a/tests/unittests/flows/llm_flows/test_base_llm_flow.py b/tests/unittests/flows/llm_flows/test_base_llm_flow.py index 5dfbb78cdcf..0888326d1d6 100644 --- a/tests/unittests/flows/llm_flows/test_base_llm_flow.py +++ b/tests/unittests/flows/llm_flows/test_base_llm_flow.py @@ -1784,3 +1784,45 @@ async def test_transfer_to_sibling_from_non_llm_agent_allowed(): # Assert assert agent is not None assert agent.name == 'child2' + + +@pytest.mark.asyncio +async def test_postprocess_live_skips_none_function_response_event(): + """When every live function call defers, no None event must be yielded. + + handle_function_calls_live returns None if all calls are long-running, and + yielding that None downstream crashes the live receive loop. + """ + from google.adk.flows.llm_flows import base_llm_flow as blf + + agent = Agent(name='test_agent', model='gemini-2.0-flash') + invocation_context = await testing_utils.create_invocation_context( + agent=agent + ) + flow = BaseLlmFlowForTesting() + + fc_part = types.Part( + function_call=types.FunctionCall(name='lro', id='1', args={}) + ) + content = types.Content(role='model', parts=[fc_part]) + model_response_event = Event( + invocation_id=invocation_context.invocation_id, + author=agent.name, + content=content, + ) + llm_request = LlmRequest(model='gemini-2.0-flash') + llm_response = LlmResponse(content=content) + + with mock.patch.object( + blf.functions, + 'handle_function_calls_live', + new=AsyncMock(return_value=None), + ): + events = [ + event + async for event in flow._postprocess_live( + invocation_context, llm_request, llm_response, model_response_event + ) + ] + + assert all(event is not None for event in events) From 6aad10df0e49d1a6a2d07f8bfeea3109b44fc475 Mon Sep 17 00:00:00 2001 From: George Weale Date: Wed, 1 Jul 2026 13:53:24 -0700 Subject: [PATCH 10/18] feat: expose SKIP_THOUGHT_SIGNATURE_VALIDATOR constant Gemini thinking models require a thought_signature on generated parts, and the backend rejects replayed parts that lack one. Callers who synthesize conversation history must set b'skip_thought_signature_validator' on the fabricated part to bypass validation, and several ADK consumers hardcode that byte-string independently. Expose it once so they can depend on a single constant. Co-authored-by: George Weale PiperOrigin-RevId: 941277709 --- src/google/adk/utils/content_utils.py | 8 ++++++ tests/unittests/utils/test_content_utils.py | 32 +++++++++++++++++++++ 2 files changed, 40 insertions(+) create mode 100644 tests/unittests/utils/test_content_utils.py diff --git a/src/google/adk/utils/content_utils.py b/src/google/adk/utils/content_utils.py index 344970ec653..dcda0102a87 100644 --- a/src/google/adk/utils/content_utils.py +++ b/src/google/adk/utils/content_utils.py @@ -16,6 +16,14 @@ from google.genai import types +SKIP_THOUGHT_SIGNATURE_VALIDATOR: bytes = b'skip_thought_signature_validator' +"""Placeholder ``Part.thought_signature`` that bypasses backend validation. + +Set it on a part you synthesize yourself (a model turn or tool call/response +the model never produced) so the Gemini backend accepts the fabricated part +instead of rejecting it for a missing signature. +""" + def is_audio_part(part: types.Part) -> bool: return ( diff --git a/tests/unittests/utils/test_content_utils.py b/tests/unittests/utils/test_content_utils.py new file mode 100644 index 00000000000..cf441a178a0 --- /dev/null +++ b/tests/unittests/utils/test_content_utils.py @@ -0,0 +1,32 @@ +# 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 __future__ import annotations + +from google.adk.utils.content_utils import SKIP_THOUGHT_SIGNATURE_VALIDATOR +from google.genai import types + + +def test_skip_thought_signature_validator_wire_value(): + # The backend recognizes this exact byte string to bypass validation; + # changing it would break every replayed synthetic part. + assert SKIP_THOUGHT_SIGNATURE_VALIDATOR == b'skip_thought_signature_validator' + + +def test_skip_thought_signature_validator_assignable_to_part(): + part = types.Part( + text='injected', + thought_signature=SKIP_THOUGHT_SIGNATURE_VALIDATOR, + ) + assert part.thought_signature == SKIP_THOUGHT_SIGNATURE_VALIDATOR From e76df3d4d1991836b78c1666a8f1218683c5d6c2 Mon Sep 17 00:00:00 2001 From: Miguel Miranda Dias <7780875+pandego@users.noreply.github.com> Date: Wed, 1 Jul 2026 14:01:08 -0700 Subject: [PATCH 11/18] fix: use exchanged OAuth credential in ApplicationIntegrationToolset Merge https://github.com/google/adk-python/pull/4659 ## Summary - update `ApplicationIntegrationToolset.get_tools()` to propagate `AuthConfig.exchanged_auth_credential` to generated `IntegrationConnectorTool` instances - ensure connector tools use exchanged OAuth credentials (including runtime access token) instead of only the raw auth credential - add regression coverage for exchanged credential propagation Fixes #4553 ## Testing plan - `uv run --extra test pytest -q tests/unittests/tools/application_integration_tool/test_application_integration_toolset.py tests/unittests/tools/application_integration_tool/test_integration_connector_tool.py` ## Notes - This is a minimal fix scoped to connector-tool credential wiring only. Co-authored-by: Shangjie Chen COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/4659 from pandego:fix/4553-ait-exchanged-auth-credential 1f9c9478a35a3839c6e16f031a49211e9cd368d5 PiperOrigin-RevId: 941281800 --- .../application_integration_toolset.py | 76 +++++++++++++++---- .../test_application_integration_toolset.py | 65 ++++++++++++++++ 2 files changed, 128 insertions(+), 13 deletions(-) diff --git a/src/google/adk/tools/application_integration_tool/application_integration_toolset.py b/src/google/adk/tools/application_integration_tool/application_integration_toolset.py index b8565f6b9be..030a73aa8c5 100644 --- a/src/google/adk/tools/application_integration_tool/application_integration_toolset.py +++ b/src/google/adk/tools/application_integration_tool/application_integration_toolset.py @@ -15,6 +15,8 @@ from __future__ import annotations import logging +from typing import Any +from typing import cast from typing import List from typing import Optional from typing import Union @@ -29,6 +31,7 @@ from ...auth.auth_credential import ServiceAccountCredential from ...auth.auth_schemes import AuthScheme from ...auth.auth_tool import AuthConfig +from ..base_tool import BaseTool from ..base_toolset import BaseToolset from ..base_toolset import ToolPredicate from ..openapi_tool.auth.auth_helpers import service_account_scheme_credential @@ -43,7 +46,7 @@ # TODO: Apply a common toolset interface -class ApplicationIntegrationToolset(BaseToolset): +class ApplicationIntegrationToolset(BaseToolset): # type: ignore[misc] """ApplicationIntegrationToolset generates tools from a given Application Integration or Integration Connector resource. @@ -182,11 +185,15 @@ def __init__( "Invalid request, Either integration or (connection and" " (entity_operations or actions)) should be provided." ) - self._openapi_toolset = None - self._tools = [] + self._openapi_toolset: Optional[OpenAPIToolset] = None + self._tools: list[IntegrationConnectorTool] = [] self._parse_spec_to_toolset(spec, connection_details) - def _parse_spec_to_toolset(self, spec_dict, connection_details): + def _parse_spec_to_toolset( + self, + spec_dict: dict[str, Any], + connection_details: dict[str, Any], + ) -> None: """Parses the spec dict to OpenAPI toolset.""" if self._service_account_json: sa_credential = ServiceAccountCredential.model_validate_json( @@ -270,21 +277,64 @@ def _parse_spec_to_toolset(self, spec_dict, connection_details): ) ) + def _clone_connector_tool_with_auth_credential( + self, + tool: IntegrationConnectorTool, + auth_credential: AuthCredential, + ) -> IntegrationConnectorTool: + return IntegrationConnectorTool( + name=tool.name, + description=tool.description, + connection_name=tool._connection_name, + connection_host=tool._connection_host, + connection_service_name=tool._connection_service_name, + entity=tool._entity, + action=tool._action, + operation=tool._operation, + rest_api_tool=tool._rest_api_tool, + auth_scheme=tool._auth_scheme, + auth_credential=auth_credential, + ) + @override async def get_tools( self, readonly_context: Optional[ReadonlyContext] = None, - ) -> List[RestApiTool]: - return ( - [ - tool - for tool in self._tools - if self._is_tool_selected(tool, readonly_context) - ] - if self._openapi_toolset is None - else await self._openapi_toolset.get_tools(readonly_context) + ) -> List[BaseTool]: + if self._openapi_toolset is not None: + return cast( + List[BaseTool], + await self._openapi_toolset.get_tools(readonly_context), + ) + + exchanged_auth_credential = ( + self._auth_config.exchanged_auth_credential + if self._auth_config + else None ) + selected_tools = [ + tool + for tool in self._tools + if self._is_tool_selected(tool, readonly_context) + ] + + if not exchanged_auth_credential: + return selected_tools + + resolved_tools: List[BaseTool] = [] + for tool in selected_tools: + if isinstance(tool, IntegrationConnectorTool) and tool._auth_scheme: + resolved_tools.append( + self._clone_connector_tool_with_auth_credential( + tool, exchanged_auth_credential + ) + ) + else: + resolved_tools.append(tool) + + return resolved_tools + @override async def close(self) -> None: if self._openapi_toolset: diff --git a/tests/unittests/tools/application_integration_tool/test_application_integration_toolset.py b/tests/unittests/tools/application_integration_tool/test_application_integration_toolset.py index 787bb8ce661..5d51b29e25f 100644 --- a/tests/unittests/tools/application_integration_tool/test_application_integration_toolset.py +++ b/tests/unittests/tools/application_integration_tool/test_application_integration_toolset.py @@ -671,3 +671,68 @@ async def test_init_with_connection_with_auth_override_disabled_and_custom_auth( assert (await toolset.get_tools())[0]._operation == "EXECUTE_ACTION" assert not (await toolset.get_tools())[0]._auth_scheme assert not (await toolset.get_tools())[0]._auth_credential + + +@pytest.mark.asyncio +async def test_get_tools_uses_exchanged_auth_credential_when_available( + project, + location, + mock_integration_client, + mock_connections_client, + mock_openapi_action_spec_parser, + connection_details_auth_override_enabled, +): + connection_name = "test-connection" + actions_list = ["create"] + mock_connections_client.return_value.get_connection_details.return_value = ( + connection_details_auth_override_enabled + ) + + oauth2_data_google_cloud = { + "type": "oauth2", + "flows": { + "authorizationCode": { + "authorizationUrl": "https://test-url/o/oauth2/auth", + "tokenUrl": "https://test-url/token", + "scopes": { + "https://test-url/auth/test-scope": "test scope", + }, + } + }, + } + + oauth2_scheme = dict_to_auth_scheme(oauth2_data_google_cloud) + raw_auth_credential = AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth( + client_id="test-client-id", + client_secret="test-client-secret", + ), + ) + + toolset = ApplicationIntegrationToolset( + project, + location, + connection=connection_name, + actions=actions_list, + auth_scheme=oauth2_scheme, + auth_credential=raw_auth_credential, + ) + + exchanged_auth_credential = AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth( + client_id="test-client-id", + client_secret="test-client-secret", + access_token="exchanged-access-token", + ), + ) + toolset._auth_config.exchanged_auth_credential = exchanged_auth_credential + + original_tool = toolset._tools[0] + tools = await toolset.get_tools() + + assert len(tools) == 1 + assert tools[0] is not original_tool + assert tools[0]._auth_credential == exchanged_auth_credential + assert original_tool._auth_credential == raw_auth_credential From be23327b03ebfbdfce969e8922a0d8e10349c99d Mon Sep 17 00:00:00 2001 From: George Weale Date: Wed, 1 Jul 2026 14:15:32 -0700 Subject: [PATCH 12/18] chore: remove the issue triaging agent The agent ran on a schedule with a write-scoped GitHub token while feeding attacker-controlled issue title and body straight into its prompt, so a crafted issue could steer its labeling and owner-assignment actions. Removing the agent and its workflow removes that exposure. Co-authored-by: George Weale PiperOrigin-RevId: 941289396 --- .github/workflows/triage.yml | 68 ------- .../adk_team/adk_triaging_agent/README.md | 108 ---------- .../adk_team/adk_triaging_agent/__init__.py | 15 -- .../adk_team/adk_triaging_agent/main.py | 185 ------------------ .../adk_team/adk_triaging_agent/settings.py | 35 ---- .../adk_team/adk_triaging_agent/utils.py | 61 ------ 6 files changed, 472 deletions(-) delete mode 100644 .github/workflows/triage.yml delete mode 100644 contributing/samples/adk_team/adk_triaging_agent/README.md delete mode 100755 contributing/samples/adk_team/adk_triaging_agent/__init__.py delete mode 100644 contributing/samples/adk_team/adk_triaging_agent/main.py delete mode 100644 contributing/samples/adk_team/adk_triaging_agent/settings.py delete mode 100644 contributing/samples/adk_team/adk_triaging_agent/utils.py diff --git a/.github/workflows/triage.yml b/.github/workflows/triage.yml deleted file mode 100644 index ed0f851497f..00000000000 --- a/.github/workflows/triage.yml +++ /dev/null @@ -1,68 +0,0 @@ -# Copyright 2025 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. - -name: ADK Issue Triaging Agent - -on: - issues: - types: [opened, labeled] - schedule: - # Run every 6 hours to triage untriaged issues - - cron: '0 */6 * * *' - -jobs: - agent-triage-issues: - runs-on: ubuntu-latest - # Run for: - # - Scheduled runs (batch processing) - # - New issues (need component labeling) - # - Issues labeled with "planned" (need owner assignment) - if: >- - github.repository == 'google/adk-python' && ( - github.event_name == 'schedule' || - github.event.action == 'opened' - ) - permissions: - issues: write - contents: read - - steps: - - name: Checkout repository - uses: actions/checkout@v6 - - - name: Set up Python - uses: actions/setup-python@v6 - with: - python-version: '3.11' - - - name: Install dependencies - run: | - python -m pip install --upgrade pip - pip install requests google-adk - - - name: Run Triaging Script - env: - GITHUB_TOKEN: ${{ secrets.ADK_TRIAGE_AGENT }} - GOOGLE_API_KEY: ${{ secrets.GOOGLE_API_KEY }} - GOOGLE_GENAI_USE_VERTEXAI: 0 - OWNER: ${{ github.repository_owner }} - REPO: ${{ github.event.repository.name }} - INTERACTIVE: 0 - EVENT_NAME: ${{ github.event_name }} # 'issues', 'schedule', etc. - ISSUE_NUMBER: ${{ github.event.issue.number }} - ISSUE_TITLE: ${{ github.event.issue.title }} - ISSUE_BODY: ${{ github.event.issue.body }} - ISSUE_COUNT_TO_PROCESS: '3' # Process 3 issues at a time on schedule - PYTHONPATH: contributing/samples/adk_team - run: python -m adk_triaging_agent.main diff --git a/contributing/samples/adk_team/adk_triaging_agent/README.md b/contributing/samples/adk_team/adk_triaging_agent/README.md deleted file mode 100644 index cb0fe3a8262..00000000000 --- a/contributing/samples/adk_team/adk_triaging_agent/README.md +++ /dev/null @@ -1,108 +0,0 @@ -# ADK Issue Triaging Assistant - -The ADK Issue Triaging Assistant is a Python-based agent designed to help manage and triage GitHub issues for the `google/adk-python` repository. It uses a large language model to analyze issues, recommend appropriate component labels, set issue types, and assign owners based on predefined rules. - -This agent can be operated in two distinct modes: an interactive mode for local use or as a fully automated GitHub Actions workflow. - -______________________________________________________________________ - -## Triaging Workflow - -The agent performs different actions based on the issue state: - -| Condition | Actions | -| ------------------------------------------------- | -------------------------------------------------- | -| Issue without component label | Add component label + Set issue type (Bug/Feature) | -| Issue with "planned" label but no assignee | Assign owner based on component label | -| Issue with "planned" label AND no component label | Add component label + Set type + Assign owner | - -### Component Labels - -The agent can assign the following component labels, each mapped to an owner: - -- `a2a`, `agent engine`, `auth`, `bq`, `core`, `documentation`, `eval`, `live`, `mcp`, `models`, `services`, `tools`, `tracing`, `web`, `workflow` - -### Issue Types - -Based on the issue content, the agent will set the issue type to: - -- **Bug**: For bug reports -- **Feature**: For feature requests - -______________________________________________________________________ - -## Interactive Mode - -This mode allows you to run the agent locally to review its recommendations in real-time before any changes are made to your repository's issues. - -### Features - -- **Web Interface**: The agent's interactive mode can be rendered in a web browser using the ADK's `adk web` command. -- **User Approval**: In interactive mode, the agent is instructed to ask for your confirmation before applying labels or assigning owners. - -### Running in Interactive Mode - -To run the agent in interactive mode, first set the required environment variables. Then, execute the following command in your terminal: - -```bash -adk web -``` - -This will start a local server and provide a URL to access the agent's web interface in your browser. - -______________________________________________________________________ - -## GitHub Workflow Mode - -For automated, hands-off issue triaging, the agent can be integrated directly into your repository's CI/CD pipeline using a GitHub Actions workflow. - -### Workflow Triggers - -The GitHub workflow is configured to run on specific triggers: - -1. **New Issues (`opened`)**: When a new issue is created, the agent adds an appropriate component label and sets the issue type. - -1. **Planned Label Added (`labeled` with "planned")**: When an issue is labeled as "planned", the agent assigns an owner based on the component label. If the issue doesn't have a component label yet, the agent will also add one. - -1. **Scheduled Runs**: The workflow runs every 6 hours to process any issues that need triaging (either missing component labels or missing assignees for "planned" issues). - -### Automated Actions - -When running as part of the GitHub workflow, the agent operates non-interactively: - -- **Component Labeling**: Automatically applies the most appropriate component label -- **Issue Type Setting**: Sets the issue type to Bug or Feature based on content -- **Owner Assignment**: Only assigns owners for issues marked as "planned" - -This behavior is configured by setting the `INTERACTIVE` environment variable to `0` in the workflow file. - -### Workflow Configuration - -The workflow is defined in a YAML file (`.github/workflows/triage.yml`). This file contains the steps to check out the code, set up the Python environment, install dependencies, and run the triaging script with the necessary environment variables and secrets. - -______________________________________________________________________ - -## Setup and Configuration - -Whether running in interactive or workflow mode, the agent requires the following setup. - -### Dependencies - -The agent requires the following Python libraries. - -```bash -pip install --upgrade pip -pip install google-adk requests -``` - -### Environment Variables - -The following environment variables are required for the agent to connect to the necessary services. - -- `GITHUB_TOKEN`: **(Required)** A GitHub Personal Access Token with `issues:write` permissions. Needed for both interactive and workflow modes. -- `GOOGLE_API_KEY`: **(Required)** Your API key for the Gemini API. Needed for both interactive and workflow modes. -- `OWNER`: The GitHub organization or username that owns the repository (e.g., `google`). In the workflow, this is automatically set from the repository context. -- `REPO`: The name of the GitHub repository (e.g., `adk-python`). In the workflow, this is automatically set from the repository context. -- `INTERACTIVE`: Controls the agent's interaction mode. For the automated workflow, this is set to `0`. For interactive mode, it should be set to `1` or left unset. - -For local execution in interactive mode, you can place these variables in a `.env` file in the project's root directory. For the GitHub workflow, they should be configured as repository secrets. diff --git a/contributing/samples/adk_team/adk_triaging_agent/__init__.py b/contributing/samples/adk_team/adk_triaging_agent/__init__.py deleted file mode 100755 index 4015e47d6e4..00000000000 --- a/contributing/samples/adk_team/adk_triaging_agent/__init__.py +++ /dev/null @@ -1,15 +0,0 @@ -# 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 . import agent diff --git a/contributing/samples/adk_team/adk_triaging_agent/main.py b/contributing/samples/adk_team/adk_triaging_agent/main.py deleted file mode 100644 index fcdac832ac5..00000000000 --- a/contributing/samples/adk_team/adk_triaging_agent/main.py +++ /dev/null @@ -1,185 +0,0 @@ -# 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. - -import asyncio -import time - -from adk_triaging_agent import agent -from adk_triaging_agent.agent import LABEL_TO_OWNER -from adk_triaging_agent.settings import EVENT_NAME -from adk_triaging_agent.settings import GITHUB_BASE_URL -from adk_triaging_agent.settings import ISSUE_BODY -from adk_triaging_agent.settings import ISSUE_COUNT_TO_PROCESS -from adk_triaging_agent.settings import ISSUE_NUMBER -from adk_triaging_agent.settings import ISSUE_TITLE -from adk_triaging_agent.settings import OWNER -from adk_triaging_agent.settings import REPO -from adk_triaging_agent.utils import get_request -from adk_triaging_agent.utils import parse_number_string -from google.adk.agents.run_config import RunConfig -from google.adk.runners import InMemoryRunner -from google.adk.runners import Runner -from google.genai import types -import requests - -APP_NAME = "adk_triage_app" -USER_ID = "adk_triage_user" - - -async def fetch_specific_issue_details(issue_number: int): - """Fetches details for a single issue if it needs triaging.""" - url = f"{GITHUB_BASE_URL}/repos/{OWNER}/{REPO}/issues/{issue_number}" - print(f"Fetching details for specific issue: {url}") - - try: - issue_data = get_request(url) - labels = issue_data.get("labels", []) - label_names = {label["name"] for label in labels} - assignees = issue_data.get("assignees", []) - - # Check issue state - component_labels = set(LABEL_TO_OWNER.keys()) - has_planned = "planned" in label_names - existing_component_labels = label_names & component_labels - has_component = bool(existing_component_labels) - has_assignee = len(assignees) > 0 - - # Determine what actions are needed - needs_component_label = not has_component - needs_owner = not has_assignee - - if needs_component_label or needs_owner: - print( - f"Issue #{issue_number} needs triaging. " - f"needs_component_label={needs_component_label}, " - f"needs_owner={needs_owner}" - ) - return { - "number": issue_data["number"], - "title": issue_data["title"], - "body": issue_data.get("body", ""), - "has_planned_label": has_planned, - "has_component_label": has_component, - "existing_component_label": ( - list(existing_component_labels)[0] - if existing_component_labels - else None - ), - "needs_component_label": needs_component_label, - "needs_owner": needs_owner, - } - else: - print(f"Issue #{issue_number} is already fully triaged. Skipping.") - return None - except requests.exceptions.RequestException as e: - print(f"Error fetching issue #{issue_number}: {e}") - if hasattr(e, "response") and e.response is not None: - print(f"Response content: {e.response.text}") - return None - - -async def call_agent_async( - runner: Runner, user_id: str, session_id: str, prompt: str -) -> str: - """Call the agent asynchronously with the user's prompt.""" - content = types.Content( - role="user", parts=[types.Part.from_text(text=prompt)] - ) - - final_response_text = "" - async for event in runner.run_async( - user_id=user_id, - session_id=session_id, - new_message=content, - run_config=RunConfig(save_input_blobs_as_artifacts=False), - ): - if ( - event.content - and event.content.parts - and hasattr(event.content.parts[0], "text") - and event.content.parts[0].text - ): - print(f"** {event.author} (ADK): {event.content.parts[0].text}") - if event.author == agent.root_agent.name: - final_response_text += event.content.parts[0].text - - return final_response_text - - -async def main(): - runner = InMemoryRunner( - agent=agent.root_agent, - app_name=APP_NAME, - ) - session = await runner.session_service.create_session( - user_id=USER_ID, - app_name=APP_NAME, - ) - - if EVENT_NAME == "issues" and ISSUE_NUMBER: - print(f"EVENT: Processing specific issue due to '{EVENT_NAME}' event.") - issue_number = parse_number_string(ISSUE_NUMBER) - if not issue_number: - print(f"Error: Invalid issue number received: {ISSUE_NUMBER}.") - return - - specific_issue = await fetch_specific_issue_details(issue_number) - if specific_issue is None: - print( - f"No issue details found for #{issue_number} that needs triaging," - " or an error occurred. Skipping agent interaction." - ) - return - - issue_title = ISSUE_TITLE or specific_issue["title"] - issue_body = ISSUE_BODY or specific_issue["body"] - needs_component_label = specific_issue.get("needs_component_label", True) - needs_owner = specific_issue.get("needs_owner", False) - existing_component_label = specific_issue.get("existing_component_label") - - prompt = ( - f"Triage GitHub issue #{issue_number}.\n\n" - f'Title: "{issue_title}"\n' - f'Body: "{issue_body}"\n\n' - f"Issue state: needs_component_label={needs_component_label}, " - f"needs_owner={needs_owner}, " - f"existing_component_label={existing_component_label}" - ) - else: - print(f"EVENT: Processing batch of issues (event: {EVENT_NAME}).") - issue_count = parse_number_string(ISSUE_COUNT_TO_PROCESS, default_value=3) - prompt = ( - f"Please use 'list_untriaged_issues' to find {issue_count} issues that" - " need triaging, then triage each one according to your instructions." - ) - - response = await call_agent_async(runner, USER_ID, session.id, prompt) - print(f"<<<< Agent Final Output: {response}\n") - - -if __name__ == "__main__": - start_time = time.time() - print( - f"Start triaging {OWNER}/{REPO} issues at" - f" {time.strftime('%Y-%m-%d %H:%M:%S', time.gmtime(start_time))}" - ) - print("-" * 80) - asyncio.run(main()) - print("-" * 80) - end_time = time.time() - print( - "Triaging finished at" - f" {time.strftime('%Y-%m-%d %H:%M:%S', time.gmtime(end_time))}", - ) - print("Total script execution time:", f"{end_time - start_time:.2f} seconds") diff --git a/contributing/samples/adk_team/adk_triaging_agent/settings.py b/contributing/samples/adk_team/adk_triaging_agent/settings.py deleted file mode 100644 index fdc8b4e0338..00000000000 --- a/contributing/samples/adk_team/adk_triaging_agent/settings.py +++ /dev/null @@ -1,35 +0,0 @@ -# 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. - -import os - -from dotenv import load_dotenv - -load_dotenv(override=True) - -GITHUB_BASE_URL = "https://api.github.com" - -GITHUB_TOKEN = os.getenv("GITHUB_TOKEN") -if not GITHUB_TOKEN: - raise ValueError("GITHUB_TOKEN environment variable not set") - -OWNER = os.getenv("OWNER", "google") -REPO = os.getenv("REPO", "adk-python") -EVENT_NAME = os.getenv("EVENT_NAME") -ISSUE_NUMBER = os.getenv("ISSUE_NUMBER") -ISSUE_TITLE = os.getenv("ISSUE_TITLE") -ISSUE_BODY = os.getenv("ISSUE_BODY") -ISSUE_COUNT_TO_PROCESS = os.getenv("ISSUE_COUNT_TO_PROCESS") - -IS_INTERACTIVE = os.environ.get("INTERACTIVE", "1").lower() in ["true", "1"] diff --git a/contributing/samples/adk_team/adk_triaging_agent/utils.py b/contributing/samples/adk_team/adk_triaging_agent/utils.py deleted file mode 100644 index 8c5aa9b19dc..00000000000 --- a/contributing/samples/adk_team/adk_triaging_agent/utils.py +++ /dev/null @@ -1,61 +0,0 @@ -# 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 typing import Any - -from adk_triaging_agent.settings import GITHUB_TOKEN -import requests - -headers = { - "Authorization": f"token {GITHUB_TOKEN}", - "Accept": "application/vnd.github.v3+json", -} - - -def get_request( - url: str, params: dict[str, Any] | None = None -) -> dict[str, Any]: - if params is None: - params = {} - response = requests.get(url, headers=headers, params=params, timeout=60) - response.raise_for_status() - return response.json() - - -def post_request(url: str, payload: Any) -> dict[str, Any]: - response = requests.post(url, headers=headers, json=payload, timeout=60) - response.raise_for_status() - return response.json() - - -def patch_request(url: str, payload: Any) -> dict[str, Any]: - response = requests.patch(url, headers=headers, json=payload, timeout=60) - response.raise_for_status() - return response.json() - - -def error_response(error_message: str) -> dict[str, Any]: - return {"status": "error", "message": error_message} - - -def parse_number_string(number_str: str, default_value: int = 0) -> int: - """Parse a number from the given string.""" - try: - return int(number_str) - except ValueError: - print( - f"Warning: Invalid number string: {number_str}. Defaulting to" - f" {default_value}." - ) - return default_value From edb0fd2d8796443bd141a610725497a5f8829af3 Mon Sep 17 00:00:00 2001 From: Max Ind Date: Wed, 1 Jul 2026 14:42:47 -0700 Subject: [PATCH 13/18] refactor(telemetry): extract invocation span into _instrumentation.record_invocation Co-authored-by: Max Ind PiperOrigin-RevId: 941303834 --- src/google/adk/runners.py | 13 ++++++++++-- src/google/adk/telemetry/_instrumentation.py | 21 ++++++++++++++++++++ 2 files changed, 32 insertions(+), 2 deletions(-) diff --git a/src/google/adk/runners.py b/src/google/adk/runners.py index 55ab638085a..9802801dcad 100644 --- a/src/google/adk/runners.py +++ b/src/google/adk/runners.py @@ -56,6 +56,7 @@ from .sessions.base_session_service import BaseSessionService from .sessions.base_session_service import GetSessionConfig from .sessions.session import Session +from .telemetry import _instrumentation from .telemetry.tracing import tracer from .tools.base_toolset import BaseToolset from .utils._debug_output import print_event @@ -66,6 +67,10 @@ logger = logging.getLogger('google_adk.' + __name__) +# Silence unused warning. +# tracer is imported for backwards compatibility, to avoid breaking change in the API. +_ = tracer + def _find_active_task_isolation_scope(session) -> Optional[str]: """Walk session backwards; find the active paused task agent's scope. @@ -454,7 +459,9 @@ async def _run_node_async( Events flow through ic._event_queue via NodeRunner. """ - with tracer.start_as_current_span('invocation'): + with _instrumentation.record_invocation( + entrypoint_node=node or self.agent, conversation_id=session_id + ): # 1. Setup session = await self._get_or_create_session( user_id=user_id, session_id=session_id @@ -1040,7 +1047,9 @@ async def _run_with_trace( new_message: Optional[types.Content] = None, invocation_id: Optional[str] = None, ) -> AsyncGenerator[Event, None]: - with tracer.start_as_current_span('invocation'): + with _instrumentation.record_invocation( + entrypoint_node=self.agent, conversation_id=session_id + ): session = await self._get_or_create_session( user_id=user_id, session_id=session_id, diff --git a/src/google/adk/telemetry/_instrumentation.py b/src/google/adk/telemetry/_instrumentation.py index 2284fb839dc..1f8cc31feba 100644 --- a/src/google/adk/telemetry/_instrumentation.py +++ b/src/google/adk/telemetry/_instrumentation.py @@ -20,6 +20,7 @@ import sys import time from typing import AsyncIterator +from typing import Iterator from typing import TYPE_CHECKING from opentelemetry import trace @@ -35,6 +36,7 @@ from ..models.llm_request import LlmRequest from ..models.llm_response import LlmResponse from ..tools.base_tool import BaseTool + from ..workflow._base_node import BaseNode logger = logging.getLogger("google_adk." + __name__) @@ -69,6 +71,25 @@ def _get_elapsed_s( return time.monotonic() - fallback_start +@contextlib.contextmanager +def record_invocation( + entrypoint_node: BaseNode | None, + conversation_id: str, +) -> Iterator[None]: + """Top-level ``invocation`` span for a runner invocation. + + Args: + entrypoint_node: The runner's root agent/node. + conversation_id: Session/conversation id. + + Yields: + Nothing; the span is active for the duration of the block. + """ + del entrypoint_node, conversation_id # Unused until schema v2 lands. + with tracing.tracer.start_as_current_span("invocation"): + yield + + @dataclasses.dataclass class TelemetryContext: """Stores all telemetry related state.""" From c01e5380206b55428d4a20ea16a1fa7cfbc8e5db Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Wed, 1 Jul 2026 15:10:15 -0700 Subject: [PATCH 14/18] fix(mcp): await async MCP header providers Merge https://github.com/google/adk-python/pull/6105 ## Summary - await async/awaitable MCP header providers in async execution paths. - add unit tests verifying the async header provider functionality. Fixes #6090 Co-authored-by: Kathy Wu COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/6105 from he-yufeng:fix/await-async-header-provider 086d352afa96071228ec482aa5499fe0ee7a57b6 PiperOrigin-RevId: 941317964 --- src/google/adk/tools/mcp_tool/mcp_tool.py | 49 ++++++++++--------- src/google/adk/tools/mcp_tool/mcp_toolset.py | 13 +++-- .../unittests/tools/mcp_tool/test_mcp_tool.py | 35 +++++++++++++ .../tools/mcp_tool/test_mcp_toolset.py | 26 ++++++++++ 4 files changed, 96 insertions(+), 27 deletions(-) diff --git a/src/google/adk/tools/mcp_tool/mcp_tool.py b/src/google/adk/tools/mcp_tool/mcp_tool.py index 26e12bc11ea..4c06451ccff 100644 --- a/src/google/adk/tools/mcp_tool/mcp_tool.py +++ b/src/google/adk/tools/mcp_tool/mcp_tool.py @@ -16,16 +16,13 @@ import asyncio import base64 +from collections.abc import Awaitable import inspect import logging from typing import Any from typing import Callable -from typing import Dict -from typing import List -from typing import Optional from typing import Protocol from typing import runtime_checkable -from typing import Union import warnings from fastapi.openapi.models import APIKeyIn @@ -104,9 +101,9 @@ def __call__( self, tool_name: str, *, - callback_context: Optional[CallbackContext] = None, + callback_context: CallbackContext | None = None, **kwargs: Any, - ) -> Optional[ProgressFnT]: + ) -> ProgressFnT | None: """Create a progress callback for a specific tool. Args: @@ -139,15 +136,17 @@ def __init__( *, mcp_tool: McpBaseTool, mcp_session_manager: MCPSessionManager, - auth_scheme: Optional[AuthScheme] = None, - auth_credential: Optional[AuthCredential] = None, - require_confirmation: Union[bool, Callable[..., bool]] = False, - header_provider: Optional[ - Callable[[ReadonlyContext], Dict[str, str]] - ] = None, - progress_callback: Optional[ - Union[ProgressFnT, ProgressCallbackFactory] - ] = None, + auth_scheme: AuthScheme | None = None, + auth_credential: AuthCredential | None = None, + require_confirmation: bool | Callable[..., bool] = False, + header_provider: ( + Callable[ + [ReadonlyContext], + dict[str, str] | Awaitable[dict[str, str]], + ] + | None + ) = None, + progress_callback: ProgressFnT | ProgressCallbackFactory | None = None, ): """Initializes an McpTool. @@ -225,7 +224,7 @@ def raw_mcp_tool(self) -> McpBaseTool: return self._mcp_tool @property - def visibility(self) -> List[str]: + def visibility(self) -> list[str]: """Returns the visibility if this MCP tool meta has one.""" meta = getattr(self.raw_mcp_tool, "meta", None) if not meta or not isinstance(meta, dict): @@ -238,7 +237,7 @@ def visibility(self) -> List[str]: return [] @property - def mcp_app_resource_uri(self) -> Optional[str]: + def mcp_app_resource_uri(self) -> str | None: """Returns the MCP App UI resource URI if this tool has one. MCP Apps declare a UI resource via `meta.ui.resourceUri` in the tool @@ -379,7 +378,7 @@ async def run_async( @override async def _run_async_impl( self, *, args, tool_context: ToolContext, credential: AuthCredential - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """Runs the tool asynchronously. Args: @@ -396,8 +395,10 @@ async def _run_async_impl( dynamic_headers = self._header_provider( ReadonlyContext(tool_context._invocation_context) # pylint: disable=protected-access ) + if inspect.isawaitable(dynamic_headers): + dynamic_headers = await dynamic_headers - headers: Dict[str, str] = {} + headers: dict[str, str] = {} if auth_headers: headers.update(auth_headers) if dynamic_headers: @@ -406,7 +407,7 @@ async def _run_async_impl( # Propagate trace context in the _meta field as sprcified by MCP protocol. # See https://agentclientprotocol.com/protocol/extensibility#the-meta-field - trace_carrier: Dict[str, str] = {} + trace_carrier: dict[str, str] = {} propagate.get_global_textmap().inject(carrier=trace_carrier) meta_trace_context = trace_carrier if trace_carrier else None @@ -468,7 +469,7 @@ async def _run_async_impl( ) return result - def _detect_error_in_response(self, response: Any) -> Optional[str]: + def _detect_error_in_response(self, response: Any) -> str | None: """Telemetry hook: returns an error type if the response indicates an error.""" if isinstance(response, dict) and response.get("isError"): return "MCP_TOOL_ERROR" @@ -476,7 +477,7 @@ def _detect_error_in_response(self, response: Any) -> Optional[str]: def _resolve_progress_callback( self, tool_context: ToolContext - ) -> Optional[ProgressFnT]: + ) -> ProgressFnT | None: """Resolve the progress callback for the current invocation. If progress_callback is a ProgressCallbackFactory, call it to create @@ -510,7 +511,7 @@ def _resolve_progress_callback( async def _get_headers( self, tool_context: ToolContext, credential: AuthCredential - ) -> Optional[dict[str, str]]: + ) -> dict[str, str] | None: """Extracts authentication headers from credentials. Args: @@ -524,7 +525,7 @@ async def _get_headers( ValueError: If API key authentication is configured for non-header location. """ - headers: Optional[dict[str, str]] = None + headers: dict[str, str] | None = None if credential: if credential.oauth2: headers = {"Authorization": f"Bearer {credential.oauth2.access_token}"} diff --git a/src/google/adk/tools/mcp_tool/mcp_toolset.py b/src/google/adk/tools/mcp_tool/mcp_toolset.py index faf615c7d08..d736b00c5fa 100644 --- a/src/google/adk/tools/mcp_tool/mcp_toolset.py +++ b/src/google/adk/tools/mcp_tool/mcp_toolset.py @@ -16,6 +16,7 @@ import asyncio import base64 +import inspect import logging import sys from typing import Any @@ -107,9 +108,13 @@ def __init__( auth_scheme: Optional[AuthScheme] = None, auth_credential: Optional[AuthCredential] = None, require_confirmation: Union[bool, Callable[..., bool]] = False, - header_provider: Optional[ - Callable[[ReadonlyContext], Dict[str, str]] - ] = None, + header_provider: ( + Callable[ + [ReadonlyContext], + dict[str, str] | Awaitable[dict[str, str]], + ] + | None + ) = None, progress_callback: Optional[ Union[ProgressFnT, ProgressCallbackFactory] ] = None, @@ -293,6 +298,8 @@ async def _execute_with_session( # Add headers from header_provider if available if self._header_provider and readonly_context: provider_headers = self._header_provider(readonly_context) + if inspect.isawaitable(provider_headers): + provider_headers = await provider_headers if provider_headers: headers.update(provider_headers) diff --git a/tests/unittests/tools/mcp_tool/test_mcp_tool.py b/tests/unittests/tools/mcp_tool/test_mcp_tool.py index 8d1af6808b4..fefd04f190b 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_tool.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_tool.py @@ -953,6 +953,41 @@ async def test_run_async_impl_with_header_provider_no_auth(self): "test_tool", arguments=args, progress_callback=None, meta=None ) + @pytest.mark.asyncio + async def test_run_async_impl_with_async_header_provider_no_auth(self): + """Test running tool with an async header_provider and no authentication.""" + expected_headers = {"X-Tenant-ID": "test-tenant"} + + async def header_provider(_context): + return expected_headers + + tool = MCPTool( + mcp_tool=self.mock_mcp_tool, + mcp_session_manager=self.mock_session_manager, + header_provider=header_provider, + ) + + mcp_response = CallToolResult( + content=[TextContent(type="text", text="response text")] + ) + self.mock_session.call_tool = AsyncMock(return_value=mcp_response) + + tool_context = Mock(spec=ToolContext) + tool_context._invocation_context = Mock() + args = {"param1": "test_value"} + + result = await tool._run_async_impl( + args=args, tool_context=tool_context, credential=None + ) + + assert result == mcp_response.model_dump(exclude_none=True, mode="json") + self.mock_session_manager.create_session.assert_called_once_with( + headers=expected_headers + ) + self.mock_session.call_tool.assert_called_once_with( + "test_tool", arguments=args, progress_callback=None, meta=None + ) + @pytest.mark.asyncio async def test_run_async_impl_with_header_provider_and_oauth2(self): """Test running tool with header_provider and OAuth2 auth.""" diff --git a/tests/unittests/tools/mcp_tool/test_mcp_toolset.py b/tests/unittests/tools/mcp_tool/test_mcp_toolset.py index 20ec612e8c0..0c1700ae5f6 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_toolset.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_toolset.py @@ -303,6 +303,32 @@ async def test_get_tools_with_header_provider(self): headers=expected_headers ) + @pytest.mark.asyncio + async def test_get_tools_with_async_header_provider(self): + """Test get_tools with an async header_provider.""" + mock_tools = [MockMCPTool("tool1"), MockMCPTool("tool2")] + self.mock_session.list_tools = AsyncMock( + return_value=MockListToolsResult(mock_tools) + ) + mock_readonly_context = Mock(spec=ReadonlyContext) + expected_headers = {"X-Tenant-ID": "test-tenant"} + + async def header_provider(_context): + return expected_headers + + toolset = McpToolset( + connection_params=self.mock_stdio_params, + header_provider=header_provider, + ) + toolset._mcp_session_manager = self.mock_session_manager + + tools = await toolset.get_tools(readonly_context=mock_readonly_context) + + assert len(tools) == 2 + self.mock_session_manager.create_session.assert_called_once_with( + headers=expected_headers + ) + @pytest.mark.asyncio async def test_close_success(self): """Test successful cleanup.""" From 5735690d551ec2948e7ca6b347bb6ca0f21643e2 Mon Sep 17 00:00:00 2001 From: Shangjie Chen Date: Wed, 1 Jul 2026 15:10:50 -0700 Subject: [PATCH 15/18] refactor: Align single-turn subagent branch scoping with function_call_id Co-authored-by: Shangjie Chen PiperOrigin-RevId: 941318373 --- .../tests/weather_cupertino.json | 2 +- .../tests/weather_sanjose.json | 2 +- .../tests/weather_tokyo.json | 2 +- .../tests/gaming_1000.json | 26 +++++++++---------- .../tests/gaming_1000_large.json | 22 ++++++++-------- .../single_turn_sub_agent/tests/go.json | 14 +++++----- src/google/adk/tools/agent_tool.py | 11 +++++++- 7 files changed, 44 insertions(+), 35 deletions(-) diff --git a/contributing/samples/core/input_output_schema/tests/weather_cupertino.json b/contributing/samples/core/input_output_schema/tests/weather_cupertino.json index a036ca8fb7d..0502c9f7187 100644 --- a/contributing/samples/core/input_output_schema/tests/weather_cupertino.json +++ b/contributing/samples/core/input_output_schema/tests/weather_cupertino.json @@ -42,7 +42,7 @@ }, { "author": "weather_agent", - "branch": "weather_agent@1", + "branch": "weather_agent@fc-1", "content": { "parts": [ { diff --git a/contributing/samples/core/input_output_schema/tests/weather_sanjose.json b/contributing/samples/core/input_output_schema/tests/weather_sanjose.json index 8b76e47addb..9888e02b675 100644 --- a/contributing/samples/core/input_output_schema/tests/weather_sanjose.json +++ b/contributing/samples/core/input_output_schema/tests/weather_sanjose.json @@ -42,7 +42,7 @@ }, { "author": "weather_agent", - "branch": "weather_agent@1", + "branch": "weather_agent@fc-1", "content": { "parts": [ { diff --git a/contributing/samples/core/input_output_schema/tests/weather_tokyo.json b/contributing/samples/core/input_output_schema/tests/weather_tokyo.json index e531c244ab0..1b06901a20c 100644 --- a/contributing/samples/core/input_output_schema/tests/weather_tokyo.json +++ b/contributing/samples/core/input_output_schema/tests/weather_tokyo.json @@ -42,7 +42,7 @@ }, { "author": "weather_agent", - "branch": "weather_agent@1", + "branch": "weather_agent@fc-1", "content": { "parts": [ { diff --git a/contributing/samples/multi_agent/single_turn_sub_agent/tests/gaming_1000.json b/contributing/samples/multi_agent/single_turn_sub_agent/tests/gaming_1000.json index 9112d546dbe..c442c759f02 100644 --- a/contributing/samples/multi_agent/single_turn_sub_agent/tests/gaming_1000.json +++ b/contributing/samples/multi_agent/single_turn_sub_agent/tests/gaming_1000.json @@ -77,7 +77,7 @@ }, { "author": "phone_recommender", - "branch": "phone_recommender@1", + "branch": "phone_recommender@fc-1", "content": { "parts": [ { @@ -102,7 +102,7 @@ }, { "author": "phone_recommender", - "branch": "phone_recommender@1", + "branch": "phone_recommender@fc-1", "content": { "parts": [ { @@ -125,7 +125,7 @@ }, { "author": "phone_recommender", - "branch": "phone_recommender@1", + "branch": "phone_recommender@fc-1", "content": { "parts": [ { @@ -150,7 +150,7 @@ }, { "author": "phone_recommender", - "branch": "phone_recommender@1", + "branch": "phone_recommender@fc-1", "content": { "parts": [ { @@ -173,7 +173,7 @@ }, { "author": "phone_recommender", - "branch": "phone_recommender@1", + "branch": "phone_recommender@fc-1", "content": { "parts": [ { @@ -198,7 +198,7 @@ }, { "author": "phone_recommender", - "branch": "phone_recommender@1", + "branch": "phone_recommender@fc-1", "content": { "parts": [ { @@ -221,7 +221,7 @@ }, { "author": "phone_recommender", - "branch": "phone_recommender@1", + "branch": "phone_recommender@fc-1", "content": { "parts": [ { @@ -246,7 +246,7 @@ }, { "author": "phone_recommender", - "branch": "phone_recommender@1", + "branch": "phone_recommender@fc-1", "content": { "parts": [ { @@ -269,7 +269,7 @@ }, { "author": "phone_recommender", - "branch": "phone_recommender@1", + "branch": "phone_recommender@fc-1", "content": { "parts": [ { @@ -294,7 +294,7 @@ }, { "author": "phone_recommender", - "branch": "phone_recommender@1", + "branch": "phone_recommender@fc-1", "content": { "parts": [ { @@ -317,7 +317,7 @@ }, { "author": "phone_recommender", - "branch": "phone_recommender@1", + "branch": "phone_recommender@fc-1", "content": { "parts": [ { @@ -351,7 +351,7 @@ } }, "author": "phone_recommender", - "branch": "phone_recommender@1", + "branch": "phone_recommender@fc-1", "content": { "parts": [ { @@ -376,7 +376,7 @@ }, { "author": "phone_recommender", - "branch": "phone_recommender@1", + "branch": "phone_recommender@fc-1", "content": { "parts": [ { diff --git a/contributing/samples/multi_agent/single_turn_sub_agent/tests/gaming_1000_large.json b/contributing/samples/multi_agent/single_turn_sub_agent/tests/gaming_1000_large.json index 93be74f9816..9904bb65e1c 100644 --- a/contributing/samples/multi_agent/single_turn_sub_agent/tests/gaming_1000_large.json +++ b/contributing/samples/multi_agent/single_turn_sub_agent/tests/gaming_1000_large.json @@ -44,7 +44,7 @@ }, { "author": "phone_recommender", - "branch": "phone_recommender@1", + "branch": "phone_recommender@fc-1", "content": { "parts": [ { @@ -69,7 +69,7 @@ }, { "author": "phone_recommender", - "branch": "phone_recommender@1", + "branch": "phone_recommender@fc-1", "content": { "parts": [ { @@ -92,7 +92,7 @@ }, { "author": "phone_recommender", - "branch": "phone_recommender@1", + "branch": "phone_recommender@fc-1", "content": { "parts": [ { @@ -117,7 +117,7 @@ }, { "author": "phone_recommender", - "branch": "phone_recommender@1", + "branch": "phone_recommender@fc-1", "content": { "parts": [ { @@ -140,7 +140,7 @@ }, { "author": "phone_recommender", - "branch": "phone_recommender@1", + "branch": "phone_recommender@fc-1", "content": { "parts": [ { @@ -165,7 +165,7 @@ }, { "author": "phone_recommender", - "branch": "phone_recommender@1", + "branch": "phone_recommender@fc-1", "content": { "parts": [ { @@ -188,7 +188,7 @@ }, { "author": "phone_recommender", - "branch": "phone_recommender@1", + "branch": "phone_recommender@fc-1", "content": { "parts": [ { @@ -213,7 +213,7 @@ }, { "author": "phone_recommender", - "branch": "phone_recommender@1", + "branch": "phone_recommender@fc-1", "content": { "parts": [ { @@ -236,7 +236,7 @@ }, { "author": "phone_recommender", - "branch": "phone_recommender@1", + "branch": "phone_recommender@fc-1", "content": { "parts": [ { @@ -270,7 +270,7 @@ } }, "author": "phone_recommender", - "branch": "phone_recommender@1", + "branch": "phone_recommender@fc-1", "content": { "parts": [ { @@ -295,7 +295,7 @@ }, { "author": "phone_recommender", - "branch": "phone_recommender@1", + "branch": "phone_recommender@fc-1", "content": { "parts": [ { diff --git a/contributing/samples/multi_agent/single_turn_sub_agent/tests/go.json b/contributing/samples/multi_agent/single_turn_sub_agent/tests/go.json index 8eea2f7d703..1c69be34e97 100644 --- a/contributing/samples/multi_agent/single_turn_sub_agent/tests/go.json +++ b/contributing/samples/multi_agent/single_turn_sub_agent/tests/go.json @@ -110,7 +110,7 @@ }, { "author": "phone_recommender", - "branch": "phone_recommender@1", + "branch": "phone_recommender@fc-1", "content": { "parts": [ { @@ -144,7 +144,7 @@ }, { "author": "phone_recommender", - "branch": "phone_recommender@1", + "branch": "phone_recommender@fc-1", "content": { "parts": [ { @@ -176,7 +176,7 @@ }, { "author": "phone_recommender", - "branch": "phone_recommender@1", + "branch": "phone_recommender@fc-1", "content": { "parts": [ { @@ -219,7 +219,7 @@ }, { "author": "phone_recommender", - "branch": "phone_recommender@1", + "branch": "phone_recommender@fc-1", "content": { "parts": [ { @@ -260,7 +260,7 @@ }, { "author": "phone_recommender", - "branch": "phone_recommender@1", + "branch": "phone_recommender@fc-1", "content": { "parts": [ { @@ -294,7 +294,7 @@ } }, "author": "phone_recommender", - "branch": "phone_recommender@1", + "branch": "phone_recommender@fc-1", "content": { "parts": [ { @@ -319,7 +319,7 @@ }, { "author": "phone_recommender", - "branch": "phone_recommender@1", + "branch": "phone_recommender@fc-1", "content": { "parts": [ { diff --git a/src/google/adk/tools/agent_tool.py b/src/google/adk/tools/agent_tool.py index 8054d0ff8c1..55fe7ff6038 100644 --- a/src/google/adk/tools/agent_tool.py +++ b/src/google/adk/tools/agent_tool.py @@ -361,9 +361,18 @@ async def run_async( else: node_input = args.get('request') + # Align subagent branch scoping with node execution (Node as Tool) using function_call_id. + fc_id = tool_context.function_call_id + base_branch = tool_context.get_invocation_context().branch + segment = f'{self.agent.name}@{fc_id}' + tool_branch = f'{base_branch}.{segment}' if base_branch else segment + try: return await tool_context.run_node( - self.agent, node_input=node_input, use_sub_branch=True + self.agent, + node_input=node_input, + override_branch=tool_branch, + use_sub_branch=False, ) except Exception as e: return f'Error running sub-agent: {e}' From 99ba8ce2d976ef93864aa956a6eb0e1ec407cb03 Mon Sep 17 00:00:00 2001 From: kkj333 <7047741+kkj333@users.noreply.github.com> Date: Wed, 1 Jul 2026 15:11:01 -0700 Subject: [PATCH 16/18] fix(firestore): preserve update timestamps in list_sessions Merge https://github.com/google/adk-python/pull/5640 ### Link to Issue or Description of Change **1. Link to an existing issue (if applicable):** Closes: #5632 ### Problem `FirestoreSessionService.list_sessions()` always sets `Session.last_update_time` to `0.0`, even when Firestore documents contain `updateTime`. ### Solution Use the same timestamp conversion logic as `get_session()` in `list_sessions()`: - if `updateTime` is a `datetime`, use `.timestamp()` - otherwise, attempt `float(updateTime)` - fallback to `0.0` when conversion fails This keeps list responses consistent with `get_session()` and preserves accurate update timestamps. ### Testing Plan **Unit Tests:** - [x] I have added or updated unit tests for my change. - [x] All unit tests pass locally. Updated tests: - `test_list_sessions_with_user_id` - `test_list_sessions_without_user_id` Local pytest summary: - `tests/unittests/integrations/firestore/test_firestore_session_service.py` - `17 passed` **Manual End-to-End (E2E) Tests:** Reproduced with a mocked Firestore client where session docs include `updateTime`; before fix `last_update_time` was `0.0`, after fix it matches the Firestore value. ### Checklist - [x] I have read the [CONTRIBUTING.md](https://github.com/google/adk-python/blob/main/CONTRIBUTING.md) document. - [x] I have performed a self-review of my own code. - [x] I have commented my code, particularly in hard-to-understand areas. - [x] I have added tests that prove my fix is effective or that my feature works. - [x] New and existing unit tests pass locally with my changes. - [x] I have manually tested my changes end-to-end. - [ ] Any dependent changes have been merged and published in downstream modules. Co-authored-by: George Weale COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/5640 from kkj333:fix/firestore-list-sessions-update-time-5632 9586cc6b66f6a3ce36d093b61e8106b65b5f033b PiperOrigin-RevId: 941318502 --- .../firestore/firestore_session_service.py | 28 +++---- .../test_firestore_session_service.py | 76 +++++++++++++++++++ 2 files changed, 90 insertions(+), 14 deletions(-) diff --git a/src/google/adk/integrations/firestore/firestore_session_service.py b/src/google/adk/integrations/firestore/firestore_session_service.py index 5da4191a91c..3263e101c56 100644 --- a/src/google/adk/integrations/firestore/firestore_session_service.py +++ b/src/google/adk/integrations/firestore/firestore_session_service.py @@ -62,6 +62,18 @@ DEFAULT_USER_STATE_COLLECTION = "user_states" +def _to_last_update_time(update_time: Any) -> float: + """Converts a Firestore updateTime value to epoch seconds, or 0.0.""" + if not update_time: + return 0.0 + if isinstance(update_time, datetime): + return update_time.timestamp() + try: + return float(update_time) + except (ValueError, TypeError): + return 0.0 + + class FirestoreSessionService(BaseSessionService): # type: ignore[misc] """Session service that uses Google Cloud Firestore as the backend. @@ -322,18 +334,6 @@ async def get_session( merged_state = self._merge_state(app_state, user_state, session_state) - # Convert timestamp - update_time = data.get("updateTime") - last_update_time = 0.0 - if update_time: - if isinstance(update_time, datetime): - last_update_time = update_time.timestamp() - else: - try: - last_update_time = float(update_time) - except (ValueError, TypeError): - pass - current_revision = data.get("revision", 0) session = Session( id=session_id, @@ -341,7 +341,7 @@ async def get_session( user_id=user_id, state=merged_state, events=events, - last_update_time=last_update_time, + last_update_time=_to_last_update_time(data.get("updateTime")), ) session._storage_update_marker = str(current_revision) return session @@ -418,7 +418,7 @@ def _iter_sessions_data() -> Iterator[dict[str, Any]]: user_id=data["userId"], state=merged, events=[], - last_update_time=0.0, + last_update_time=_to_last_update_time(data.get("updateTime")), ) ) diff --git a/tests/unittests/integrations/firestore/test_firestore_session_service.py b/tests/unittests/integrations/firestore/test_firestore_session_service.py index 50b1e0b02c6..b0bd9ea71eb 100644 --- a/tests/unittests/integrations/firestore/test_firestore_session_service.py +++ b/tests/unittests/integrations/firestore/test_firestore_session_service.py @@ -14,6 +14,8 @@ from __future__ import annotations +from datetime import datetime +from datetime import timezone import json from unittest import mock @@ -387,6 +389,7 @@ async def test_list_sessions_with_user_id(mock_firestore_client): "appName": app_name, "userId": user_id, "state": {"session_key": "session_val"}, + "updateTime": 1234567890.0, } app_state_coll = mock.MagicMock() @@ -442,6 +445,77 @@ def collection_side_effect(name): assert session.state["session_key"] == "session_val" assert session.state["app:app_key"] == "app_val" assert session.state["user:user_key"] == "user_val" + assert session.last_update_time == 1234567890.0 + + +@pytest.mark.asyncio +async def test_list_sessions_preserves_datetime_update_time( + mock_firestore_client, +): + """list_sessions converts datetime updateTime to last_update_time.""" + service = FirestoreSessionService(client=mock_firestore_client) + app_name = "test_app" + user_id = "test_user" + update_time = datetime(2024, 6, 15, 10, 30, 0, tzinfo=timezone.utc) + + session_doc = mock.MagicMock() + session_doc.to_dict.return_value = { + "id": "session1", + "appName": app_name, + "userId": user_id, + "state": {}, + "updateTime": update_time, + } + + app_state_coll = mock.MagicMock() + user_state_coll = mock.MagicMock() + sessions_coll = mock.MagicMock() + + def collection_side_effect(name): + if name == service.app_state_collection: + return app_state_coll + elif name == service.user_state_collection: + return user_state_coll + elif name == service.root_collection: + return sessions_coll + return mock.MagicMock() + + mock_firestore_client.collection.side_effect = collection_side_effect + + app_doc = mock.MagicMock() + app_doc.exists = False + app_doc.to_dict.return_value = {} + app_doc_ref = mock.MagicMock() + app_state_coll.document.return_value = app_doc_ref + app_doc_ref.get = mock.AsyncMock(return_value=app_doc) + + user_doc = mock.MagicMock() + user_doc.exists = False + user_doc.to_dict.return_value = {} + user_app_doc = mock.MagicMock() + user_state_coll.document.return_value = user_app_doc + users_coll = mock.MagicMock() + user_app_doc.collection.return_value = users_coll + user_doc_ref = mock.MagicMock() + users_coll.document.return_value = user_doc_ref + user_doc_ref.get = mock.AsyncMock(return_value=user_doc) + + app_doc_in_root = mock.MagicMock() + sessions_coll.document.return_value = app_doc_in_root + users_coll = mock.MagicMock() + app_doc_in_root.collection.return_value = users_coll + user_doc_in_users = mock.MagicMock() + users_coll.document.return_value = user_doc_in_users + sessions_subcoll = mock.MagicMock() + user_doc_in_users.collection.return_value = sessions_subcoll + sessions_query = mock.MagicMock() + sessions_subcoll.where.return_value = sessions_query + sessions_query.get = mock.AsyncMock(return_value=[session_doc]) + + response = await service.list_sessions(app_name=app_name, user_id=user_id) + + assert len(response.sessions) == 1 + assert response.sessions[0].last_update_time == update_time.timestamp() @pytest.mark.asyncio @@ -455,6 +529,7 @@ async def test_list_sessions_without_user_id(mock_firestore_client): "appName": app_name, "userId": "user1", "state": {"session_key": "session_val"}, + "updateTime": 1234567890.0, } mock_firestore_client.collection_group.return_value.where.return_value.get = ( @@ -503,6 +578,7 @@ async def mock_get_all(refs): assert session.id == "session1" assert session.state["app:app_key"] == "app_val" assert session.state["user:user_key"] == "user_val" + assert session.last_update_time == 1234567890.0 mock_firestore_client.collection_group.assert_called_once_with("sessions") mock_firestore_client.collection_group.return_value.where.assert_called_once_with( From ffe41f050c147d5bff91b065b61df14c471b3bd6 Mon Sep 17 00:00:00 2001 From: George Weale Date: Wed, 1 Jul 2026 15:12:24 -0700 Subject: [PATCH 17/18] fix: use mTLS endpoint for Google OAuth2 token requests When a client certificate is configured, OAuth2 token exchange and refresh now present the certificate and target the *.mtls.googleapis.com endpoint so Context-Aware Access / token binding is honored. Gated to *.googleapis.com token endpoints with a client cert available; third-party providers and non-cert environments are unchanged. Co-authored-by: George Weale PiperOrigin-RevId: 941319195 --- src/google/adk/auth/oauth2_credential_util.py | 31 ++++-- src/google/adk/utils/_mtls_utils.py | 101 ++++++++++++++++-- .../auth/test_oauth2_credential_util.py | 99 +++++++++++++++++ tests/unittests/utils/test_mtls_utils.py | 99 +++++++++++++++++ 4 files changed, 311 insertions(+), 19 deletions(-) diff --git a/src/google/adk/auth/oauth2_credential_util.py b/src/google/adk/auth/oauth2_credential_util.py index 978d645a059..0123b804b6c 100644 --- a/src/google/adk/auth/oauth2_credential_util.py +++ b/src/google/adk/auth/oauth2_credential_util.py @@ -22,6 +22,7 @@ from authlib.oauth2.rfc6749 import OAuth2Token from fastapi.openapi.models import OAuth2 +from ..utils import _mtls_utils from ..utils.feature_decorator import experimental from .auth_credential import AuthCredential from .auth_schemes import AuthScheme @@ -83,18 +84,28 @@ def create_oauth2_session( # Scope is intentionally omitted: token exchange and refresh don't require # it per RFC 6749, and some providers reject it on these requests. - return ( - OAuth2Session( - auth_credential.oauth2.client_id, - auth_credential.oauth2.client_secret, - redirect_uri=auth_credential.oauth2.redirect_uri, - state=auth_credential.oauth2.state, - token_endpoint_auth_method=auth_credential.oauth2.token_endpoint_auth_method, - code_challenge_method=auth_credential.oauth2.code_challenge_method, - ), - token_endpoint, + session = OAuth2Session( + auth_credential.oauth2.client_id, + auth_credential.oauth2.client_secret, + redirect_uri=auth_credential.oauth2.redirect_uri, + state=auth_credential.oauth2.state, + token_endpoint_auth_method=auth_credential.oauth2.token_endpoint_auth_method, + code_challenge_method=auth_credential.oauth2.code_challenge_method, ) + # When a client certificate is configured, route Google token requests through + # the mTLS endpoint and present the cert so Context-Aware Access / token + # binding is honored. Non-Google providers and non-cert environments keep the + # existing behavior. + if ( + _mtls_utils.is_non_mtls_googleapis_endpoint(token_endpoint) + and _mtls_utils.use_client_cert_effective() + ): + if _mtls_utils.configure_session_for_mtls(session): + token_endpoint = _mtls_utils.effective_googleapis_endpoint(token_endpoint) + + return session, token_endpoint + @experimental def update_credential_with_tokens( diff --git a/src/google/adk/utils/_mtls_utils.py b/src/google/adk/utils/_mtls_utils.py index 9ba51462922..58848627cd6 100644 --- a/src/google/adk/utils/_mtls_utils.py +++ b/src/google/adk/utils/_mtls_utils.py @@ -17,10 +17,22 @@ from __future__ import annotations import enum +import logging import os +from typing import TYPE_CHECKING +from urllib.parse import urlsplit +from urllib.parse import urlunsplit from google.auth.transport import mtls +if TYPE_CHECKING: + import requests + +logger = logging.getLogger("google_adk." + __name__) + +_GOOGLEAPIS_SUFFIX = ".googleapis.com" +_MTLS_GOOGLEAPIS_SUFFIX = ".mtls.googleapis.com" + class MtlsEndpoint(enum.Enum): """Enum for the mTLS endpoint setting.""" @@ -30,10 +42,21 @@ class MtlsEndpoint(enum.Enum): NEVER = "never" +def _mtls_endpoint_setting() -> MtlsEndpoint: + """Returns the GOOGLE_API_USE_MTLS_ENDPOINT setting, defaulting to AUTO.""" + setting = os.getenv( + "GOOGLE_API_USE_MTLS_ENDPOINT", MtlsEndpoint.AUTO.value + ).lower() + try: + return MtlsEndpoint(setting) + except ValueError: + return MtlsEndpoint.AUTO + + def use_client_cert_effective() -> bool: """Returns whether client certificate should be used for mTLS.""" try: - return mtls.should_use_client_cert() + return bool(mtls.should_use_client_cert()) except (ImportError, AttributeError): return ( os.getenv("GOOGLE_API_USE_CLIENT_CERTIFICATE", "false").lower() @@ -53,16 +76,76 @@ def get_api_endpoint( mtls_template: Template for mTLS regional endpoint (e.g. "secretmanager.{location}.rep.mtls.googleapis.com"). """ - use_mtls_endpoint_str = os.getenv( - "GOOGLE_API_USE_MTLS_ENDPOINT", MtlsEndpoint.AUTO.value - ).lower() - try: - use_mtls_endpoint = MtlsEndpoint(use_mtls_endpoint_str) - except ValueError: - use_mtls_endpoint = MtlsEndpoint.AUTO - + use_mtls_endpoint = _mtls_endpoint_setting() if (use_mtls_endpoint == MtlsEndpoint.ALWAYS) or ( use_mtls_endpoint == MtlsEndpoint.AUTO and use_client_cert_effective() ): return mtls_template.format(location=location) return default_template.format(location=location) + + +def is_non_mtls_googleapis_endpoint(url: str) -> bool: + """Returns whether url points at a *.googleapis.com host without the mTLS infix.""" + if not url: + return False + host = urlsplit(url).hostname or "" + return ( + host.endswith(_GOOGLEAPIS_SUFFIX) and _MTLS_GOOGLEAPIS_SUFFIX not in host + ) + + +def effective_googleapis_endpoint(url: str) -> str: + """Rewrites a *.googleapis.com url to its .mtls.googleapis.com variant. + + Honors GOOGLE_API_USE_MTLS_ENDPOINT=never as an opt-out. Hosts that are not + googleapis.com hosts, or are already mTLS hosts, are returned unchanged so + non-Google providers are never affected. + """ + if not is_non_mtls_googleapis_endpoint(url): + return url + if _mtls_endpoint_setting() == MtlsEndpoint.NEVER: + return url + parsed = urlsplit(url) + host = parsed.hostname or "" + new_host = host[: -len(_GOOGLEAPIS_SUFFIX)] + _MTLS_GOOGLEAPIS_SUFFIX + netloc = f"{new_host}:{parsed.port}" if parsed.port else new_host + return urlunsplit( + (parsed.scheme, netloc, parsed.path, parsed.query, parsed.fragment) + ) + + +def configure_session_for_mtls(session: requests.Session) -> bool: + """Mounts a mutual-TLS adapter on a requests session when a client cert exists. + + authlib's OAuth2Session is a requests.Session but not a google-auth + AuthorizedSession, so it lacks configure_mtls_channel(). This replicates that + method's effect: load the application-default client certificate and mount an + adapter that presents it on https connections. + + Returns True if a client certificate was found and the adapter was mounted. + """ + try: + from google.auth import exceptions as ga_exceptions + from google.auth.transport import _mtls_helper + from google.auth.transport.requests import _MutualTlsAdapter + except ImportError: + return False + + cert_source = ( + mtls.default_client_cert_source() + if mtls.has_default_client_cert_source() + else None + ) + try: + is_mtls, cert, key = _mtls_helper.get_client_cert_and_key(cert_source) + except (ImportError, ga_exceptions.GoogleAuthError) as e: + logger.warning( + "Could not load client certificate for mTLS; falling back to non-mTLS" + " token request: %s", + e, + ) + return False + + if is_mtls: + session.mount("https://", _MutualTlsAdapter(cert, key)) + return bool(is_mtls) diff --git a/tests/unittests/auth/test_oauth2_credential_util.py b/tests/unittests/auth/test_oauth2_credential_util.py index 17a52002554..87dc3af8e99 100644 --- a/tests/unittests/auth/test_oauth2_credential_util.py +++ b/tests/unittests/auth/test_oauth2_credential_util.py @@ -15,6 +15,7 @@ import time from typing import Optional from unittest.mock import Mock +from unittest.mock import patch from authlib.oauth2.rfc6749 import OAuth2Token from fastapi.openapi.models import OAuth2 @@ -151,6 +152,104 @@ def test_create_oauth2_session_missing_credentials(self): assert client is None assert token_endpoint is None + def _google_openid_scheme(self) -> OpenIdConnectWithConfig: + """OpenID Connect scheme that uses Google's OAuth2 token endpoint.""" + return OpenIdConnectWithConfig( + type_="openIdConnect", + openId_connect_url=( + "https://accounts.google.com/.well-known/openid_configuration" + ), + authorization_endpoint="https://accounts.google.com/o/oauth2/v2/auth", + token_endpoint="https://oauth2.googleapis.com/token", + scopes=["openid"], + ) + + @patch.dict("os.environ", {}, clear=True) + @patch("google.adk.utils._mtls_utils.configure_session_for_mtls") + @patch("google.adk.utils._mtls_utils.use_client_cert_effective") + def test_create_oauth2_session_google_endpoint_uses_mtls( + self, mock_use_cert, mock_configure + ): + """Google token endpoint is switched to mTLS when a cert is mounted.""" + mock_use_cert.return_value = True + mock_configure.return_value = True + credential = create_oauth2_auth_credential( + auth_type=AuthCredentialTypes.OAUTH2 + ) + + client, token_endpoint = create_oauth2_session( + self._google_openid_scheme(), credential + ) + + assert client is not None + assert token_endpoint == "https://oauth2.mtls.googleapis.com/token" + mock_configure.assert_called_once_with(client) + + @patch.dict("os.environ", {}, clear=True) + @patch("google.adk.utils._mtls_utils.configure_session_for_mtls") + @patch("google.adk.utils._mtls_utils.use_client_cert_effective") + def test_create_oauth2_session_google_endpoint_no_cert_keeps_plain( + self, mock_use_cert, mock_configure + ): + """Without a client cert the plain Google endpoint is kept.""" + mock_use_cert.return_value = False + credential = create_oauth2_auth_credential( + auth_type=AuthCredentialTypes.OAUTH2 + ) + + _, token_endpoint = create_oauth2_session( + self._google_openid_scheme(), credential + ) + + assert token_endpoint == "https://oauth2.googleapis.com/token" + mock_configure.assert_not_called() + + @patch.dict("os.environ", {}, clear=True) + @patch("google.adk.utils._mtls_utils.configure_session_for_mtls") + @patch("google.adk.utils._mtls_utils.use_client_cert_effective") + def test_create_oauth2_session_cert_unavailable_keeps_plain( + self, mock_use_cert, mock_configure + ): + """If the adapter cannot be mounted, the endpoint is not switched.""" + mock_use_cert.return_value = True + mock_configure.return_value = False + credential = create_oauth2_auth_credential( + auth_type=AuthCredentialTypes.OAUTH2 + ) + + client, token_endpoint = create_oauth2_session( + self._google_openid_scheme(), credential + ) + + assert token_endpoint == "https://oauth2.googleapis.com/token" + mock_configure.assert_called_once_with(client) + + @patch.dict("os.environ", {}, clear=True) + @patch("google.adk.utils._mtls_utils.configure_session_for_mtls") + @patch("google.adk.utils._mtls_utils.use_client_cert_effective") + def test_create_oauth2_session_non_google_endpoint_skips_mtls( + self, mock_use_cert, mock_configure + ): + """Non-Google providers are never switched to an mTLS endpoint.""" + mock_use_cert.return_value = True + credential = create_oauth2_auth_credential( + auth_type=AuthCredentialTypes.OAUTH2 + ) + scheme = OpenIdConnectWithConfig( + type_="openIdConnect", + openId_connect_url=( + "https://example.com/.well-known/openid_configuration" + ), + authorization_endpoint="https://example.com/auth", + token_endpoint="https://example.com/token", + scopes=["openid"], + ) + + _, token_endpoint = create_oauth2_session(scheme, credential) + + assert token_endpoint == "https://example.com/token" + mock_configure.assert_not_called() + @pytest.mark.parametrize( "token_endpoint_auth_method, expected_auth_method", [ diff --git a/tests/unittests/utils/test_mtls_utils.py b/tests/unittests/utils/test_mtls_utils.py index cfe381edd69..9a2572a437c 100644 --- a/tests/unittests/utils/test_mtls_utils.py +++ b/tests/unittests/utils/test_mtls_utils.py @@ -19,6 +19,7 @@ from unittest.mock import patch from google.adk.utils import _mtls_utils +from google.auth import exceptions as ga_exceptions import pytest _DEFAULT_TEMPLATE = "service.{location}.rep.googleapis.com" @@ -118,3 +119,101 @@ def test_get_api_endpoint_invalid_fallback_to_auto( ) assert endpoint == _MTLS_TEMPLATE.format(location=_LOCATION) mock_use_client_cert.assert_called_once() + + @pytest.mark.parametrize( + "url, expected", + [ + ("https://oauth2.googleapis.com/token", True), + ("https://openidconnect.googleapis.com/v1/userinfo", True), + ("https://oauth2.mtls.googleapis.com/token", False), + ("https://example.com/token", False), + ("https://accounts.google.com/o/oauth2/v2/auth", False), + (None, False), + ("", False), + ], + ) + def test_is_non_mtls_googleapis_endpoint(self, url, expected): + assert _mtls_utils.is_non_mtls_googleapis_endpoint(url) is expected + + @patch.dict("os.environ", {}, clear=True) + @pytest.mark.parametrize( + "url, expected", + [ + ( + "https://oauth2.googleapis.com/token", + "https://oauth2.mtls.googleapis.com/token", + ), + ( + "https://openidconnect.googleapis.com/v1/userinfo", + "https://openidconnect.mtls.googleapis.com/v1/userinfo", + ), + # Non-Google providers are never rewritten. + ("https://example.com/token", "https://example.com/token"), + # Already-mTLS hosts are left alone. + ( + "https://oauth2.mtls.googleapis.com/token", + "https://oauth2.mtls.googleapis.com/token", + ), + ], + ) + def test_effective_googleapis_endpoint_rewrites(self, url, expected): + assert _mtls_utils.effective_googleapis_endpoint(url) == expected + + @patch.dict( + "os.environ", {"GOOGLE_API_USE_MTLS_ENDPOINT": "never"}, clear=True + ) + def test_effective_googleapis_endpoint_never_opts_out(self): + url = "https://oauth2.googleapis.com/token" + assert _mtls_utils.effective_googleapis_endpoint(url) == url + + @patch.dict("os.environ", {}, clear=True) + def test_effective_googleapis_endpoint_preserves_query(self): + url = "https://iam.googleapis.com/v1/token?foo=bar" + assert ( + _mtls_utils.effective_googleapis_endpoint(url) + == "https://iam.mtls.googleapis.com/v1/token?foo=bar" + ) + + @patch("google.auth.transport.mtls.has_default_client_cert_source") + @patch("google.auth.transport._mtls_helper.get_client_cert_and_key") + @patch("google.auth.transport.requests._MutualTlsAdapter") + def test_configure_session_for_mtls_mounts_adapter( + self, mock_adapter, mock_get_cert, mock_has_cert_source + ): + mock_has_cert_source.return_value = False + mock_get_cert.return_value = (True, b"cert", b"key") + session = MagicMock() + + result = _mtls_utils.configure_session_for_mtls(session) + + assert result is True + mock_adapter.assert_called_once_with(b"cert", b"key") + session.mount.assert_called_once_with("https://", mock_adapter.return_value) + + @patch("google.auth.transport.mtls.has_default_client_cert_source") + @patch("google.auth.transport._mtls_helper.get_client_cert_and_key") + def test_configure_session_for_mtls_no_cert( + self, mock_get_cert, mock_has_cert_source + ): + mock_has_cert_source.return_value = False + mock_get_cert.return_value = (False, None, None) + session = MagicMock() + + result = _mtls_utils.configure_session_for_mtls(session) + + assert result is False + session.mount.assert_not_called() + + @patch("google.auth.transport.mtls.has_default_client_cert_source") + @patch("google.auth.transport._mtls_helper.get_client_cert_and_key") + def test_configure_session_for_mtls_cert_error_falls_back( + self, mock_get_cert, mock_has_cert_source + ): + mock_has_cert_source.return_value = False + mock_get_cert.side_effect = ga_exceptions.ClientCertError("boom") + session = MagicMock() + + result = _mtls_utils.configure_session_for_mtls(session) + + assert result is False + session.mount.assert_not_called() From 391101050eafa6f7a1e1dfc533970eb17bf99672 Mon Sep 17 00:00:00 2001 From: George Weale Date: Wed, 1 Jul 2026 15:16:01 -0700 Subject: [PATCH 18/18] chore: re-enable and fix the mypy delta presubmit Scope the delta check to the files the CL changed via the kokoro presubmit_request, and decouple the workspace copy so mypy module resolution is stable. Co-authored-by: George Weale PiperOrigin-RevId: 941321034 --- pyproject.toml | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 3daff0eba28..0fd1049383a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -292,8 +292,11 @@ single_line_exclusions = [] known_third_party = [ "a2a", "google.adk" ] [tool.mypy] +mypy_path = [ "src" ] exclude = [ "contributing/samples/", "tests/" ] -follow_imports = "skip" +namespace_packages = true +explicit_package_bases = true +follow_imports = "normal" python_version = "3.11" disable_error_code = [ "import-not-found", "import-untyped", "unused-ignore" ] strict = true