From 6b800026aa89eb29c736ec5c95897617eccc7b23 Mon Sep 17 00:00:00 2001 From: DABH Date: Mon, 27 Jul 2026 15:28:26 -0500 Subject: [PATCH 1/2] Support ToolContext session state in activity_tool via ToolContextSnapshot MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Activities wrapped with activity_tool could not access the ADK ToolContext: declaring a tool_context parameter put it in the LLM-facing tool schema and then failed at runtime trying to serialize the live ToolContext as an activity argument (#1470). An activity can now declare a parameter named tool_context annotated with the new ToolContextSnapshot dataclass. Exactly like a native ADK function tool's tool_context parameter, it is excluded from the tool schema (ADK reserves the name) and filled at invocation time — with a serializable snapshot of the live ToolContext (session state as a plain dict, plus the function-call id) taken workflow-side before the activity is scheduled. The live ToolContext never crosses the activity boundary. Annotating the parameter with an ADK context type raises an actionable error at wrap time instead of failing at serialization time. The snapshot is one-way by design: activities may run on different workers, so session-state modifications must happen workflow-side using information returned from the activity. Closes #1470 --- .../contrib/google_adk_agents/README.md | 34 +++ .../contrib/google_adk_agents/workflow.py | 116 ++++++++- .../test_adk_tool_context.py | 235 ++++++++++++++++++ 3 files changed, 384 insertions(+), 1 deletion(-) create mode 100644 tests/contrib/google_adk_agents/test_adk_tool_context.py diff --git a/temporalio/contrib/google_adk_agents/README.md b/temporalio/contrib/google_adk_agents/README.md index 40ebb9aee..e9a526cff 100644 --- a/temporalio/contrib/google_adk_agents/README.md +++ b/temporalio/contrib/google_adk_agents/README.md @@ -166,6 +166,40 @@ worker = Worker( ) ``` +### Reading Session State in Activity Tools + +ADK's live `ToolContext` holds non-serializable objects, so it cannot be an +activity argument. To read the serializable subset from an activity-backed +tool, declare a parameter named `tool_context` annotated with +`ToolContextSnapshot`: + +```python +from temporalio import activity +from temporalio.contrib.google_adk_agents.workflow import ( + ToolContextSnapshot, + activity_tool, +) + + +@activity.defn +async def get_weather(query: str, tool_context: ToolContextSnapshot) -> dict: + db_url = tool_context.state.get("url", "") + ... + + +weather_tool = activity_tool(get_weather, start_to_close_timeout=timedelta(seconds=30)) +``` + +Exactly like a native ADK function tool's `tool_context` parameter, it is +excluded from the LLM-facing tool schema; at invocation the wrapper snapshots +the live `ToolContext` (session state as a plain dict, plus the function-call +id) and passes it to the activity. + +The snapshot is one-way: mutations inside the activity do not propagate back +to the session, because the activity may run on a different worker. To modify +session state, return the needed information from the activity and apply it in +workflow-side code (for example an ADK callback or a plain tool function). + ### Local ADK Runs The same agent definitions can also be exercised outside Temporal with diff --git a/temporalio/contrib/google_adk_agents/workflow.py b/temporalio/contrib/google_adk_agents/workflow.py index b1d150391..d722ad93d 100644 --- a/temporalio/contrib/google_adk_agents/workflow.py +++ b/temporalio/contrib/google_adk_agents/workflow.py @@ -2,11 +2,113 @@ import functools import inspect -from typing import Any, Callable +import typing +from dataclasses import dataclass, field +from typing import Any, Callable, cast import temporalio.workflow from temporalio import workflow +_TOOL_CONTEXT_PARAM = "tool_context" + + +@dataclass(frozen=True) +class ToolContextSnapshot: + """Serializable snapshot of the ADK ``ToolContext`` for activity-backed tools. + + .. warning:: + This class is experimental and may change in future versions. + Use with caution in production environments. + + ADK's ``ToolContext`` holds live, non-serializable objects, so it cannot + cross the activity boundary: activity inputs are sent to the server and + may run on a different worker than the workflow. This snapshot carries the + serializable subset instead. + + Declare a parameter named ``tool_context`` annotated with this type on an + activity wrapped by :func:`activity_tool`: + + .. code-block:: python + + @activity.defn + async def get_weather(query: str, tool_context: ToolContextSnapshot) -> dict: + db_url = tool_context.state.get("url", "") + ... + + Exactly like a native ADK function tool's ``tool_context`` parameter, it + is excluded from the LLM-facing tool schema and filled in at invocation + time — here with a snapshot taken from the live ``ToolContext`` before the + activity is scheduled. + + The snapshot is one-way: mutating it inside the activity does not + propagate back to the session. To modify session state, return the needed + information from the activity and apply it in workflow-side code (for + example an ADK callback or a plain tool function). + + Attributes: + state: The session state visible to this tool call, as a plain dict. + function_call_id: The id of the function call being handled, when + available. + """ + + state: dict[str, Any] = field(default_factory=dict) + function_call_id: str | None = None + + +def _tool_context_parameter(activity_def: Callable) -> inspect.Parameter | None: + """Returns the activity's ``tool_context`` parameter after validating it. + + ADK reserves the parameter name ``tool_context`` for context injection and + never includes it in the LLM-facing tool schema, so a parameter with that + name can only ever be filled by the integration. It must be annotated with + :class:`ToolContextSnapshot` (or left unannotated). + """ + parameter = inspect.signature(activity_def).parameters.get(_TOOL_CONTEXT_PARAM) + if parameter is None: + return None + if parameter.annotation is inspect.Parameter.empty: + return parameter + try: + resolved = typing.get_type_hints(activity_def).get( + _TOOL_CONTEXT_PARAM, parameter.annotation + ) + except Exception: + resolved = parameter.annotation + if resolved is ToolContextSnapshot: + return parameter + # Accept an optional annotation too (ToolContextSnapshot | None). + if set(typing.get_args(resolved)) == {ToolContextSnapshot, type(None)}: + return parameter + annotation_name = getattr(resolved, "__name__", str(resolved)) + if getattr(resolved, "__module__", "").startswith("google.adk"): + raise ValueError( + f"Activity '{activity_def.__name__}' declares 'tool_context:" + f" {annotation_name}', but ADK context objects are not serializable" + " and cannot be activity arguments. Annotate the parameter with" + " ToolContextSnapshot instead to receive the serializable subset" + " (session state and function-call id)." + ) + raise ValueError( + f"Activity '{activity_def.__name__}' has a 'tool_context' parameter" + f" annotated with {annotation_name}. The name 'tool_context' is" + " reserved by ADK for context injection; annotate the parameter with" + " ToolContextSnapshot to receive the serializable subset of the tool" + " context." + ) + + +def _snapshot_tool_context(tool_context: Any) -> ToolContextSnapshot: + """Builds the serializable snapshot from a live ADK ``ToolContext``.""" + state: dict[str, Any] = {} + state_object = getattr(tool_context, "state", None) + if state_object is not None: + to_dict = getattr(state_object, "to_dict", None) + state = dict(cast(Any, to_dict() if callable(to_dict) else state_object)) + return ToolContextSnapshot( + state=state, + function_call_id=getattr(tool_context, "function_call_id", None), + ) + def activity_tool(activity_def: Callable, **kwargs: Any) -> Callable: """Decorator/Wrapper to wrap a Temporal Activity as an ADK Tool. @@ -17,10 +119,22 @@ def activity_tool(activity_def: Callable, **kwargs: Any) -> Callable: This ensures the activity's signature is preserved for ADK's tool schema generation while marking it as a tool that executes via 'workflow.execute_activity'. + + If the activity declares a parameter named ``tool_context``, it must be + annotated with :class:`ToolContextSnapshot`. ADK excludes the parameter + from the tool schema and injects the live ``ToolContext`` into the + wrapper, which passes the activity a serializable snapshot of it (session + state and function-call id) in that parameter's position. """ + tool_context_param = _tool_context_parameter(activity_def) @functools.wraps(activity_def) async def wrapper(*args: Any, **kw: Any): + # ADK injects the live ToolContext by name; replace it with the + # serializable snapshot the activity actually declares. + if tool_context_param is not None and _TOOL_CONTEXT_PARAM in kw: + kw[_TOOL_CONTEXT_PARAM] = _snapshot_tool_context(kw[_TOOL_CONTEXT_PARAM]) + # Inspect signature to bind arguments sig = inspect.signature(activity_def) bound = sig.bind(*args, **kw) diff --git a/tests/contrib/google_adk_agents/test_adk_tool_context.py b/tests/contrib/google_adk_agents/test_adk_tool_context.py new file mode 100644 index 000000000..eabab3a8c --- /dev/null +++ b/tests/contrib/google_adk_agents/test_adk_tool_context.py @@ -0,0 +1,235 @@ +# 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 ToolContextSnapshot injection into activity-backed tools. + +Covers https://github.com/temporalio/sdk-python/issues/1470: activities +wrapped with activity_tool can read the serializable subset of the ADK +ToolContext (session state and function-call id) without the live, +non-serializable ToolContext ever crossing the activity boundary, and +without the parameter leaking into the LLM-facing tool schema. +""" + +import uuid +from collections.abc import AsyncGenerator +from datetime import timedelta +from typing import Any + +import pytest +from google.adk import Agent +from google.adk.models import BaseLlm, LLMRegistry +from google.adk.models.llm_request import LlmRequest +from google.adk.models.llm_response import LlmResponse +from google.adk.runners import InMemoryRunner +from google.adk.tools.function_tool import FunctionTool +from google.adk.tools.tool_context import ToolContext +from google.adk.utils.context_utils import Aclosing +from google.genai import types +from google.genai.types import Content, FunctionCall, Part + +from temporalio import activity, workflow +from temporalio.client import Client +from temporalio.contrib.google_adk_agents import GoogleAdkPlugin, TemporalModel +from temporalio.contrib.google_adk_agents.workflow import ( + ToolContextSnapshot, + activity_tool, +) +from temporalio.worker import Worker + +TASK_QUEUE = "adk-tool-context-task-queue" + + +@activity.defn +async def lookup_weather( + city: str, tool_context: ToolContextSnapshot, units: str = "celsius" +) -> str: + """Activity that reads tool configuration from session state. + + Mirrors the shape reported in issue #1470, with tool_context deliberately + in the middle of the parameter list to prove positional slotting. + """ + db_url = tool_context.state.get("db_url", "") + has_function_call_id = "yes" if tool_context.function_call_id else "no" + return f"{city}|{units}|{db_url}|fc={has_function_call_id}" + + +def weather_agent(model_name: str) -> Agent: + return Agent( + name="state_agent", + model=TemporalModel(model_name), + tools=[ + activity_tool(lookup_weather, start_to_close_timeout=timedelta(seconds=30)) + ], + ) + + +class StateToolModel(BaseLlm): + """Scripted model: call lookup_weather once, then echo its response.""" + + async def generate_content_async( + self, llm_request: LlmRequest, stream: bool = False + ) -> AsyncGenerator[LlmResponse, None]: + tool_response = None + for content in llm_request.contents: + for part in content.parts or []: + if ( + part.function_response is not None + and part.function_response.name == "lookup_weather" + ): + tool_response = part.function_response + if tool_response is None: + yield LlmResponse( + content=Content( + role="model", + parts=[ + Part( + function_call=FunctionCall( + name="lookup_weather", args={"city": "NYC"} + ) + ) + ], + ) + ) + else: + yield LlmResponse( + content=Content( + role="model", + parts=[Part(text=f"done:{tool_response.response}")], + ) + ) + + @classmethod + def supported_models(cls) -> list[str]: + return ["state_tool_model"] + + +async def run_state_agent(model_name: str) -> str: + """Runs the agent against a session seeded with state; returns final text.""" + runner = InMemoryRunner(agent=weather_agent(model_name), app_name="test_app") + session = await runner.session_service.create_session( + app_name="test_app", + user_id="test", + state={"db_url": "postgres://config"}, + ) + final_text = "" + async with Aclosing( + runner.run_async( + user_id="test", + session_id=session.id, + new_message=types.Content( + role="user", parts=[types.Part(text="weather in NYC?")] + ), + ) + ) as agen: + async for event in agen: + if event.content and event.content.parts and event.content.parts[0].text: + final_text = event.content.parts[0].text + return final_text + + +@workflow.defn +class StateToolWorkflow: + @workflow.run + async def run(self, model_name: str) -> str: + return await run_state_agent(model_name) + + +@pytest.mark.asyncio +async def test_activity_tool_receives_tool_context_snapshot(client: Client): + new_config = client.config() + new_config["plugins"] = [GoogleAdkPlugin()] + client = Client(**new_config) + + async with Worker( + client, + task_queue=TASK_QUEUE, + activities=[lookup_weather], + workflows=[StateToolWorkflow], + max_cached_workflows=0, + ): + LLMRegistry.register(StateToolModel) + result = await client.execute_workflow( + StateToolWorkflow.run, + "state_tool_model", + id=f"tool-context-{uuid.uuid4()}", + task_queue=TASK_QUEUE, + execution_timeout=timedelta(seconds=30), + ) + # The activity saw the session state, the default parameter value, and a + # populated function-call id — none of which came from the LLM. + assert "NYC|celsius|postgres://config|fc=yes" in result + + +@pytest.mark.asyncio +async def test_activity_tool_snapshot_outside_workflow(): + """Local ADK runs (no Temporal) receive the same snapshot.""" + LLMRegistry.register(StateToolModel) + result = await run_state_agent("state_tool_model") + assert "NYC|celsius|postgres://config|fc=yes" in result + + +def _declared_properties(tool: FunctionTool) -> dict[str, Any]: + """Property names in the LLM-facing declaration, across schema styles.""" + declaration = tool._get_declaration() + assert declaration is not None + if declaration.parameters_json_schema is not None: + return declaration.parameters_json_schema.get("properties", {}) + assert declaration.parameters is not None + return declaration.parameters.properties or {} + + +def test_tool_schema_excludes_tool_context(): + """The tool_context parameter never appears in the LLM-facing schema.""" + tool = FunctionTool( + func=activity_tool(lookup_weather, start_to_close_timeout=timedelta(seconds=30)) + ) + properties = _declared_properties(tool) + assert "city" in properties + assert "units" in properties + assert "tool_context" not in properties + + +def test_activity_tool_rejects_adk_tool_context_annotation(): + """Annotating with the live ADK ToolContext gives an actionable error.""" + + @activity.defn + async def bad_tool(query: str, tool_context: ToolContext) -> str: # pyright: ignore[reportUnusedParameter] + return query + + with pytest.raises(ValueError, match="ToolContextSnapshot"): + activity_tool(bad_tool, start_to_close_timeout=timedelta(seconds=30)) + + +def test_activity_tool_rejects_other_tool_context_annotation(): + """The reserved name with an unrelated annotation gives an actionable error.""" + + @activity.defn + async def confused_tool(query: str, tool_context: dict[str, Any]) -> str: # pyright: ignore[reportUnusedParameter] + return query + + with pytest.raises(ValueError, match="reserved by ADK"): + activity_tool(confused_tool, start_to_close_timeout=timedelta(seconds=30)) + + +def test_activity_tool_without_tool_context_unchanged(): + """Activities without a tool_context parameter keep their exact schema.""" + + @activity.defn + async def plain_tool(query: str) -> str: + return query + + tool = FunctionTool( + func=activity_tool(plain_tool, start_to_close_timeout=timedelta(seconds=30)) + ) + assert set(_declared_properties(tool)) == {"query"} From 093f27952eb88e473e2c5b716b72751d09c1e7a2 Mon Sep 17 00:00:00 2001 From: DABH Date: Tue, 28 Jul 2026 14:37:30 -0500 Subject: [PATCH 2/2] Harden tool_context validation, close annotation-based injection bypass MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ADK detects the context parameter annotation-first across all parameters (find_context_parameter), falling back to the name 'tool_context', so the previous name-only validation missed cases where ADK would still inject the live, non-serializable context: - Reject an ADK context annotation on any parameter, not just 'tool_context' (e.g. 'ctx: ToolContext' previously passed wrap-time validation, then failed payload serialization at runtime — the exact failure mode from #1470 — and, when combined with a 'tool_context: ToolContextSnapshot' parameter, leaked that parameter into the LLM-facing tool schema as required). - Reject ToolContextSnapshot on a parameter not named 'tool_context' (ADK never injects it there, so it would leak into the tool schema). - Reject an unannotated 'tool_context' (public docs already required the annotation; without it the activity decoded the snapshot as a plain dict under Temporal but received a ToolContextSnapshot in local runs). - Route Optional[ToolContext] to the ADK-specific error message and render union annotations readably. Docs: state serializability/payload-size constraints under Temporal and read-only snapshot semantics (nested values may alias live session state in local runs); add missing timedelta import to the README snippet. Tests: new wrap-time acceptance/rejection matrix (Optional form, ADK context under any name, misplaced snapshot, unannotated), legacy (non-JSON-schema) declaration path, context-only tool schema, mixed-type session state round-trip, and an in-activity marker proving the snapshot crosses a real activity boundary; drop incorrect copyright header. Verified against google-adk 2.2.0 (locked) and 2.5.0. --- .../contrib/google_adk_agents/README.md | 23 ++- .../contrib/google_adk_agents/workflow.py | 150 +++++++++++++---- .../test_adk_tool_context.py | 151 +++++++++++++++--- 3 files changed, 261 insertions(+), 63 deletions(-) diff --git a/temporalio/contrib/google_adk_agents/README.md b/temporalio/contrib/google_adk_agents/README.md index e9a526cff..4a15c39d6 100644 --- a/temporalio/contrib/google_adk_agents/README.md +++ b/temporalio/contrib/google_adk_agents/README.md @@ -174,6 +174,8 @@ tool, declare a parameter named `tool_context` annotated with `ToolContextSnapshot`: ```python +from datetime import timedelta + from temporalio import activity from temporalio.contrib.google_adk_agents.workflow import ( ToolContextSnapshot, @@ -193,12 +195,21 @@ weather_tool = activity_tool(get_weather, start_to_close_timeout=timedelta(secon Exactly like a native ADK function tool's `tool_context` parameter, it is excluded from the LLM-facing tool schema; at invocation the wrapper snapshots the live `ToolContext` (session state as a plain dict, plus the function-call -id) and passes it to the activity. - -The snapshot is one-way: mutations inside the activity do not propagate back -to the session, because the activity may run on a different worker. To modify -session state, return the needed information from the activity and apply it in -workflow-side code (for example an ADK callback or a plain tool function). +id) and passes it to the activity. Annotating any parameter with a live ADK +context type raises `ValueError` at wrap time, since ADK would inject the +non-serializable context into it regardless of its name. + +When running under Temporal, the entire session state crosses the activity +boundary: every value in it must be serializable by the configured data +converter and the total size must fit within payload limits, even for keys +the tool never reads. Local ADK runs pass the snapshot in memory and have no +such constraint. + +The snapshot is one-way and should be treated as read-only: mutations inside +the activity do not propagate back to the session, because the activity may +run on a different worker. To modify session state, return the needed +information from the activity and apply it in workflow-side code (for example +an ADK callback or a plain tool function). ### Local ADK Runs diff --git a/temporalio/contrib/google_adk_agents/workflow.py b/temporalio/contrib/google_adk_agents/workflow.py index d722ad93d..13bca4de3 100644 --- a/temporalio/contrib/google_adk_agents/workflow.py +++ b/temporalio/contrib/google_adk_agents/workflow.py @@ -2,6 +2,7 @@ import functools import inspect +import types import typing from dataclasses import dataclass, field from typing import Any, Callable, cast @@ -25,8 +26,9 @@ class ToolContextSnapshot: may run on a different worker than the workflow. This snapshot carries the serializable subset instead. - Declare a parameter named ``tool_context`` annotated with this type on an - activity wrapped by :func:`activity_tool`: + Declare a parameter named ``tool_context`` annotated with this type (or + ``ToolContextSnapshot | None``) on an activity wrapped by + :func:`activity_tool`: .. code-block:: python @@ -40,8 +42,17 @@ async def get_weather(query: str, tool_context: ToolContextSnapshot) -> dict: time — here with a snapshot taken from the live ``ToolContext`` before the activity is scheduled. - The snapshot is one-way: mutating it inside the activity does not - propagate back to the session. To modify session state, return the needed + When running under Temporal, the entire session state crosses the + activity boundary: every value in it must be serializable by the + configured data converter and the total size must fit within payload + limits, even for keys the tool never reads. A non-serializable value + fails the workflow task when the activity is scheduled. Local ADK runs + pass the snapshot in memory and have no such constraint. + + The snapshot is one-way and should be treated as read-only: mutating it + inside the activity does not propagate back to the session (and in local + runs nested values may alias the live session state, so mutating them can + corrupt the session). To modify session state, return the needed information from the activity and apply it in workflow-side code (for example an ADK callback or a plain tool function). @@ -55,39 +66,59 @@ async def get_weather(query: str, tool_context: ToolContextSnapshot) -> dict: function_call_id: str | None = None -def _tool_context_parameter(activity_def: Callable) -> inspect.Parameter | None: - """Returns the activity's ``tool_context`` parameter after validating it. +def _annotation_members(annotation: Any) -> tuple[Any, ...]: + """Returns a union annotation's members, or the annotation itself.""" + origin = typing.get_origin(annotation) + if origin is typing.Union or origin is types.UnionType: # pyright: ignore[reportDeprecated] + return typing.get_args(annotation) + return (annotation,) + + +def _annotation_display(members: tuple[Any, ...]) -> str: + """Renders annotation members for error messages.""" + return " | ".join( + "None" if member is type(None) else getattr(member, "__name__", str(member)) + for member in members + ) + + +def _adk_context_error( + activity_def: Callable, name: str, annotation_name: str +) -> ValueError: + """Builds the error for a parameter annotated with a live ADK context.""" + return ValueError( + f"Activity '{activity_def.__name__}' declares '{name}:" + f" {annotation_name}', but ADK context objects are not serializable" + " and cannot be activity arguments. Declare a parameter named" + " 'tool_context' annotated with ToolContextSnapshot instead to receive" + " the serializable subset (session state and function-call id)." + ) - ADK reserves the parameter name ``tool_context`` for context injection and - never includes it in the LLM-facing tool schema, so a parameter with that - name can only ever be filled by the integration. It must be annotated with - :class:`ToolContextSnapshot` (or left unannotated). + +def _validated_tool_context_parameter( + activity_def: Callable, parameter: inspect.Parameter, members: tuple[Any, ...] +) -> inspect.Parameter: + """Returns the ``tool_context`` parameter after validating its annotation. + + The parameter must be annotated with :class:`ToolContextSnapshot` or + ``ToolContextSnapshot | None``. """ - parameter = inspect.signature(activity_def).parameters.get(_TOOL_CONTEXT_PARAM) - if parameter is None: - return None if parameter.annotation is inspect.Parameter.empty: - return parameter - try: - resolved = typing.get_type_hints(activity_def).get( - _TOOL_CONTEXT_PARAM, parameter.annotation - ) - except Exception: - resolved = parameter.annotation - if resolved is ToolContextSnapshot: - return parameter - # Accept an optional annotation too (ToolContextSnapshot | None). - if set(typing.get_args(resolved)) == {ToolContextSnapshot, type(None)}: - return parameter - annotation_name = getattr(resolved, "__name__", str(resolved)) - if getattr(resolved, "__module__", "").startswith("google.adk"): raise ValueError( - f"Activity '{activity_def.__name__}' declares 'tool_context:" - f" {annotation_name}', but ADK context objects are not serializable" - " and cannot be activity arguments. Annotate the parameter with" - " ToolContextSnapshot instead to receive the serializable subset" + f"Activity '{activity_def.__name__}' has an unannotated" + " 'tool_context' parameter. Annotate it with ToolContextSnapshot" + " to receive the serializable subset of the ADK tool context" " (session state and function-call id)." ) + non_none = [member for member in members if member is not type(None)] + if non_none == [ToolContextSnapshot]: + return parameter + annotation_name = _annotation_display(members) + if any( + getattr(member, "__module__", "").startswith("google.adk") + for member in non_none + ): + raise _adk_context_error(activity_def, _TOOL_CONTEXT_PARAM, annotation_name) raise ValueError( f"Activity '{activity_def.__name__}' has a 'tool_context' parameter" f" annotated with {annotation_name}. The name 'tool_context' is" @@ -97,6 +128,52 @@ def _tool_context_parameter(activity_def: Callable) -> inspect.Parameter | None: ) +def _tool_context_parameter(activity_def: Callable) -> inspect.Parameter | None: + """Validates context-related parameters and returns the ``tool_context`` one. + + ADK injects the live context into the first parameter annotated with an + ADK context type (regardless of name), or failing that into one named + ``tool_context``, and excludes that parameter from the LLM-facing tool + schema. Live context objects cannot cross the activity boundary, so the + only supported declaration is a parameter named ``tool_context`` annotated + with :class:`ToolContextSnapshot` (or ``ToolContextSnapshot | None``); + anything else ADK would treat as a context parameter is rejected at wrap + time, as is a misplaced ToolContextSnapshot annotation that would leak + into the tool schema. + """ + try: + hints = typing.get_type_hints(activity_def) + except Exception: + hints = {} + adk_context: type[Any] | None + try: + from google.adk.tools.tool_context import ToolContext + + adk_context = ToolContext + except ImportError: + adk_context = None + tool_context_parameter: inspect.Parameter | None = None + for name, parameter in inspect.signature(activity_def).parameters.items(): + members = _annotation_members(hints.get(name, parameter.annotation)) + if name == _TOOL_CONTEXT_PARAM: + tool_context_parameter = _validated_tool_context_parameter( + activity_def, parameter, members + ) + elif adk_context is not None and any( + member is adk_context for member in members + ): + raise _adk_context_error(activity_def, name, _annotation_display(members)) + elif any(member is ToolContextSnapshot for member in members): + raise ValueError( + f"Activity '{activity_def.__name__}' annotates parameter" + f" '{name}' with ToolContextSnapshot, but ADK only injects the" + " tool context into a parameter named 'tool_context'; under" + " any other name it would appear in the LLM-facing tool" + " schema. Rename the parameter to 'tool_context'." + ) + return tool_context_parameter + + def _snapshot_tool_context(tool_context: Any) -> ToolContextSnapshot: """Builds the serializable snapshot from a live ADK ``ToolContext``.""" state: dict[str, Any] = {} @@ -121,10 +198,13 @@ def activity_tool(activity_def: Callable, **kwargs: Any) -> Callable: while marking it as a tool that executes via 'workflow.execute_activity'. If the activity declares a parameter named ``tool_context``, it must be - annotated with :class:`ToolContextSnapshot`. ADK excludes the parameter - from the tool schema and injects the live ``ToolContext`` into the - wrapper, which passes the activity a serializable snapshot of it (session - state and function-call id) in that parameter's position. + annotated with :class:`ToolContextSnapshot` (or ``ToolContextSnapshot | + None``). ADK excludes the parameter from the tool schema and injects the + live ``ToolContext`` into the wrapper, which passes the activity a + serializable snapshot of it (session state and function-call id) in that + parameter's position. Annotating any parameter with a live ADK context + type raises ``ValueError`` at wrap time, since ADK would inject the + non-serializable context into it regardless of its name. """ tool_context_param = _tool_context_parameter(activity_def) diff --git a/tests/contrib/google_adk_agents/test_adk_tool_context.py b/tests/contrib/google_adk_agents/test_adk_tool_context.py index eabab3a8c..be5b8a891 100644 --- a/tests/contrib/google_adk_agents/test_adk_tool_context.py +++ b/tests/contrib/google_adk_agents/test_adk_tool_context.py @@ -1,17 +1,3 @@ -# 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 ToolContextSnapshot injection into activity-backed tools. Covers https://github.com/temporalio/sdk-python/issues/1470: activities @@ -24,10 +10,11 @@ import uuid from collections.abc import AsyncGenerator from datetime import timedelta -from typing import Any +from typing import Any, Optional # pyright: ignore[reportDeprecated] import pytest from google.adk import Agent +from google.adk.features import FeatureName, override_feature_enabled from google.adk.models import BaseLlm, LLMRegistry from google.adk.models.llm_request import LlmRequest from google.adk.models.llm_response import LlmResponse @@ -49,6 +36,12 @@ TASK_QUEUE = "adk-tool-context-task-queue" +SESSION_STATE: dict[str, Any] = { + "db_url": "postgres://config", + "retries": 3, + "regions": {"primary": "us-east1", "replicas": ["eu-west1"]}, +} + @activity.defn async def lookup_weather( @@ -57,11 +50,16 @@ async def lookup_weather( """Activity that reads tool configuration from session state. Mirrors the shape reported in issue #1470, with tool_context deliberately - in the middle of the parameter list to prove positional slotting. + in the middle of the parameter list to prove positional slotting. The + returned string also records whether the code ran inside a real activity, + so tests can tell the activity boundary was actually crossed (or not). """ db_url = tool_context.state.get("db_url", "") + retries = tool_context.state.get("retries", -1) + region = tool_context.state.get("regions", {}).get("primary", "") has_function_call_id = "yes" if tool_context.function_call_id else "no" - return f"{city}|{units}|{db_url}|fc={has_function_call_id}" + in_act = "yes" if activity.in_activity() else "no" + return f"{city}|{units}|{db_url}|{retries}|{region}|fc={has_function_call_id}|act={in_act}" def weather_agent(model_name: str) -> Agent: @@ -120,7 +118,7 @@ async def run_state_agent(model_name: str) -> str: session = await runner.session_service.create_session( app_name="test_app", user_id="test", - state={"db_url": "postgres://config"}, + state=SESSION_STATE, ) final_text = "" async with Aclosing( @@ -166,9 +164,11 @@ async def test_activity_tool_receives_tool_context_snapshot(client: Client): task_queue=TASK_QUEUE, execution_timeout=timedelta(seconds=30), ) - # The activity saw the session state, the default parameter value, and a - # populated function-call id — none of which came from the LLM. - assert "NYC|celsius|postgres://config|fc=yes" in result + # The activity saw mixed-type session state (string, int, nested dict), + # the default parameter value, and a populated function-call id — none of + # which came from the LLM — and act=yes proves the snapshot crossed a + # real activity boundary rather than running inline in the workflow. + assert "NYC|celsius|postgres://config|3|us-east1|fc=yes|act=yes" in result @pytest.mark.asyncio @@ -176,7 +176,7 @@ async def test_activity_tool_snapshot_outside_workflow(): """Local ADK runs (no Temporal) receive the same snapshot.""" LLMRegistry.register(StateToolModel) result = await run_state_agent("state_tool_model") - assert "NYC|celsius|postgres://config|fc=yes" in result + assert "NYC|celsius|postgres://config|3|us-east1|fc=yes|act=no" in result def _declared_properties(tool: FunctionTool) -> dict[str, Any]: @@ -200,6 +200,66 @@ def test_tool_schema_excludes_tool_context(): assert "tool_context" not in properties +def test_tool_schema_excludes_tool_context_legacy_declaration(): + """Exclusion also holds on the legacy (non-JSON-schema) declaration path.""" + override_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL, False) + try: + tool = FunctionTool( + func=activity_tool( + lookup_weather, start_to_close_timeout=timedelta(seconds=30) + ) + ) + declaration = tool._get_declaration() + assert declaration is not None + assert declaration.parameters is not None + properties = declaration.parameters.properties or {} + assert "city" in properties + assert "units" in properties + assert "tool_context" not in properties + finally: + # The flag is default-on across the supported google-adk range. + override_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL, True) + + +def test_tool_schema_context_only_parameter(): + """A tool whose only parameter is tool_context exposes no LLM arguments.""" + + @activity.defn + async def ctx_only_tool(tool_context: ToolContextSnapshot) -> str: + return str(tool_context.state) + + tool = FunctionTool( + func=activity_tool(ctx_only_tool, start_to_close_timeout=timedelta(seconds=30)) + ) + declaration = tool._get_declaration() + if declaration is not None: + json_properties = (declaration.parameters_json_schema or {}).get( + "properties", {} + ) + legacy_properties = ( + (declaration.parameters.properties or {}) if declaration.parameters else {} + ) + assert not json_properties + assert not legacy_properties + + +def test_activity_tool_accepts_optional_snapshot_annotation(): + """ToolContextSnapshot | None is accepted and still excluded from the schema.""" + + @activity.defn + async def optional_tool( + query: str, + tool_context: ToolContextSnapshot | None = None, # pyright: ignore[reportUnusedParameter] + ) -> str: + return query + + tool = FunctionTool( + func=activity_tool(optional_tool, start_to_close_timeout=timedelta(seconds=30)) + ) + properties = _declared_properties(tool) + assert set(properties) == {"query"} + + def test_activity_tool_rejects_adk_tool_context_annotation(): """Annotating with the live ADK ToolContext gives an actionable error.""" @@ -211,6 +271,53 @@ async def bad_tool(query: str, tool_context: ToolContext) -> str: # pyright: ig activity_tool(bad_tool, start_to_close_timeout=timedelta(seconds=30)) +def test_activity_tool_rejects_optional_adk_tool_context_annotation(): + """Optional[ToolContext] is rejected with the ADK-specific message.""" + + @activity.defn + async def optional_bad_tool( + query: str, + tool_context: Optional[ToolContext] = None, # pyright: ignore[reportUnusedParameter, reportDeprecated] + ) -> str: + return query + + with pytest.raises(ValueError, match="not serializable"): + activity_tool(optional_bad_tool, start_to_close_timeout=timedelta(seconds=30)) + + +def test_activity_tool_rejects_adk_context_under_any_name(): + """ADK injects into any param annotated with a context type, so all are rejected.""" + + @activity.defn + async def sneaky_tool(query: str, ctx: ToolContext) -> str: # pyright: ignore[reportUnusedParameter] + return query + + with pytest.raises(ValueError, match="not serializable"): + activity_tool(sneaky_tool, start_to_close_timeout=timedelta(seconds=30)) + + +def test_activity_tool_rejects_snapshot_under_other_name(): + """ToolContextSnapshot on a differently-named param would leak into the schema.""" + + @activity.defn + async def misnamed_tool(query: str, snap: ToolContextSnapshot) -> str: # pyright: ignore[reportUnusedParameter] + return query + + with pytest.raises(ValueError, match="named 'tool_context'"): + activity_tool(misnamed_tool, start_to_close_timeout=timedelta(seconds=30)) + + +def test_activity_tool_rejects_unannotated_tool_context(): + """The reserved name without an annotation gives an actionable error.""" + + @activity.defn + async def untyped_tool(query: str, tool_context) -> str: # type: ignore[no-untyped-def] # pyright: ignore[reportUnusedParameter, reportMissingParameterType] + return query + + with pytest.raises(ValueError, match="unannotated 'tool_context'"): + activity_tool(untyped_tool, start_to_close_timeout=timedelta(seconds=30)) + + def test_activity_tool_rejects_other_tool_context_annotation(): """The reserved name with an unrelated annotation gives an actionable error."""