security: restrict pickle deserialization in v0 schema and migration tool by daridor9 · Pull Request #5635 · google/adk-python

daridor9 · 2026-05-07T18:18:44Z

Summary

This PR adds a RestrictedUnpickler to prevent arbitrary code execution via crafted pickle payloads in the v0 schema runtime path and the migration tool.

Fixes #5634

Changes

New: `_safe_unpickle.py`

RestrictedUnpickler subclass that allowlists only known-safe types:
- ADK model types (google.adk.*)
- Google GenAI types (google.genai.*)
- Pydantic internals (pydantic.*, pydantic_core.*)
- Standard Python builtins and collections/datetime/copyreg
Blocks all other globals (e.g., os, subprocess, posix, builtins.__import__)
safe_loads() function as drop-in replacement for pickle.loads()
Escape hatch: ADK_ALLOW_UNSAFE_V0_PICKLE=1 env var for databases with non-standard pickled types (logs deprecation warning)

Patched: `v0.py` (Sink 1 — runtime read path)

DynamicPickleType.process_result_value() now calls safe_loads() instead of pickle.loads()
Affects MySQL, Cloud Spanner, and SQLite dialects

Patched: `migrate_from_sqlalchemy_pickle.py` (Sink 2 — migration path)

_row_to_event() now calls safe_loads() instead of pickle.loads()
Ensures adk migrate session is safe to run on untrusted v0 databases

Testing

Existing v0 databases with legitimate EventActions pickle data will continue to work — all ADK types are allowlisted
Malicious pickle payloads (e.g., os.system, subprocess.Popen) are blocked with a clear error message
The env var escape hatch provides a migration path for edge cases

Security Context

Reported to Google VRP: 2026-04-24
GHSA filed
CVSS 3.1: 9.9 Critical (AV:N/AC:L/PR:L/UI:N/S:C/C:H/I:H/A:H)
90-day coordinated disclosure deadline: 2026-07-23

Submitted by SPR{K}3 Security Research (@daridor9)

google-cla · 2026-05-07T18:18:56Z

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

adk-bot · 2026-05-07T18:20:13Z

Hello @daridor9, thank you for your contribution! This is a valuable security fix.

Before we can proceed with the review, could you please sign the Contributor License Agreement (CLA)? The cla/google check is currently failing. You can find details on how to sign it in our contribution guidelines.

Once the CLA is signed, we can move forward with labeling and reviewing your PR.

Response from ADK Triaging Agent

rohityan · 2026-05-08T22:42:21Z

Hi @daridor9 , Thank you for your contribution! We appreciate you taking the time to submit this pull request. Please fix formatting errors by running autoformat.sh

daridor9 · 2026-05-11T05:30:05Z

Done — ran autoformat.sh (pyink + isort) across all three files. See commits 422e1d8, 1407e8d, ac60e57.

rohityan · 2026-05-29T18:47:14Z

Hi @daridor9 , please fix mypy-diff errors.

daridor9 · 2026-05-30T05:41:52Z

Hi @rohityan , Pushed 5f22404 addressing the mypy-diff failure.

mypy-diff: added type annotations to DynamicPickleType.result_processor and its inner process() in v0.py (the two [no-untyped-def] errors).
pre-commit: ran pyink on test_safe_unpickle.py, which was 4-space indented.

Verified both locally against the branch merged with current main: mypy-diff reports no new errors, and pre-commit passes all hooks. The CI runs are awaiting maintainer approval — ready for re-run whenever you have a moment.

daridor9 · 2026-06-03T07:46:01Z

Hi @rohityan — just following up. Formatting is done. Is there anything specific outstanding? Happy to address whatever is needed to move this forward. Best, Dan

Ref: google#5634

Replace raw pickle.loads() with safe_loads() in DynamicPickleType.process_result_value(). Ref: google#5634

Covers potential enum values in state_delta/agent_state dict[str, Any] fields. Ref: google#5634

Ref: google#5634

This test suite verifies that the RestrictedUnpickler correctly blocks malicious pickle payloads and ensures legitimate data can be safely unpickled. It also checks the behavior of the unpickler when an environment variable allows unsafe pickling.

