fix: route workflow IDs and retry jitter through platform seams

Merge https://github.com/google/adk-python/pull/6468

PiperOrigin-RevId: 966214760
This commit is contained in:
DABH
2026-08-17 15:28:19 -07:00
committed by Copybara-Service
parent bf2efd71d6
commit 8f85107cca
11 changed files with 278 additions and 29 deletions
+2 -2
View File
@@ -15,13 +15,13 @@ from __future__ import annotations
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.
+10
View File
@@ -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',
]
+51
View File
@@ -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.
"""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. Note
that the provider callable is evaluated on every `get_random()` call; to
preserve RNG sequence state across calls, return an existing
`random.Random` instance rather than constructing a new one inside the
callable (e.g., `rng = random.Random(42); set_random_provider(lambda:
rng)` instead of `set_random_provider(lambda: random.Random(42))`).
"""
_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()()
+2 -2
View File
@@ -19,7 +19,6 @@ from __future__ import annotations
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 typing_extensions import override
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 @@ class _ToolNode(BaseNode):
) -> 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
@@ -16,8 +16,7 @@ from __future__ import annotations
"""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
@@ -86,7 +85,9 @@ def _get_retry_delay(
# bound but collapse every overshooting draw onto exactly max_delay,
# firing the retries jitter exists to spread out at the same instant.
delay = min(delay, max_delay / (1.0 + jitter))
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 min(delay, max_delay)
@@ -0,0 +1,56 @@
# 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.
Frameworks that replay agent workflows (e.g. durable execution engines)
install a deterministic id provider; the generated interrupt_id must be
stable across replays because recorded user responses reference it.
"""
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"
+51
View File
@@ -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())
@@ -20,6 +20,7 @@ from typing import Any
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
@@ -29,8 +30,8 @@ from google.adk.workflow import Edge
from google.adk.workflow import START
from google.adk.workflow._errors import NodeTimeoutError
from google.adk.workflow._graph import Graph
from google.adk.workflow._node import node
from google.adk.workflow._node import Node
from google.adk.workflow._node import node
from google.adk.workflow._node_status import NodeStatus
from google.adk.workflow._retry_config import RetryConfig
from google.adk.workflow._workflow import Workflow
@@ -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)
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.assert_called_once_with(-2.0, 2.0)
# 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 = [
@@ -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.
Frameworks that replay agent workflows (e.g. durable execution engines)
install a deterministic id provider; the generated function_call_id must be
stable across replays.
"""
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])
@@ -19,6 +19,7 @@ from typing import Any
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
@@ -29,8 +30,8 @@ from google.adk.workflow import BaseNode
from google.adk.workflow import Edge
from google.adk.workflow import START
from google.adk.workflow._graph import Graph
from google.adk.workflow._node import node
from google.adk.workflow._node import Node
from google.adk.workflow._node import node
from google.adk.workflow._node_status import NodeStatus
from google.adk.workflow._retry_config import RetryConfig
from google.adk.workflow._workflow import Workflow
@@ -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) == [
(
@@ -16,6 +16,7 @@ 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
@@ -92,6 +93,29 @@ class TestGetRetryDelay:
assert at_cap / len(delays) < 0.01
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.
Frameworks that replay agent workflows (e.g. durable execution engines)
install a deterministic random provider; the computed delay must then be
reproducible across replays.
"""
config = RetryConfig(initial_delay=10.0, backoff_factor=1.0, jitter=0.5)
state = NodeState(attempt_count=1)
rng = random.Random(42)
adk_platform.set_random_provider(lambda: rng)
try:
expected_rng = random.Random(42)
expected_delays = [
max(0.0, 10.0 + expected_rng.uniform(-5.0, 5.0)) for _ in range(5)
]
delays = [_get_retry_delay(config, state) for _ in range(5)]
assert delays == expected_delays
finally:
adk_platform.reset_random_provider()
class TestShouldRetryNode: