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:
@@ -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.
|
||||
|
||||
|
||||
@@ -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',
|
||||
]
|
||||
|
||||
@@ -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()()
|
||||
@@ -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"
|
||||
@@ -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:
|
||||
|
||||
|
||||
Reference in New Issue
Block a user