Add tests Ref: google/adk-python#5634for EventActions serialization and deserialization.

…test file - Add Any type hints to DynamicPickleType.result_processor and its inner process() to clear mypy-diff [no-untyped-def]. - Reformat test_safe_unpickle.py to 2-space pyink style.

daridor9 · 2026-06-04T04:24:38Z

Hi, @rohityan main now includes a fix for the migration path (migrate_from_sqlalchemy_pickle.py). This PR has been rebased accordingly and no longer modifies that file.

The runtime sink in v0.py (DynamicPickleType.process_result_value) remains unpatched in main — every live agent session read still calls raw pickle.loads() on database bytes. This PR closes that gap.

Remaining scope: _safe_unpickle.py (shared restricted unpickler) + v0.py (runtime read path) + tests.

Co-authored-by: Sasha Sobran <asobran@google.com> PiperOrigin-RevId: 917516394

PiperOrigin-RevId: 917525676

Co-authored-by: Kathy Wu <wukathy@google.com> PiperOrigin-RevId: 918037985

Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 918503151

Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 918539981

Bump version to 2.2.0. Co-authored-by: Shangjie Chen <deanchen@google.com> PiperOrigin-RevId: 927500912

Port of GitHub PR: google#5927 Validate normalized relative paths in skill extraction code to prevent directory traversal via malicious GCS skill resource names. Co-authored-by: Yifan Wang <wanyif@google.com> PiperOrigin-RevId: 927503030

…sessions Co-authored-by: Kathy Wu <wukathy@google.com> PiperOrigin-RevId: 927552187

…pi_server Expose more configuration options in CLI and API server to support Gemini Enterprise: - Add `--trigger_sources` to `adk deploy agent_engine` and `adk api_server`. - Add `--express_mode` and `--gemini_enterprise_app_name` to `adk api_server`. - Add service URIs (`--artifact_service_uri`, `--memory_service_uri`, `--session_service_uri`) to `cli_deploy` and `api_server`. - Add `TopSpanProcessor` and `get_propagated_context` to telemetry for Agent Engine integration. - Update `pyproject.toml` dependencies. - Add and update unit tests for the new CLI options and telemetry. Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 927557476

Co-authored-by: Shangjie Chen <deanchen@google.com> PiperOrigin-RevId: 927571309

