From 115c7839f42597dbe69e42a14efe607691f8580d Mon Sep 17 00:00:00 2001 From: DABH Date: Mon, 27 Jul 2026 16:05:47 -0500 Subject: [PATCH] fix: route workflow IDs and retry jitter through platform seams MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ADK routes event IDs, invocation IDs, timestamps, and client function-call IDs through the injectable google.adk.platform providers so that embedders can customize how generated values are produced — for example to make agent runs deterministic and replayable. Three call sites on the workflow HITL/retry path still reach for the stdlib directly, which bypasses any installed provider: - RequestInput.interrupt_id defaults to uuid.uuid4(). Recorded user responses reference this ID, so an ID that regenerates differently on a replayed run breaks resume matching. - _ToolNode mints ToolContext.function_call_id with uuid.uuid4(). - Retry backoff jitter draws from the global random module (jitter defaults to 1.0 for any RetryConfig), making computed delays unreproducible. Add a google.adk.platform random seam (get_random / set_random_provider / reset_random_provider, exported via google.adk.platform per the visibility convention) mirroring the existing time and uuid providers, and route the three call sites through the platform seams. Default behavior is unchanged. The two existing jitter tests patched random.uniform globally; they now inject a mock through set_random_provider, exercising the new seam. --- src/google/adk/events/request_input.py | 4 +- src/google/adk/platform/__init__.py | 10 ++++ src/google/adk/platform/_random.py | 46 +++++++++++++++ src/google/adk/workflow/_tool_node.py | 4 +- src/google/adk/workflow/utils/_retry_utils.py | 7 ++- tests/unittests/events/test_request_input.py | 57 +++++++++++++++++++ tests/unittests/platform/test_random.py | 51 +++++++++++++++++ .../workflow/test_node_runner_failure.py | 32 ++++++----- tests/unittests/workflow/test_tool_node.py | 48 ++++++++++++++++ .../workflow/test_workflow_failures.py | 18 +++--- .../workflow/utils/test_retry_utils.py | 21 +++++++ 11 files changed, 270 insertions(+), 28 deletions(-) create mode 100644 src/google/adk/platform/_random.py create mode 100644 tests/unittests/events/test_request_input.py create mode 100644 tests/unittests/platform/test_random.py diff --git a/src/google/adk/events/request_input.py b/src/google/adk/events/request_input.py index 1c42fe00402..f3311ec9c77 100644 --- a/src/google/adk/events/request_input.py +++ b/src/google/adk/events/request_input.py @@ -15,13 +15,13 @@ from typing import Any from typing import Optional -import uuid from pydantic import alias_generators from pydantic import BaseModel from pydantic import ConfigDict from pydantic import Field +from ..platform import uuid as platform_uuid from ..utils._schema_utils import SchemaType @@ -39,7 +39,7 @@ class RequestInput(BaseModel): "The ID of the interrupt, usually a function call ID. This is used" " to identify the interrupt that the input is for." ), - default_factory=lambda: str(uuid.uuid4()), + default_factory=platform_uuid.new_uuid, ) """The ID of the interrupt, usually a function call ID. diff --git a/src/google/adk/platform/__init__.py b/src/google/adk/platform/__init__.py index 58d482ea386..09d8ea84b44 100644 --- a/src/google/adk/platform/__init__.py +++ b/src/google/adk/platform/__init__.py @@ -11,3 +11,13 @@ # 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 ._random import get_random +from ._random import reset_random_provider +from ._random import set_random_provider + +__all__ = [ + 'get_random', + 'reset_random_provider', + 'set_random_provider', +] diff --git a/src/google/adk/platform/_random.py b/src/google/adk/platform/_random.py new file mode 100644 index 00000000000..9c2462d7c2d --- /dev/null +++ b/src/google/adk/platform/_random.py @@ -0,0 +1,46 @@ +# 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. + +"""Platform module for abstracting random number generation.""" + +from __future__ import annotations + +from contextvars import ContextVar +import random +from typing import Callable + +_default_random: random.Random = random.Random() +_default_random_provider: Callable[[], random.Random] = lambda: _default_random +_random_provider_context_var: ContextVar[Callable[[], random.Random]] = ( + ContextVar("random_provider", default=_default_random_provider) +) + + +def set_random_provider(provider: Callable[[], random.Random]) -> None: + """Sets the provider for the random number generator. + + Args: + provider: A callable that returns the `random.Random` instance to use. + """ + _random_provider_context_var.set(provider) + + +def reset_random_provider() -> None: + """Resets the random provider to its default implementation.""" + _random_provider_context_var.set(_default_random_provider) + + +def get_random() -> random.Random: + """Returns the random number generator.""" + return _random_provider_context_var.get()() diff --git a/src/google/adk/workflow/_tool_node.py b/src/google/adk/workflow/_tool_node.py index 68626e51e58..c77056e1a6d 100644 --- a/src/google/adk/workflow/_tool_node.py +++ b/src/google/adk/workflow/_tool_node.py @@ -19,7 +19,6 @@ from collections.abc import AsyncGenerator import json from typing import Any -import uuid from google.genai import types from pydantic import ConfigDict @@ -28,6 +27,7 @@ from ..agents.context import Context from ..events.event import Event +from ..platform import uuid as platform_uuid from ..tools.base_tool import BaseTool from ..tools.tool_context import ToolContext from ..utils.content_utils import extract_text_from_content @@ -66,7 +66,7 @@ async def _run_impl( ) -> AsyncGenerator[Any, None]: tool_context = ToolContext( invocation_context=ctx.get_invocation_context(), - function_call_id=str(uuid.uuid4()), + function_call_id=platform_uuid.new_uuid(), ) args = node_input diff --git a/src/google/adk/workflow/utils/_retry_utils.py b/src/google/adk/workflow/utils/_retry_utils.py index d9c0c0870b6..05f932daad4 100644 --- a/src/google/adk/workflow/utils/_retry_utils.py +++ b/src/google/adk/workflow/utils/_retry_utils.py @@ -16,8 +16,7 @@ """Utility functions for retrying nodes in a workflow.""" -import random - +from ...platform import _random as platform_random from .._node_state import NodeState from .._retry_config import RetryConfig @@ -82,7 +81,9 @@ def _get_retry_delay( delay = min(delay, max_delay) if jitter > 0.0: - random_offset = random.uniform(-jitter * delay, jitter * delay) + random_offset = platform_random.get_random().uniform( + -jitter * delay, jitter * delay + ) delay = max(0.0, delay + random_offset) return delay diff --git a/tests/unittests/events/test_request_input.py b/tests/unittests/events/test_request_input.py new file mode 100644 index 00000000000..819bc5729c2 --- /dev/null +++ b/tests/unittests/events/test_request_input.py @@ -0,0 +1,57 @@ +# 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. + +"""Unit tests for the RequestInput event model.""" + +from __future__ import annotations + +import uuid + +from google.adk.events.request_input import RequestInput +from google.adk.platform import uuid as platform_uuid + + +class TestRequestInputInterruptId: + + def teardown_method(self) -> None: + platform_uuid.reset_id_provider() + + def test_default_interrupt_id_is_a_uuid(self): + """Without a custom provider, the default interrupt_id is a uuid.""" + request = RequestInput() + + # Should be parseable as uuid + uuid.UUID(request.interrupt_id) + + def test_default_interrupt_id_uses_platform_id_provider(self): + """The default interrupt_id is minted via the platform uuid seam. + + Embedders can install a custom id provider to control how IDs are + generated (e.g. seeded or sequential IDs for reproducible runs); the + default interrupt_id must come from that provider because recorded + user responses reference it by value. + """ + platform_uuid.set_id_provider(lambda: "deterministic-id") + + request = RequestInput() + + assert request.interrupt_id == "deterministic-id" + + def test_explicit_interrupt_id_is_preserved(self): + """An explicitly provided interrupt_id bypasses the provider.""" + platform_uuid.set_id_provider(lambda: "deterministic-id") + + request = RequestInput(interrupt_id="explicit-id") + + assert request.interrupt_id == "explicit-id" diff --git a/tests/unittests/platform/test_random.py b/tests/unittests/platform/test_random.py new file mode 100644 index 00000000000..8462a153ea7 --- /dev/null +++ b/tests/unittests/platform/test_random.py @@ -0,0 +1,51 @@ +# 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. + +"""Unit tests for the platform random module.""" + +import random +import unittest + +from google.adk import platform as adk_platform + + +class TestRandom(unittest.TestCase): + + def tearDown(self) -> None: + # Reset provider to default after each test + adk_platform.reset_random_provider() + + def test_default_random_provider(self) -> None: + # Verify it returns a random.Random instance producing values in range + rng = adk_platform.get_random() + self.assertIsInstance(rng, random.Random) + value = rng.uniform(0.0, 1.0) + self.assertGreaterEqual(value, 0.0) + self.assertLessEqual(value, 1.0) + + def test_custom_random_provider(self) -> None: + # Test override + seeded = random.Random(42) + adk_platform.set_random_provider(lambda: seeded) + self.assertIs(adk_platform.get_random(), seeded) + expected = random.Random(42).uniform(0.0, 1.0) + self.assertEqual(adk_platform.get_random().uniform(0.0, 1.0), expected) + + def test_reset_random_provider(self) -> None: + seeded = random.Random(42) + adk_platform.set_random_provider(lambda: seeded) + adk_platform.reset_random_provider() + self.assertIsNot(adk_platform.get_random(), seeded) + # Default provider returns a stable module-level instance + self.assertIs(adk_platform.get_random(), adk_platform.get_random()) diff --git a/tests/unittests/workflow/test_node_runner_failure.py b/tests/unittests/workflow/test_node_runner_failure.py index 39ec0acb7fc..02cefa949d2 100644 --- a/tests/unittests/workflow/test_node_runner_failure.py +++ b/tests/unittests/workflow/test_node_runner_failure.py @@ -20,6 +20,7 @@ from typing import AsyncGenerator from unittest import mock +from google.adk import platform as adk_platform from google.adk.agents.context import Context from google.adk.events.event import Event from google.adk.runners import Runner @@ -715,20 +716,23 @@ async def test_retry_applies_random_jitter(request: pytest.FixtureRequest): session = await ss.create_session(app_name=agent.name, user_id='u') msg = types.Content(parts=[types.Part(text='start')], role='user') - with ( - mock.patch('asyncio.sleep', new_callable=mock.AsyncMock) as mock_sleep, - mock.patch('random.uniform', return_value=-1.0) as mock_random, - ): - events = [] - async for event in runner.run_async( - user_id='u', session_id=session.id, new_message=msg - ): - events.append(event) - - # 4.0 + (-1.0) = 3.0 - mock_sleep.assert_any_await(3.0) - # Called with -0.5 * 4.0, 0.5 * 4.0 - mock_random.assert_called_once_with(-2.0, 2.0) + mock_random = mock.Mock() + mock_random.uniform = mock.Mock(return_value=-1.0) + adk_platform.set_random_provider(lambda: mock_random) + try: + with mock.patch('asyncio.sleep', new_callable=mock.AsyncMock) as mock_sleep: + events = [] + async for event in runner.run_async( + user_id='u', session_id=session.id, new_message=msg + ): + events.append(event) + + # 4.0 + (-1.0) = 3.0 + mock_sleep.assert_any_await(3.0) + # Called with -0.5 * 4.0, 0.5 * 4.0 + mock_random.uniform.assert_called_once_with(-2.0, 2.0) + finally: + adk_platform.reset_random_provider() results = simplify_events_with_node(events) filtered_results = [ diff --git a/tests/unittests/workflow/test_tool_node.py b/tests/unittests/workflow/test_tool_node.py index 4a49cd39914..6ebed7d8d91 100644 --- a/tests/unittests/workflow/test_tool_node.py +++ b/tests/unittests/workflow/test_tool_node.py @@ -14,9 +14,12 @@ """Tests for ToolNode input parsing and execution.""" +import itertools +import re from typing import Any from google.adk.events.event import Event +from google.adk.platform import uuid as platform_uuid from google.adk.tools.base_tool import BaseTool from google.adk.workflow import START from google.adk.workflow._tool_node import _ToolNode as ToolNode @@ -139,3 +142,48 @@ async def test_tool_node_rejects_non_dict_content(): TypeError, match="The input to ToolNode must be a dictionary" ): await _run_tool_node_wf(content) + + +@pytest.mark.asyncio +async def test_tool_node_function_call_id_uses_platform_id_provider(): + """Tests that the tool's function_call_id is minted via the platform seam. + + Embedders can install a custom id provider to control how IDs are + generated; the function_call_id handed to the tool must come from that + provider, matching how flows/llm_flows/functions.py mints the same IDs. + """ + captured_ids: list[str] = [] + + class CapturingTool(BaseTool): + """A tool that records the function_call_id it was invoked with.""" + + def __init__(self): + super().__init__(name="capturing_tool", description="Captures ids") + + async def run_async(self, *, args: dict[str, Any], tool_context) -> Any: + captured_ids.append(tool_context.function_call_id) + return {} + + tool_node = ToolNode(tool=CapturingTool()) + + def start_node(): + return Event(output={"param_a": 1}) + + wf = Workflow( + name="tool_node_id_wf", + edges=[ + (START, start_node), + (start_node, tool_node), + ], + ) + counter = itertools.count() + platform_uuid.set_id_provider(lambda: f"fixed-{next(counter)}") + try: + app_instance = testing_utils.App(name="test_app", root_agent=wf) + runner = testing_utils.InMemoryRunner(app=app_instance) + await runner.run_async("start") + finally: + platform_uuid.reset_id_provider() + + assert len(captured_ids) == 1 + assert re.fullmatch(r"fixed-\d+", captured_ids[0]) diff --git a/tests/unittests/workflow/test_workflow_failures.py b/tests/unittests/workflow/test_workflow_failures.py index 8eaccdff81a..526f78a8621 100644 --- a/tests/unittests/workflow/test_workflow_failures.py +++ b/tests/unittests/workflow/test_workflow_failures.py @@ -19,6 +19,7 @@ from typing import AsyncGenerator from unittest import mock +from google.adk import platform as adk_platform from google.adk.agents.context import Context from google.adk.apps.app import App from google.adk.events.event import Event @@ -586,13 +587,16 @@ async def test_retry_with_jitter(request: pytest.FixtureRequest): app = App(name=request.function.__name__, root_agent=agent) runner = testing_utils.InMemoryRunner(app=app) - with ( - mock.patch('asyncio.sleep', new_callable=mock.AsyncMock) as mock_sleep, - mock.patch('random.uniform', return_value=-1.0) as mock_random, - ): - events = await runner.run_async(testing_utils.get_user_content('start')) - mock_sleep.assert_any_await(3.0) - mock_random.assert_called_once_with(-2.0, 2.0) + mock_random = mock.Mock() + mock_random.uniform = mock.Mock(return_value=-1.0) + adk_platform.set_random_provider(lambda: mock_random) + try: + with mock.patch('asyncio.sleep', new_callable=mock.AsyncMock) as mock_sleep: + events = await runner.run_async(testing_utils.get_user_content('start')) + mock_sleep.assert_any_await(3.0) + mock_random.uniform.assert_called_once_with(-2.0, 2.0) + finally: + adk_platform.reset_random_provider() assert simplify_events_with_node(events) == [ ( diff --git a/tests/unittests/workflow/utils/test_retry_utils.py b/tests/unittests/workflow/utils/test_retry_utils.py index db007c27d2c..c2c19621f77 100644 --- a/tests/unittests/workflow/utils/test_retry_utils.py +++ b/tests/unittests/workflow/utils/test_retry_utils.py @@ -14,6 +14,9 @@ from __future__ import annotations +import random + +from google.adk import platform as adk_platform from google.adk.workflow._node_state import NodeState from google.adk.workflow._retry_config import RetryConfig from google.adk.workflow.utils._retry_utils import _get_retry_delay @@ -70,6 +73,24 @@ def test_adds_jitter_when_enabled(self): assert all(5.0 <= d <= 15.0 for d in delays) assert len(set(delays)) > 1 + def test_jitter_uses_platform_random_provider(self): + """Jitter is drawn via the platform random seam so it is injectable. + + Embedders can install a custom random provider (e.g. a seeded Random); + the computed delay must then be reproducible from run to run. + """ + config = RetryConfig(initial_delay=10.0, backoff_factor=1.0, jitter=0.5) + state = NodeState(attempt_count=1) + adk_platform.set_random_provider(lambda: random.Random(42)) + try: + expected_offset = random.Random(42).uniform(-5.0, 5.0) + + delays = {_get_retry_delay(config, state) for _ in range(5)} + + assert delays == {max(0.0, 10.0 + expected_offset)} + finally: + adk_platform.reset_random_provider() + class TestShouldRetryNode: