From 74daf177e4a79ee4912bc430e3f38db4a30c1d41 Mon Sep 17 00:00:00 2001 From: Dan Aridor Date: Thu, 7 May 2026 21:17:45 +0300 Subject: [PATCH 01/11] security: add RestrictedUnpickler for safe v0 pickle deserialization Ref: google/adk-python#5634 --- .../adk/sessions/schemas/_safe_unpickle.py | 84 +++++++++++++++++++ 1 file changed, 84 insertions(+) create mode 100644 src/google/adk/sessions/schemas/_safe_unpickle.py diff --git a/src/google/adk/sessions/schemas/_safe_unpickle.py b/src/google/adk/sessions/schemas/_safe_unpickle.py new file mode 100644 index 0000000000..9018953692 --- /dev/null +++ b/src/google/adk/sessions/schemas/_safe_unpickle.py @@ -0,0 +1,84 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Restricted unpickler for safe deserialization of v0 EventActions data. + +The v0 schema stored EventActions as pickle blobs. This module provides a +safe deserialization path that only allows known ADK and standard types, +blocking arbitrary code execution via crafted pickle payloads. + +See: https://docs.python.org/3/library/pickle.html#restricting-globals +""" + +from __future__ import annotations + +import io +import logging +import os +import pickle +from typing import Any + +logger = logging.getLogger("google_adk." + __name__) + +_ALLOWED_MODULE_PREFIXES: tuple[str, ...] = ( + "google.adk.", + "google.genai.", + "pydantic.", + "pydantic_core.", +) + +_ALLOWED_GLOBALS: dict[str, set[str]] = { + "builtins": { + "dict", "list", "set", "tuple", "frozenset", "bytes", + "bytearray", "True", "False", "None", "type", "object", + "complex", "slice", "range", "int", "float", "str", "bool", + }, + "collections": {"OrderedDict", "defaultdict"}, + "datetime": {"datetime", "date", "time", "timedelta", "timezone"}, + "copy_reg": {"_reconstructor"}, + "copyreg": {"_reconstructor", "__newobj__"}, + "_codecs": {"encode"}, +} + + +class _RestrictedUnpickler(pickle.Unpickler): + """Unpickler that only allows reconstruction of known-safe types.""" + + def find_class(self, module: str, name: str) -> Any: + for prefix in _ALLOWED_MODULE_PREFIXES: + if module.startswith(prefix): + return super().find_class(module, name) + allowed_names = _ALLOWED_GLOBALS.get(module) + if allowed_names and name in allowed_names: + return super().find_class(module, name) + raise pickle.UnpicklingError( + f"Blocked unsafe pickle global: {module}.{name}. " + f"If this is a legitimate ADK type, please file an issue at " + f"https://github.com/google/adk-python/issues" + ) + + +def safe_loads(data: bytes) -> Any: + """Deserialize pickle bytes using a restricted unpickler. + + If ADK_ALLOW_UNSAFE_V0_PICKLE=1 is set, falls back to unrestricted + pickle.loads() for compatibility. A deprecation warning is logged. + """ + if os.environ.get("ADK_ALLOW_UNSAFE_V0_PICKLE") == "1": + logger.warning( + "ADK_ALLOW_UNSAFE_V0_PICKLE is set - using unrestricted " + "pickle.loads(). This is unsafe and will be removed in a " + "future release. Migrate to the v1 JSON schema." + ) + return pickle.loads(data) # noqa: S301 + return _RestrictedUnpickler(io.BytesIO(data)).load() From 41df6998a265fa42639156276bf88c8b2617463b Mon Sep 17 00:00:00 2001 From: Dan Aridor Date: Thu, 7 May 2026 21:17:59 +0300 Subject: [PATCH 02/11] security: use RestrictedUnpickler in v0 runtime read path Replace raw pickle.loads() with safe_loads() in DynamicPickleType.process_result_value(). Ref: google/adk-python#5634 --- src/google/adk/sessions/schemas/v0.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/google/adk/sessions/schemas/v0.py b/src/google/adk/sessions/schemas/v0.py index e4a4368c6d..62ee72b607 100644 --- a/src/google/adk/sessions/schemas/v0.py +++ b/src/google/adk/sessions/schemas/v0.py @@ -31,6 +31,8 @@ import json import logging import pickle + +from ._safe_unpickle import safe_loads as _safe_pickle_loads from typing import Any from typing import Optional @@ -114,7 +116,7 @@ def process_result_value(self, value, dialect): """Ensures the raw bytes from the database are unpickled back into a Python object.""" if value is not None: if dialect.name in ("spanner+spanner", "mysql"): - return pickle.loads(value) + return _safe_pickle_loads(value) return value From 99c79fbe3810b2f0d61bc930b11f2c4eb40898c5 Mon Sep 17 00:00:00 2001 From: Dan Aridor Date: Thu, 7 May 2026 21:22:54 +0300 Subject: [PATCH 03/11] chore: retrigger CLA check after signing From e0ccf64c22cca3c1f86f81981c4a21599ad54612 Mon Sep 17 00:00:00 2001 From: Dan Aridor Date: Thu, 7 May 2026 22:50:56 +0300 Subject: [PATCH 04/11] security: add enum types to RestrictedUnpickler allowlist Covers potential enum values in state_delta/agent_state dict[str, Any] fields. Ref: google/adk-python#5634 --- src/google/adk/sessions/schemas/_safe_unpickle.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/google/adk/sessions/schemas/_safe_unpickle.py b/src/google/adk/sessions/schemas/_safe_unpickle.py index 9018953692..8309462d4b 100644 --- a/src/google/adk/sessions/schemas/_safe_unpickle.py +++ b/src/google/adk/sessions/schemas/_safe_unpickle.py @@ -48,6 +48,7 @@ "copy_reg": {"_reconstructor"}, "copyreg": {"_reconstructor", "__newobj__"}, "_codecs": {"encode"}, + "enum": {"__new__", "Enum", "IntEnum", "StrEnum"}, } From 7aeba1b0e4d6c4fef99e3d1a69401bccdf31a564 Mon Sep 17 00:00:00 2001 From: Dan Aridor Date: Sat, 9 May 2026 10:10:19 +0300 Subject: [PATCH 05/11] style: run pyink + isort on _safe_unpickle.py Ref: google/adk-python#5634 --- .../adk/sessions/schemas/_safe_unpickle.py | 72 +++++++++++-------- 1 file changed, 44 insertions(+), 28 deletions(-) diff --git a/src/google/adk/sessions/schemas/_safe_unpickle.py b/src/google/adk/sessions/schemas/_safe_unpickle.py index 8309462d4b..a61fcfb5b6 100644 --- a/src/google/adk/sessions/schemas/_safe_unpickle.py +++ b/src/google/adk/sessions/schemas/_safe_unpickle.py @@ -39,9 +39,25 @@ _ALLOWED_GLOBALS: dict[str, set[str]] = { "builtins": { - "dict", "list", "set", "tuple", "frozenset", "bytes", - "bytearray", "True", "False", "None", "type", "object", - "complex", "slice", "range", "int", "float", "str", "bool", + "dict", + "list", + "set", + "tuple", + "frozenset", + "bytes", + "bytearray", + "True", + "False", + "None", + "type", + "object", + "complex", + "slice", + "range", + "int", + "float", + "str", + "bool", }, "collections": {"OrderedDict", "defaultdict"}, "datetime": {"datetime", "date", "time", "timedelta", "timezone"}, @@ -53,33 +69,33 @@ class _RestrictedUnpickler(pickle.Unpickler): - """Unpickler that only allows reconstruction of known-safe types.""" + """Unpickler that only allows reconstruction of known-safe types.""" - def find_class(self, module: str, name: str) -> Any: - for prefix in _ALLOWED_MODULE_PREFIXES: - if module.startswith(prefix): - return super().find_class(module, name) - allowed_names = _ALLOWED_GLOBALS.get(module) - if allowed_names and name in allowed_names: - return super().find_class(module, name) - raise pickle.UnpicklingError( - f"Blocked unsafe pickle global: {module}.{name}. " - f"If this is a legitimate ADK type, please file an issue at " - f"https://github.com/google/adk-python/issues" - ) + def find_class(self, module: str, name: str) -> Any: + for prefix in _ALLOWED_MODULE_PREFIXES: + if module.startswith(prefix): + return super().find_class(module, name) + allowed_names = _ALLOWED_GLOBALS.get(module) + if allowed_names and name in allowed_names: + return super().find_class(module, name) + raise pickle.UnpicklingError( + f"Blocked unsafe pickle global: {module}.{name}. " + "If this is a legitimate ADK type, please file an issue at " + "https://github.com/google/adk-python/issues" + ) def safe_loads(data: bytes) -> Any: - """Deserialize pickle bytes using a restricted unpickler. + """Deserialize pickle bytes using a restricted unpickler. - If ADK_ALLOW_UNSAFE_V0_PICKLE=1 is set, falls back to unrestricted - pickle.loads() for compatibility. A deprecation warning is logged. - """ - if os.environ.get("ADK_ALLOW_UNSAFE_V0_PICKLE") == "1": - logger.warning( - "ADK_ALLOW_UNSAFE_V0_PICKLE is set - using unrestricted " - "pickle.loads(). This is unsafe and will be removed in a " - "future release. Migrate to the v1 JSON schema." - ) - return pickle.loads(data) # noqa: S301 - return _RestrictedUnpickler(io.BytesIO(data)).load() + If ADK_ALLOW_UNSAFE_V0_PICKLE=1 is set, falls back to unrestricted + pickle.loads() for compatibility. A deprecation warning is logged. + """ + if os.environ.get("ADK_ALLOW_UNSAFE_V0_PICKLE") == "1": + logger.warning( + "ADK_ALLOW_UNSAFE_V0_PICKLE is set - using unrestricted " + "pickle.loads(). This is unsafe and will be removed in a " + "future release. Migrate to the v1 JSON schema." + ) + return pickle.loads(data) # noqa: S301 + return _RestrictedUnpickler(io.BytesIO(data)).load() From 37b8fed27c2a6e162cddf31891b910c399fd382f Mon Sep 17 00:00:00 2001 From: Dan Aridor Date: Sat, 9 May 2026 10:10:21 +0300 Subject: [PATCH 06/11] style: run pyink + isort on v0.py Ref: google/adk-python#5634 --- src/google/adk/sessions/schemas/v0.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/google/adk/sessions/schemas/v0.py b/src/google/adk/sessions/schemas/v0.py index 62ee72b607..5b378cb7a8 100644 --- a/src/google/adk/sessions/schemas/v0.py +++ b/src/google/adk/sessions/schemas/v0.py @@ -31,8 +31,6 @@ import json import logging import pickle - -from ._safe_unpickle import safe_loads as _safe_pickle_loads from typing import Any from typing import Optional @@ -59,6 +57,7 @@ from ...events.event import Event from ...events.event_actions import EventActions from ..session import Session +from ._safe_unpickle import safe_loads as _safe_pickle_loads from .shared import DEFAULT_MAX_KEY_LENGTH from .shared import DEFAULT_MAX_VARCHAR_LENGTH from .shared import DynamicJSON From fc1cad565e542b35f994b6670c981d7d749df5c2 Mon Sep 17 00:00:00 2001 From: Dan Aridor Date: Sat, 9 May 2026 10:10:22 +0300 Subject: [PATCH 07/11] style: run pyink + isort on migrate_from_sqlalchemy_pickle.py Ref: google/adk-python#5634 --- .../adk/sessions/migration/migrate_from_sqlalchemy_pickle.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py b/src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py index 65a78c9401..334fd9b907 100644 --- a/src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py +++ b/src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py @@ -30,6 +30,7 @@ from google.adk.sessions import _session_util from google.adk.sessions.migration import _schema_check_utils from google.adk.sessions.schemas import v1 +from google.adk.sessions.schemas._safe_unpickle import safe_loads as _safe_pickle_loads from google.genai import types import sqlalchemy from sqlalchemy import create_engine From c5218fa21a4dcc8c361c5169cf6d0f1b3a233d80 Mon Sep 17 00:00:00 2001 From: Dan Aridor Date: Mon, 11 May 2026 18:16:06 +0300 Subject: [PATCH 08/11] security: override result_processor for SQLite safe deserialization --- src/google/adk/sessions/schemas/v0.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/src/google/adk/sessions/schemas/v0.py b/src/google/adk/sessions/schemas/v0.py index 5b378cb7a8..486791b790 100644 --- a/src/google/adk/sessions/schemas/v0.py +++ b/src/google/adk/sessions/schemas/v0.py @@ -111,6 +111,19 @@ def process_bind_param(self, value, dialect): return pickle.dumps(value) return value + def result_processor(self, dialect, coltype): + if dialect.name in ("mysql", "spanner+spanner"): + return super().result_processor(dialect, coltype) + + def process(value): + if value is None: + return None + if isinstance(value, memoryview): + value = bytes(value) + return _safe_pickle_loads(value) + + return process + def process_result_value(self, value, dialect): """Ensures the raw bytes from the database are unpickled back into a Python object.""" if value is not None: From a9b863ecdfc23f7b80b0d473fd67aef8d7600ae1 Mon Sep 17 00:00:00 2001 From: Dan Aridor Date: Mon, 11 May 2026 18:22:14 +0300 Subject: [PATCH 09/11] Add unit tests for RestrictedUnpickler functionality 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. --- .../unittests/sessions/test_safe_unpickle.py | 139 ++++++++++++++++++ 1 file changed, 139 insertions(+) create mode 100644 tests/unittests/sessions/test_safe_unpickle.py diff --git a/tests/unittests/sessions/test_safe_unpickle.py b/tests/unittests/sessions/test_safe_unpickle.py new file mode 100644 index 0000000000..e28c98ed12 --- /dev/null +++ b/tests/unittests/sessions/test_safe_unpickle.py @@ -0,0 +1,139 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for _safe_unpickle RestrictedUnpickler.""" + +from __future__ import annotations + +import io +import os +import pickle +import struct +import unittest + +from google.adk.sessions.schemas._safe_unpickle import safe_loads + + +def _make_global_payload(module: str, func: str, *args: str) -> bytes: + """Craft a pickle stream that calls module.func(*args).""" + buf = io.BytesIO() + buf.write(pickle.PROTO + struct.pack("B", 2)) + buf.write(b"c" + f"{module}\n{func}\n".encode()) + buf.write(b"(") + for arg in args: + encoded = arg.encode("utf-8") + buf.write( + pickle.SHORT_BINUNICODE + struct.pack(" safe_loads.""" + + def _round_trip(self, obj): + return safe_loads(pickle.dumps(obj)) + + def test_string_values(self): + original = {"state_delta": {"key": "value"}, "artifact_delta": {}} + self.assertEqual(self._round_trip(original), original) + + def test_nested_dict(self): + original = { + "state_delta": { + "user_prefs": {"theme": "dark", "lang": "en"}, + "counter": 42, + }, + "artifact_delta": {"files": ["a.txt", "b.txt"]}, + } + self.assertEqual(self._round_trip(original), original) + + def test_none_and_bool(self): + original = { + "skip_summarization": True, + "requested_auth_configs": None, + "escalate": False, + } + self.assertEqual(self._round_trip(original), original) + + def test_empty_dict(self): + self.assertEqual(self._round_trip({}), {}) + + +class TestEnvVarFallback(unittest.TestCase): + """ADK_ALLOW_UNSAFE_V0_PICKLE=1 must bypass RestrictedUnpickler.""" + + _ENV_KEY = "ADK_ALLOW_UNSAFE_V0_PICKLE" + _PAYLOAD = _make_global_payload("collections", "Counter") + + def test_blocked_without_env_var(self): + old = os.environ.pop(self._ENV_KEY, None) + try: + with self.assertRaises(pickle.UnpicklingError): + safe_loads(self._PAYLOAD) + finally: + if old is not None: + os.environ[self._ENV_KEY] = old + + def test_allowed_with_env_var(self): + old = os.environ.get(self._ENV_KEY) + try: + os.environ[self._ENV_KEY] = "1" + from collections import Counter + result = safe_loads(self._PAYLOAD) + self.assertIsInstance(result, Counter) + finally: + if old is None: + os.environ.pop(self._ENV_KEY, None) + else: + os.environ[self._ENV_KEY] = old + + +if __name__ == "__main__": + unittest.main() From 9f9c67e9db18a6cd616c571cf2cb62a16516b0a2 Mon Sep 17 00:00:00 2001 From: Dan Aridor Date: Mon, 18 May 2026 20:48:45 +0300 Subject: [PATCH 10/11] test: add real EventActions round-trip smoke tests Add tests Ref: google/adk-python#5634for EventActions serialization and deserialization. --- .../unittests/sessions/test_safe_unpickle.py | 49 +++++++++++++++++++ 1 file changed, 49 insertions(+) diff --git a/tests/unittests/sessions/test_safe_unpickle.py b/tests/unittests/sessions/test_safe_unpickle.py index e28c98ed12..f1b8d5fbd9 100644 --- a/tests/unittests/sessions/test_safe_unpickle.py +++ b/tests/unittests/sessions/test_safe_unpickle.py @@ -22,6 +22,7 @@ import struct import unittest +from google.adk.events.event_actions import EventActions from google.adk.sessions.schemas._safe_unpickle import safe_loads @@ -105,6 +106,54 @@ def test_none_and_bool(self): def test_empty_dict(self): self.assertEqual(self._round_trip({}), {}) +class TestRealEventActionsRoundTrip(unittest.TestCase): + """Smoke test: real EventActions instances survive pickle -> safe_loads.""" + + def _round_trip(self, obj): + return safe_loads(pickle.dumps(obj)) + + def test_minimal_event_actions(self): + original = EventActions() + result = self._round_trip(original) + self.assertIsInstance(result, EventActions) + self.assertEqual(result.state_delta, {}) + self.assertEqual(result.artifact_delta, {}) + + def test_event_actions_with_state_delta(self): + original = EventActions( + state_delta={"user_name": "alice", "turn_count": 3, "active": True}, + artifact_delta={"report.pdf": 2}, + ) + result = self._round_trip(original) + self.assertIsInstance(result, EventActions) + self.assertEqual(result.state_delta, original.state_delta) + self.assertEqual(result.artifact_delta, original.artifact_delta) + + def test_event_actions_with_transfer_and_escalate(self): + original = EventActions( + transfer_to_agent="specialist_agent", + escalate=True, + skip_summarization=True, + ) + result = self._round_trip(original) + self.assertIsInstance(result, EventActions) + self.assertEqual(result.transfer_to_agent, "specialist_agent") + self.assertTrue(result.escalate) + self.assertTrue(result.skip_summarization) + + def test_event_actions_with_complex_state_values(self): + original = EventActions( + state_delta={ + "nested": {"a": [1, 2, 3], "b": None}, + "count": 42, + "tags": ["ml", "security"], + }, + ) + result = self._round_trip(original) + self.assertIsInstance(result, EventActions) + self.assertEqual(result.state_delta["nested"]["a"], [1, 2, 3]) + self.assertIsNone(result.state_delta["nested"]["b"]) + class TestEnvVarFallback(unittest.TestCase): """ADK_ALLOW_UNSAFE_V0_PICKLE=1 must bypass RestrictedUnpickler.""" From 5c0ba0692f4b62d665deda779afdc6bdff1dcdf1 Mon Sep 17 00:00:00 2001 From: Dan Aridor Date: Sat, 30 May 2026 08:13:30 +0300 Subject: [PATCH 11/11] fix(sessions): annotate result_processor for mypy-diff; pyink-format 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. --- .../migrate_from_sqlalchemy_pickle.py | 1 - src/google/adk/sessions/schemas/v0.py | 4 +- .../unittests/sessions/test_safe_unpickle.py | 280 +++++++++--------- 3 files changed, 143 insertions(+), 142 deletions(-) diff --git a/src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py b/src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py index 334fd9b907..65a78c9401 100644 --- a/src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py +++ b/src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py @@ -30,7 +30,6 @@ from google.adk.sessions import _session_util from google.adk.sessions.migration import _schema_check_utils from google.adk.sessions.schemas import v1 -from google.adk.sessions.schemas._safe_unpickle import safe_loads as _safe_pickle_loads from google.genai import types import sqlalchemy from sqlalchemy import create_engine diff --git a/src/google/adk/sessions/schemas/v0.py b/src/google/adk/sessions/schemas/v0.py index 486791b790..ddd8b5d39f 100644 --- a/src/google/adk/sessions/schemas/v0.py +++ b/src/google/adk/sessions/schemas/v0.py @@ -111,11 +111,11 @@ def process_bind_param(self, value, dialect): return pickle.dumps(value) return value - def result_processor(self, dialect, coltype): + def result_processor(self, dialect: Any, coltype: Any) -> Any: if dialect.name in ("mysql", "spanner+spanner"): return super().result_processor(dialect, coltype) - def process(value): + def process(value: Any) -> Any: if value is None: return None if isinstance(value, memoryview): diff --git a/tests/unittests/sessions/test_safe_unpickle.py b/tests/unittests/sessions/test_safe_unpickle.py index f1b8d5fbd9..7e58522b47 100644 --- a/tests/unittests/sessions/test_safe_unpickle.py +++ b/tests/unittests/sessions/test_safe_unpickle.py @@ -27,162 +27,164 @@ def _make_global_payload(module: str, func: str, *args: str) -> bytes: - """Craft a pickle stream that calls module.func(*args).""" - buf = io.BytesIO() - buf.write(pickle.PROTO + struct.pack("B", 2)) - buf.write(b"c" + f"{module}\n{func}\n".encode()) - buf.write(b"(") - for arg in args: - encoded = arg.encode("utf-8") - buf.write( - pickle.SHORT_BINUNICODE + struct.pack(" safe_loads.""" - - def _round_trip(self, obj): - return safe_loads(pickle.dumps(obj)) - - def test_string_values(self): - original = {"state_delta": {"key": "value"}, "artifact_delta": {}} - self.assertEqual(self._round_trip(original), original) - - def test_nested_dict(self): - original = { - "state_delta": { - "user_prefs": {"theme": "dark", "lang": "en"}, - "counter": 42, - }, - "artifact_delta": {"files": ["a.txt", "b.txt"]}, - } - self.assertEqual(self._round_trip(original), original) - - def test_none_and_bool(self): - original = { - "skip_summarization": True, - "requested_auth_configs": None, - "escalate": False, - } - self.assertEqual(self._round_trip(original), original) - - def test_empty_dict(self): - self.assertEqual(self._round_trip({}), {}) + """Legitimate EventActions data must survive pickle -> safe_loads.""" + + def _round_trip(self, obj): + return safe_loads(pickle.dumps(obj)) + + def test_string_values(self): + original = {"state_delta": {"key": "value"}, "artifact_delta": {}} + self.assertEqual(self._round_trip(original), original) + + def test_nested_dict(self): + original = { + "state_delta": { + "user_prefs": {"theme": "dark", "lang": "en"}, + "counter": 42, + }, + "artifact_delta": {"files": ["a.txt", "b.txt"]}, + } + self.assertEqual(self._round_trip(original), original) + + def test_none_and_bool(self): + original = { + "skip_summarization": True, + "requested_auth_configs": None, + "escalate": False, + } + self.assertEqual(self._round_trip(original), original) + + def test_empty_dict(self): + self.assertEqual(self._round_trip({}), {}) + class TestRealEventActionsRoundTrip(unittest.TestCase): - """Smoke test: real EventActions instances survive pickle -> safe_loads.""" - - def _round_trip(self, obj): - return safe_loads(pickle.dumps(obj)) - - def test_minimal_event_actions(self): - original = EventActions() - result = self._round_trip(original) - self.assertIsInstance(result, EventActions) - self.assertEqual(result.state_delta, {}) - self.assertEqual(result.artifact_delta, {}) - - def test_event_actions_with_state_delta(self): - original = EventActions( - state_delta={"user_name": "alice", "turn_count": 3, "active": True}, - artifact_delta={"report.pdf": 2}, - ) - result = self._round_trip(original) - self.assertIsInstance(result, EventActions) - self.assertEqual(result.state_delta, original.state_delta) - self.assertEqual(result.artifact_delta, original.artifact_delta) - - def test_event_actions_with_transfer_and_escalate(self): - original = EventActions( - transfer_to_agent="specialist_agent", - escalate=True, - skip_summarization=True, - ) - result = self._round_trip(original) - self.assertIsInstance(result, EventActions) - self.assertEqual(result.transfer_to_agent, "specialist_agent") - self.assertTrue(result.escalate) - self.assertTrue(result.skip_summarization) - - def test_event_actions_with_complex_state_values(self): - original = EventActions( - state_delta={ - "nested": {"a": [1, 2, 3], "b": None}, - "count": 42, - "tags": ["ml", "security"], - }, - ) - result = self._round_trip(original) - self.assertIsInstance(result, EventActions) - self.assertEqual(result.state_delta["nested"]["a"], [1, 2, 3]) - self.assertIsNone(result.state_delta["nested"]["b"]) + """Smoke test: real EventActions instances survive pickle -> safe_loads.""" + + def _round_trip(self, obj): + return safe_loads(pickle.dumps(obj)) + + def test_minimal_event_actions(self): + original = EventActions() + result = self._round_trip(original) + self.assertIsInstance(result, EventActions) + self.assertEqual(result.state_delta, {}) + self.assertEqual(result.artifact_delta, {}) + + def test_event_actions_with_state_delta(self): + original = EventActions( + state_delta={"user_name": "alice", "turn_count": 3, "active": True}, + artifact_delta={"report.pdf": 2}, + ) + result = self._round_trip(original) + self.assertIsInstance(result, EventActions) + self.assertEqual(result.state_delta, original.state_delta) + self.assertEqual(result.artifact_delta, original.artifact_delta) + + def test_event_actions_with_transfer_and_escalate(self): + original = EventActions( + transfer_to_agent="specialist_agent", + escalate=True, + skip_summarization=True, + ) + result = self._round_trip(original) + self.assertIsInstance(result, EventActions) + self.assertEqual(result.transfer_to_agent, "specialist_agent") + self.assertTrue(result.escalate) + self.assertTrue(result.skip_summarization) + + def test_event_actions_with_complex_state_values(self): + original = EventActions( + state_delta={ + "nested": {"a": [1, 2, 3], "b": None}, + "count": 42, + "tags": ["ml", "security"], + }, + ) + result = self._round_trip(original) + self.assertIsInstance(result, EventActions) + self.assertEqual(result.state_delta["nested"]["a"], [1, 2, 3]) + self.assertIsNone(result.state_delta["nested"]["b"]) class TestEnvVarFallback(unittest.TestCase): - """ADK_ALLOW_UNSAFE_V0_PICKLE=1 must bypass RestrictedUnpickler.""" - - _ENV_KEY = "ADK_ALLOW_UNSAFE_V0_PICKLE" - _PAYLOAD = _make_global_payload("collections", "Counter") - - def test_blocked_without_env_var(self): - old = os.environ.pop(self._ENV_KEY, None) - try: - with self.assertRaises(pickle.UnpicklingError): - safe_loads(self._PAYLOAD) - finally: - if old is not None: - os.environ[self._ENV_KEY] = old - - def test_allowed_with_env_var(self): - old = os.environ.get(self._ENV_KEY) - try: - os.environ[self._ENV_KEY] = "1" - from collections import Counter - result = safe_loads(self._PAYLOAD) - self.assertIsInstance(result, Counter) - finally: - if old is None: - os.environ.pop(self._ENV_KEY, None) - else: - os.environ[self._ENV_KEY] = old + """ADK_ALLOW_UNSAFE_V0_PICKLE=1 must bypass RestrictedUnpickler.""" + + _ENV_KEY = "ADK_ALLOW_UNSAFE_V0_PICKLE" + _PAYLOAD = _make_global_payload("collections", "Counter") + + def test_blocked_without_env_var(self): + old = os.environ.pop(self._ENV_KEY, None) + try: + with self.assertRaises(pickle.UnpicklingError): + safe_loads(self._PAYLOAD) + finally: + if old is not None: + os.environ[self._ENV_KEY] = old + + def test_allowed_with_env_var(self): + old = os.environ.get(self._ENV_KEY) + try: + os.environ[self._ENV_KEY] = "1" + from collections import Counter + + result = safe_loads(self._PAYLOAD) + self.assertIsInstance(result, Counter) + finally: + if old is None: + os.environ.pop(self._ENV_KEY, None) + else: + os.environ[self._ENV_KEY] = old if __name__ == "__main__": - unittest.main() + unittest.main()