diff --git a/temporalio/contrib/google_adk_agents/README.md b/temporalio/contrib/google_adk_agents/README.md index 40ebb9aee..4a15c39d6 100644 --- a/temporalio/contrib/google_adk_agents/README.md +++ b/temporalio/contrib/google_adk_agents/README.md @@ -166,6 +166,51 @@ 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 datetime import timedelta + +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. 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 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..13bca4de3 100644 --- a/temporalio/contrib/google_adk_agents/workflow.py +++ b/temporalio/contrib/google_adk_agents/workflow.py @@ -2,11 +2,190 @@ import functools import inspect -from typing import Any, Callable +import types +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 (or + ``ToolContextSnapshot | None``) 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. + + 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). + + 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 _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)." + ) + + +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``. + """ + if parameter.annotation is inspect.Parameter.empty: + raise ValueError( + 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" + " reserved by ADK for context injection; annotate the parameter with" + " ToolContextSnapshot to receive the serializable subset of the tool" + " context." + ) + + +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] = {} + 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 +196,25 @@ 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` (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) @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..be5b8a891 --- /dev/null +++ b/tests/contrib/google_adk_agents/test_adk_tool_context.py @@ -0,0 +1,342 @@ +"""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, 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 +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" + +SESSION_STATE: dict[str, Any] = { + "db_url": "postgres://config", + "retries": 3, + "regions": {"primary": "us-east1", "replicas": ["eu-west1"]}, +} + + +@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. 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" + 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: + 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=SESSION_STATE, + ) + 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 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 +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|3|us-east1|fc=yes|act=no" 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_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.""" + + @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_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.""" + + @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"}