OpenTelemetry Semantic Conventions v1.41.0 establishes new standards for generative AI token usage telemetry (https://github.com/open-telemetry/semantic-conventions/blob/v1.41.0/docs/registry/attributes/gen-ai.md). Specifically, reasoning tokens SHOULD now be included in the total output token count (`gen_ai.usage.output_tokens`), with dedicated attributes for reasoning and cache read counts. This change updates ADK Python's tracing telemetry to align with these conventions and achieve parity with the recent adk-go implementation (b/508209702). Summary of changes: * `gen_ai.usage.output_tokens`: Semantics updated to include reasoning tokens (`candidates_token_count + thoughts_token_count`). * `gen_ai.usage.reasoning.output_tokens`: New attribute to record reasoning token usage (`thoughts_token_count`). * `gen_ai.usage.cache_read.input_tokens`: New attribute to record cached content input tokens (`cached_content_token_count`). * Deprecated `gen_ai.usage.experimental.reasoning_tokens` has been removed in favor of the standardized semconv attribute. PiperOrigin-RevId: 927772890

LiteLLM normally returns OpenAI-compatible tool call arguments as JSON strings, but some providers can stream a complete tool call whose finalized argument payload is a Python dict literal or has unquoted object keys. Keep strict JSON as the primary path, then repair only those complete object-literal shapes so ADK can still surface the intended function call. Co-authored-by: Yifan Wang <wanyif@google.com> PiperOrigin-RevId: 927943214

Closes google#5546 Co-authored-by: Yifan Wang <wanyif@google.com> PiperOrigin-RevId: 927960159

Closes google#4975 Co-authored-by: Yifan Wang <wanyif@google.com> PiperOrigin-RevId: 927975731

Co-authored-by: Shangjie Chen <deanchen@google.com> PiperOrigin-RevId: 928214336

Co-authored-by: Shangjie Chen <deanchen@google.com> PiperOrigin-RevId: 928259520

Closes google#5580 Co-authored-by: Yifan Wang <wanyif@google.com> PiperOrigin-RevId: 928268519

The local artifact service wrote to a single <agents_dir>/.adk/artifacts root shared by all agents, while session storage is per-agent under <agents_dir>/<agent>/.adk. Route artifacts the same way so each agent's artifacts live next to its session.db. For adk web/api_server this moves the default artifact location to <agents_dir>/<agent>/.adk/artifacts. To stay backward compatible, PerAgentFileArtifactService falls back to reading the old shared root on a miss, so existing artifacts keep loading; new writes go per-agent (copy-on-write migration). create_artifact_service_from_options logs a WARNING when a legacy shared root is found, pointing users at the manual move (the users/ dir into <agent>/.adk/artifacts). adk run is unaffected (it already resolved to the per-agent path). Also harden _resolve_agent_dir to use Path.is_relative_to instead of a string-prefix check, which accepted prefix-sharing siblings (e.g. app_name ../agents_evil under an agents root). This guards both the artifact and session per-agent paths. Co-authored-by: Yifan Wang <wanyif@google.com> PiperOrigin-RevId: 928307911

Co-authored-by: Wiktoria Walczak <wwalczak@google.com> PiperOrigin-RevId: 928428367

This change introduces a `TokenUsage` class to encapsulate the extraction and formatting of token usage metadata from `types.GenerateContentResponseUsageMetadata` into OpenTelemetry attributes. This class is used across `_experimental_semconv.py`, `_metrics.py`, and `tracing.py` to ensure consistent handling of token counts, including input, output, cached, reasoning, and experimental system instruction tokens. PiperOrigin-RevId: 928624507

Co-authored-by: Wei Sun (Jack) <weisun@google.com> PiperOrigin-RevId: 928672927

Co-authored-by: Shangjie Chen <deanchen@google.com> PiperOrigin-RevId: 928691597

Align cli_tools_click.py and cli_deploy.py with the GitHub main branch to ensure all express_mode validations, adk services options, and type annotations are fully synchronized. Co-authored-by: Shangjie Chen <deanchen@google.com> PiperOrigin-RevId: 928741460

Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 928746744

Re-aligns two files with GitHub main so the piper-to-github cutover does not silently revert them: the _ALLOWED_PICKLE_GLOBALS type annotation and a stray f-string prefix in skill_toolset. Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 928768999

adk-bot · 2026-06-08T23:08:01Z

I have completed the security and architectural triage for PR #5635.

The results of this analysis have been compiled into a premium pull request critique and are available in the PR Analysis Report.

Key Highlights & Decisions

Google CLA Sign-off: ✅ Verified as successfully signed.
Security Criticality: High-priority fix addressing a CVSS 9.9 vulnerability in the legacy v0 database schema runtime by replacing unrestricted pickle.loads occurrences in v0.py with a customized _RestrictedUnpickler allowlist.
Implementation Elegance: Excellent; covers both Sqlite engines (via result_processor) and MySQL/Spanner engines (via process_result_value). Supplying the ADK_ALLOW_UNSAFE_V0_PICKLE=1 environment escape hatch provides robust fallback capabilities for legacy systems without breaking normal operations.
Code & Style Alignment: Completely compliant with typing annotations, lazy logging semantics, and structure rules.
Recommendation: Approve (No changes required).

Summary of Work Done

CLA Verification: Inspected test suite logs and status check rollup via GitHub tools, demonstrating a SUCCESS status for the cla/google checks.
Diff Extraction: Fetched and parsed the line-by-line changes across the three modified files.
Workspace Inspection: Ran target lookups on the baseline v0.py and inspected the newly introduced migrate_from_sqlalchemy_pickle.py in the main branch to ensure alignment with existing migration unpickler patterns.
Compiled Report: Authored the comprehensive evaluation artifact outlining security impacts, justification, code quality alignments, and test coverages.

daridor9 · 2026-06-12T05:11:37Z

Hi @rohityan,

Just wanted to kindly follow up on this PR (#5635).

All previous feedback has been addressed (CLA signed, formatting/pyink + isort fixed, mypy type annotations added, scope narrowed after main took the migration file changes).
The PR has been kept up-to-date with main.
An internal PR Analysis Report already approved the changes with no modifications required.

Could you please take a look and let me know if anything else is needed? Happy to make any final adjustments.

Thanks for your time!

surajksharma07 mentioned this pull request May 7, 2026

Security: Unsafe pickle.loads() in v0 schema runtime path and migration tool #5634

Open

rohityan self-assigned this May 8, 2026

rohityan added services [Component] This issue is related to runtime services, e.g. sessions, memory, artifacts, etc request clarification [Status] The maintainer need clarification or more information from the author labels May 8, 2026

daridor9 and others added 11 commits June 4, 2026 07:03

security: add RestrictedUnpickler for safe v0 pickle deserialization

74daf17

Ref: google#5634

security: use RestrictedUnpickler in v0 runtime read path

41df699

Replace raw pickle.loads() with safe_loads() in DynamicPickleType.process_result_value(). Ref: google#5634

chore: retrigger CLA check after signing

99c79fb

security: add enum types to RestrictedUnpickler allowlist

e0ccf64

Covers potential enum values in state_delta/agent_state dict[str, Any] fields. Ref: google#5634

style: run pyink + isort on _safe_unpickle.py

7aeba1b

Ref: google#5634

style: run pyink + isort on v0.py

37b8fed

Ref: google#5634

style: run pyink + isort on migrate_from_sqlalchemy_pickle.py

fc1cad5

Ref: google#5634

security: override result_processor for SQLite safe deserialization

c5218fa

test: add real EventActions round-trip smoke tests

9f9c67e

Add tests Ref: google/adk-python#5634for EventActions serialization and deserialization.

fix(sessions): annotate result_processor for mypy-diff; pyink-format …

5c0ba06

…test file - Add Any type hints to DynamicPickleType.result_processor and its inner process() to clear mypy-diff [no-untyped-def]. - Reformat test_safe_unpickle.py to 2-space pyink style.

daridor9 force-pushed the fix/restricted-unpickler-v0 branch from 5f22404 to 5c0ba06 Compare June 4, 2026 04:04

daridor9 and others added 6 commits June 7, 2026 18:59

Merge branch 'main' into fix/restricted-unpickler-v0

705e2bf

ADK changes

8098d76

Co-authored-by: Sasha Sobran <asobran@google.com> PiperOrigin-RevId: 917516394

chore: update runtime config

ba69174

PiperOrigin-RevId: 917525676

chore: Remove experimental tag from SkillToolset

abdbb30

Co-authored-by: Kathy Wu <wukathy@google.com> PiperOrigin-RevId: 918037985

fix(cli): add missing newline at end of runtime-config.json

ddc0d82

Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 918503151

fix(deps): pin yarl to <1.24

8537ccc

Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 918539981

DeanChensj and others added 21 commits June 8, 2026 16:01

chore: 2.2.0 release

7ab188c

Bump version to 2.2.0. Co-authored-by: Shangjie Chen <deanchen@google.com> PiperOrigin-RevId: 927500912

fix(models): default grounding_metadata to empty for Gemini 3.1 live …

1321bf4

…sessions Co-authored-by: Kathy Wu <wukathy@google.com> PiperOrigin-RevId: 927552187

ADK changes

3e588cf

Co-authored-by: Shangjie Chen <deanchen@google.com> PiperOrigin-RevId: 927571309

feat: raise explicit error for unsupported LiteLlm file attachments

8251dba

Closes google#5546 Co-authored-by: Yifan Wang <wanyif@google.com> PiperOrigin-RevId: 927960159

fix: preserve media blocks in ollama content flattening

63ce7d3

Closes google#4975 Co-authored-by: Yifan Wang <wanyif@google.com> PiperOrigin-RevId: 927975731

chore: fix CI pre-commit

07a66b5

Co-authored-by: Shangjie Chen <deanchen@google.com> PiperOrigin-RevId: 928214336

No public description

b29caaa

Co-authored-by: Shangjie Chen <deanchen@google.com> PiperOrigin-RevId: 928259520

fix: validate session_id and enforce ownership in delete_session

8ce0181

Closes google#5580 Co-authored-by: Yifan Wang <wanyif@google.com> PiperOrigin-RevId: 928268519

fix: fix compaction tests visibility

ba4180f

Co-authored-by: Wiktoria Walczak <wwalczak@google.com> PiperOrigin-RevId: 928428367

docs(skills): Update skill file formatting

8fa47fb

Co-authored-by: Wei Sun (Jack) <weisun@google.com> PiperOrigin-RevId: 928672927

ADK changes

c9006ad

Co-authored-by: Shangjie Chen <deanchen@google.com> PiperOrigin-RevId: 928691597

fix: restore Gemini 3.1 Flash Live 1007-error handling on Vertex

aaa1bbb

Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 928746744

Merge branch 'main' into fix/restricted-unpickler-v0

e1b513b

Jacksunwei force-pushed the main branch from 5346305 to 0337d19 Compare June 8, 2026 23:31

daridor9 added 4 commits June 9, 2026 07:00

Merge branch 'main' into fix/restricted-unpickler-v0

60ef035

Merge branch 'main' into fix/restricted-unpickler-v0

f2e2f71

Merge branch 'main' into fix/restricted-unpickler-v0

6fe02fb

Merge branch 'main' into fix/restricted-unpickler-v0

f00bc63

daridor9 added 2 commits June 12, 2026 08:11

Merge branch 'main' into fix/restricted-unpickler-v0

fb2bdba

Merge branch 'main' into fix/restricted-unpickler-v0

2e02d4e

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

security: restrict pickle deserialization in v0 schema and migration tool#5635

security: restrict pickle deserialization in v0 schema and migration tool#5635
daridor9 wants to merge 103 commits into
google:mainfrom
daridor9:fix/restricted-unpickler-v0

daridor9 commented May 7, 2026

Uh oh!

google-cla Bot commented May 7, 2026

Uh oh!

adk-bot commented May 7, 2026

Uh oh!

rohityan commented May 8, 2026

Uh oh!

daridor9 commented May 11, 2026

Uh oh!

rohityan commented May 29, 2026

Uh oh!

daridor9 commented May 30, 2026 •

edited

Loading

Uh oh!

daridor9 commented Jun 3, 2026

Uh oh!

daridor9 commented Jun 4, 2026 •

edited

Loading

Uh oh!

adk-bot commented Jun 8, 2026

Uh oh!

daridor9 commented Jun 12, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

16 participants

Conversation

daridor9 commented May 7, 2026

Summary

Changes

New: _safe_unpickle.py

Patched: v0.py (Sink 1 — runtime read path)

Patched: migrate_from_sqlalchemy_pickle.py (Sink 2 — migration path)

Testing

Security Context

Uh oh!

google-cla Bot commented May 7, 2026

Uh oh!

adk-bot commented May 7, 2026

Uh oh!

rohityan commented May 8, 2026

Uh oh!

daridor9 commented May 11, 2026

Uh oh!

rohityan commented May 29, 2026

Uh oh!

daridor9 commented May 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

daridor9 commented Jun 3, 2026

Uh oh!

daridor9 commented Jun 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

adk-bot commented Jun 8, 2026

Key Highlights & Decisions

Summary of Work Done

Uh oh!

daridor9 commented Jun 12, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

16 participants

New: `_safe_unpickle.py`

Patched: `v0.py` (Sink 1 — runtime read path)

Patched: `migrate_from_sqlalchemy_pickle.py` (Sink 2 — migration path)

daridor9 commented May 30, 2026 •

edited

Loading

daridor9 commented Jun 4, 2026 •

edited

Loading