Skip to content

fix(migration): restrict unpickling of v0 actions blobs#5866

Open
White-Mouse wants to merge 4 commits into
google:mainfrom
White-Mouse:codex/restrict-migration-unpickler
Open

fix(migration): restrict unpickling of v0 actions blobs#5866
White-Mouse wants to merge 4 commits into
google:mainfrom
White-Mouse:codex/restrict-migration-unpickler

Conversation

@White-Mouse
Copy link
Copy Markdown

@White-Mouse White-Mouse commented May 27, 2026

What

The v0 session schema stored event actions as pickled blobs. The migration helper reads raw bytes via SELECT * FROM events and previously used pickle.loads(...) directly.

This PR replaces the default load path with a restricted unpickler allowlist for builtin containers/primitives, standard ADK EventActions payloads, nested ADK core action types (AuthConfig, ToolConfirmation, EventCompaction), and the google.genai.types.Content / Part dependency classes that normal compaction payloads require.

It also adds an explicit trust toggle for legacy databases that contain custom Python objects in state_delta or other Any fields:

  • Python API: migrate(..., allow_unsafe_unpickling=True)
  • Migration runner: upgrade(..., allow_unsafe_unpickling=True)
  • CLI: adk migrate session --allow_unsafe_unpickling ...
  • Direct script: --allow_unsafe_unpickling / --allow-unsafe-unpickling

Why

pickle is not safe for untrusted inputs. Migration tooling often runs against restored/backed-up DB files or shared storage; failing closed by default reduces the blast radius if the source DB contents are compromised.

The opt-in flag keeps compatibility for users who trust their source database and need the original unsafe pickle behavior for custom legacy objects.

Associated Issue / Background

No existing GitHub issue is linked. This was found while reviewing the v0-to-v1 migration path for unsafe deserialization risks in legacy session data.

Compatibility / fail-closed boundary

Normal v0 EventActions payloads made from primitive/container fields continue to migrate. The allowlist now also covers common nested ADK action models requested during review, including requested auth configs, requested tool confirmations, and event compaction content.

Payloads that require globals outside the explicit allowlist still fail closed by default: the migration logs a warning and falls back to empty EventActions() for that event. Users can opt into the previous unsafe pickle behavior only when they trust the source database.

Verification

  • uv run pytest tests/unittests/sessions/migration/test_migration.py
    • 22 passed, 4 warnings
  • uv run mypy src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py src/google/adk/sessions/migration/migration_runner.py
    • Success: no issues found in 2 source files
  • uv run pre-commit run --files src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py src/google/adk/sessions/migration/migration_runner.py src/google/adk/cli/cli_tools_click.py tests/unittests/sessions/migration/test_migration.py
    • passed
  • git diff --check
    • passed

@google-cla
Copy link
Copy Markdown

google-cla Bot commented May 27, 2026

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 adk-bot added the services [Component] This issue is related to runtime services, e.g. sessions, memory, artifacts, etc label May 27, 2026
@adk-bot
Copy link
Copy Markdown
Collaborator

adk-bot commented May 27, 2026

Response from ADK Triaging Agent

Hello @White-Mouse, thank you for submitting this pull request!

To help us review and process your contribution more efficiently, please address the following items in accordance with our contribution guidelines:

  1. Contributor License Agreement (CLA): It looks like the CLA check has failed. Please make sure you (or your employer) have signed the Google CLA so we can proceed with the review.
  2. Testing / Verification Plan: The Verification section in your PR description is currently empty. Please update this section with your testing plan or a summary of your test results (e.g., confirming unit tests under tests/unittests/sessions/migration/test_migration.py passed).
  3. Associated Issue: If there is an existing GitHub issue related to this bug/fix, please link it in your PR description. If not, consider creating one or briefly describing the background context.

Thank you for your cooperation!

@White-Mouse White-Mouse force-pushed the codex/restrict-migration-unpickler branch from d09e22e to 71c4053 Compare May 27, 2026 14:45
@rohityan rohityan self-assigned this May 27, 2026
@wyf7107 wyf7107 requested a review from DeanChensj May 27, 2026 19:00
@rohityan
Copy link
Copy Markdown
Collaborator

Hi @White-Mouse , Thank you for your contribution! We appreciate you taking the time to submit this pull request. Can you please fix the failing mypy-diff tests before we can proceed with the review.

@rohityan rohityan added the request clarification [Status] The maintainer need clarification or more information from the author label May 27, 2026
@DeanChensj
Copy link
Copy Markdown
Collaborator

DeanChensj commented May 28, 2026

Hi @White-Mouse, thanks for the update! I appreciate the effort to move towards a "Secure by Default" approach using a restricted unpickler.

To ensure this doesn't negatively impact existing users, I'd like to suggest the following improvements to the PR:

  1. Prevent Data Loss by Expanding the Allowlist:
    The current restricted unpickler only allows EventActions. However, EventActions frequently contains nested ADK core types such as AuthConfig , ToolConfirmation, EventCompaction. Since EventCompaction also relies on google.genai.types.Content, we should include these classes (and their dependencies) in the allowlist. This ensures standard sessions migrate fully without falling back to an empty EventActions().

  2. Add an Optional "Trust" Toggle:
    Since the state_delta field can store arbitrary Python objects (Any), some users may have custom classes in their legacy data. We should provide an optional configuration (e.g., allow_unsafe_unpickling: bool = False) that allows users to explicitly opt-in to the original unpickling behavior if they trust their source database.

With these two changes, we can achieve strong security for the majority of users while maintaining full functionality and flexibility. What do you think?

@White-Mouse
Copy link
Copy Markdown
Author

White-Mouse commented May 28, 2026

Thanks for the review. I updated the PR to address both points: the restricted allowlist now covers standard nested ADK action types (AuthConfig, ToolConfirmation, EventCompaction) plus the google.genai.types.Content / Part dependency classes used by compaction payloads, and migration now has an explicit allow_unsafe_unpickling=False trust toggle for users who need to migrate custom legacy objects from a trusted source database.

I also fixed the mypy-diff blocker by adding the missing return annotation and ran the updated migration tests, targeted mypy, pre-commit, and git diff --check locally. The new pull_request workflow runs for mypy/unit/pre-commit/check-file are currently in GitHub's action_required state for this fork PR, so they need maintainer approval before they will execute on the new commit.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

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

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants