[AI-163] google_adk_agents: support ToolContext session state in activity_tool - #1683
Open
DABH wants to merge 2 commits into
Open
[AI-163] google_adk_agents: support ToolContext session state in activity_tool#1683DABH wants to merge 2 commits into
DABH wants to merge 2 commits into
Conversation
…pshot 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
DABH
force-pushed
the
activity-tool-context-state
branch
from
July 27, 2026 20:39
f05c905 to
6b80002
Compare
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What was changed
Adds
ToolContextSnapshottotemporalio.contrib.google_adk_agents.workflow: a frozen, serializable dataclass carrying the session state (as a plain dict) and the function-call id of an ADKToolContext.An activity wrapped with
activity_toolcan now declare a parameter namedtool_contextannotated withToolContextSnapshot(orToolContextSnapshot | None):Exactly like a native ADK function tool's
tool_contextparameter, it is excluded from the LLM-facing tool schema (ADK reserves the parameter name) and filled in at invocation time — the wrapper snapshots the liveToolContextworkflow-side before scheduling the activity.Because ADK detects the context parameter annotation-first across all parameters (
find_context_parameter, falling back to the nametool_context), wrap-time validation covers the whole signature rather than just the reserved name.activity_toolnow raises an actionableValueErrorat wrap time for:ctx: ToolContext,Optional[ToolContext]) — ADK would inject the live, non-serializable context there regardless of the name, which is exactly the runtime serialization failure reported in [Feature Request] Support ToolContext for ADK function tools #1470, and which additionally caused atool_context: ToolContextSnapshotsibling parameter to leak into the LLM-facing schema as a required argument;ToolContextSnapshoton a parameter not namedtool_context— ADK never injects it there, so it would leak into the tool schema;tool_context— without the type hint the activity would decode the snapshot as a plain dict under Temporal while receiving aToolContextSnapshotinstance in local ADK runs.Why
Closes #1470 (JIRA AI-163). The reporter's tools read configuration seeded into session state (
tool_context.state.get("key")), which previously required hand-writing a plain tool function around each activity.The design follows the constraints laid out in the issue discussion:
ToolContextnever crosses the activity boundary — activity inputs go to the server and may run on a different worker, so only the explicitly-called-out serializable subset is passed, extracted workflow-side before invocation. Any declaration that would let ADK inject a live context is rejected at wrap time.tool_contextname-reservation, so declarations exclude it on both the legacy and JSON-schema paths; declarations that would break that exclusion are rejected.ToolContext.state(plus the function-call id), as a single input object so more fields can be added later without breaking activity signatures. Under Temporal the entire session state crosses the activity boundary, so all state values must be serializable by the configured data converter (documented in the class docstring and README).Checklist
tests/contrib/google_adk_agents/test_adk_tool_context.py— end-to-end workflow run withmax_cached_workflows=0(mixed-type session state incl. nested dict and int, default parameter handling, populated function-call id,tool_contextmid-signature to prove positional slotting, and an in-activity marker proving the snapshot crossed a real activity boundary), local-ADK-run parity outside a workflow, schema-exclusion assertions on both the JSON-schema and legacy declaration paths plus a context-only tool, and a wrap-time acceptance/rejection matrix (ToolContextSnapshot,ToolContextSnapshot | None, ADK context under any parameter name,Optional[ToolContext], misplaced snapshot annotation, unannotatedtool_context, unrelated annotation). Suite verified against both google-adk 2.2.0 (locked) and 2.5.0: 30 passed, 5 skipped (pre-existing). pyright/basedpyright/mypy/ruff/pydocstyle clean.