Skip to content

[AI-163] google_adk_agents: support ToolContext session state in activity_tool - #1683

Open
DABH wants to merge 2 commits into
mainfrom
activity-tool-context-state
Open

[AI-163] google_adk_agents: support ToolContext session state in activity_tool#1683
DABH wants to merge 2 commits into
mainfrom
activity-tool-context-state

Conversation

@DABH

@DABH DABH commented Jul 27, 2026

Copy link
Copy Markdown
Contributor

What was changed

Adds ToolContextSnapshot to temporalio.contrib.google_adk_agents.workflow: a frozen, serializable dataclass carrying the session state (as a plain dict) and the function-call id of an ADK ToolContext.

An activity wrapped with activity_tool can now declare a parameter named tool_context annotated with ToolContextSnapshot (or ToolContextSnapshot | None):

@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 (ADK reserves the parameter name) and filled in at invocation time — the wrapper snapshots the live ToolContext workflow-side before scheduling the activity.

Because ADK detects the context parameter annotation-first across all parameters (find_context_parameter, falling back to the name tool_context), wrap-time validation covers the whole signature rather than just the reserved name. activity_tool now raises an actionable ValueError at wrap time for:

  • an ADK context annotation on any parameter (e.g. 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 a tool_context: ToolContextSnapshot sibling parameter to leak into the LLM-facing schema as a required argument;
  • ToolContextSnapshot on a parameter not named tool_context — ADK never injects it there, so it would leak into the tool schema;
  • an unannotated tool_context — without the type hint the activity would decode the snapshot as a plain dict under Temporal while receiving a ToolContextSnapshot instance 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:

  • The live ToolContext never 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.
  • Nothing leaks into the tool schema — the parameter rides ADK's own tool_context name-reservation, so declarations exclude it on both the legacy and JSON-schema paths; declarations that would break that exclusion are rejected.
  • Scoped to 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).
  • One-way by design — the snapshot is read-only; mutations inside the activity don't propagate, and session-state modifications must happen workflow-side from information the activity returns. Documented in the class docstring and README.

Checklist

  1. Closes [Feature Request] Support ToolContext for ADK function tools #1470
  2. How was this tested: new tests/contrib/google_adk_agents/test_adk_tool_context.py — end-to-end workflow run with max_cached_workflows=0 (mixed-type session state incl. nested dict and int, default parameter handling, populated function-call id, tool_context mid-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, unannotated tool_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.
  3. Any docs updates needed: contrib README updated in this PR.

@DABH
DABH requested review from a team as code owners July 27, 2026 20:28
@DABH DABH changed the title google_adk_agents: support ToolContext session state in activity_tool [AI-163] google_adk_agents: support ToolContext session state in activity_tool Jul 27, 2026
…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
DABH force-pushed the activity-tool-context-state branch from f05c905 to 6b80002 Compare July 27, 2026 20:39
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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Feature Request] Support ToolContext for ADK function tools

1 participant