fix(adk): plain bridge marker and picklable Governed* classes (#1447)

python/src/donkey_kit/integrations/adk.py

@@ -1,3 +1,3 @@
 """Google ADK adapter (BG §1.8).
 
 Supported at connection_kwargs() — not conformance-tested (BG §1.8).
@@ -34,7 +34,12 @@
 the flow's context, so without it the record stays in that copy and an
 ``after_model_callback`` reads ``UNOBSERVED``. :func:`last_call_bridged` returns
 that subclass for a model built from ``connection_kwargs()`` or
-``gemini_connection_kwargs()`` (#1218).
+``gemini_connection_kwargs()`` (#1218). The factories' classes,
+``GovernedLiteLlm``, ``GovernedGemini`` and the ``GovernedLiteLLMClient`` that
+``model()`` passes as ``llm_client``, resolve by name on this module, so pickle
+finds them. A model the factories build still holds the SDK's http client, so
+it does not pickle; nor does a plain ``LiteLlm`` given the same kwargs. A class
+:func:`last_call_bridged` builds from your own class has no name here.
 
 APPROVALS (BG §2.3, #1235): an approval-gated tool
 (``@donkey.tool(approval="required")`` in an agent's ``tools``) pauses the run
@@ -68,7 +73,7 @@
 from collections.abc import Iterable, Mapping
 from contextlib import aclosing
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Any, TypeVar
+from typing import TYPE_CHECKING, Any, ClassVar, TypeVar
 
 from ..core import _verify
 from ..core.approvals import Approval, record_decision, record_pending
@@ -317,7 +322,9 @@
         """An ADK model whose ``donkey.last_call`` reaches the context that
         started the call (#1207)."""
 
-        _donkey_last_call_bridged = True
+        # A ClassVar, so pydantic leaves it a plain class attribute rather than
+        # turning it into a private attribute.
+        _donkey_last_call_bridged: ClassVar[bool] = True
 
         def generate_content_async(self, *args: Any, **kwargs: Any) -> AsyncGenerator[Any, None]:
             bridge = open_last_call_bridge()
@@ -368,9 +375,47 @@
                     capture.tag(exc)
                     raise
 
+    GovernedLiteLLMClient.__qualname__ = GovernedLiteLLMClient.__name__
     return GovernedLiteLLMClient
 
 
+def _governed_lite_llm() -> Any:
+    from google.adk.models.lite_llm import LiteLlm
+
+    return _last_call_bridged_class(LiteLlm)
+
+
+def _governed_gemini() -> Any:
+    from google.adk.models import Gemini
+
+    return _last_call_bridged_class(Gemini)
+
+
+def _governed_lite_llm_client() -> Any:
+    from google.adk.models.lite_llm import LiteLLMClient
+
+    return _governed_llm_client_class(LiteLLMClient)
+
+
+_GOVERNED_CLASSES = {
+    "GovernedLiteLlm": _governed_lite_llm,
+    "GovernedGemini": _governed_gemini,
+    "GovernedLiteLLMClient": _governed_lite_llm_client,
+}
+
+if not TYPE_CHECKING:  # mypy would otherwise accept any name on this module
+
+    def __getattr__(name: str) -> Any:
+        """The factories' classes by name, so pickle finds them (PEP 562)."""
+        build = _GOVERNED_CLASSES.get(name)
+        if build is None:
+            raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+        try:
+            return build()
+        except ModuleNotFoundError as exc:
+            raise missing_framework_error("adk", exc.name) from exc
+
+
 def model(model: str, **kw: Any) -> LiteLlm:
     """Module-level convenience: a native ``LiteLlm`` at the proxy using a cached
     default env-configured Donkey. Equivalent to

python/tests/unit/test_adapter_async_last_call.py

@@ -36,6 +36,7 @@
 
 import asyncio
 import json
+import pickle
 from pathlib import Path
 from typing import Any
 
@@ -459,6 +460,45 @@ class Mine(Gemini):
             last_call_bridged(wrong)  # type: ignore[arg-type]
 
 
+def test_the_bridge_marker_is_a_plain_class_attribute() -> None:
+    """Not a pydantic private attribute, which is a truthy ``ModelPrivateAttr``."""
+    pytest.importorskip("google.adk")
+    pytest.importorskip("litellm")
+    from google.adk.models import Gemini
+    from google.adk.models.lite_llm import LiteLlm
+
+    from donkey_kit.integrations.adk import last_call_bridged
+
+    for base in (Gemini, LiteLlm):
+        bridged = last_call_bridged(base)
+        assert bridged.__dict__["_donkey_last_call_bridged"] is True
+        assert "_donkey_last_call_bridged" not in bridged.__private_attributes__
+
+
+async def test_the_factory_classes_resolve_by_name_so_they_pickle() -> None:
+    pytest.importorskip("google.adk")
+    pytest.importorskip("litellm")
+    from google.adk.models import Gemini
+    from google.adk.models.lite_llm import LiteLlm
+
+    from donkey_kit.integrations import adk
+
+    donkey = _donkey()
+    async with donkey:
+        lite = donkey.adk.model("gpt-4o")
+        gemini = donkey.adk.gemini("gemini-3.8-flash")
+    assert type(lite) is adk.GovernedLiteLlm
+    assert type(lite.llm_client) is adk.GovernedLiteLLMClient
+    assert type(gemini) is adk.GovernedGemini
+    for cls in (adk.GovernedLiteLlm, adk.GovernedGemini, adk.GovernedLiteLLMClient):
+        assert pickle.loads(pickle.dumps(cls)) is cls
+    plain = adk.last_call_bridged(LiteLlm)(model="openai/gpt-4o")
+    assert type(pickle.loads(pickle.dumps(plain))) is adk.GovernedLiteLlm
+    assert adk.last_call_bridged(Gemini) is adk.GovernedGemini
+    with pytest.raises(AttributeError, match="GovernedNothing"):
+        adk.GovernedNothing  # noqa: B018
+
+
 async def test_adk_factories_return_the_adk_classes() -> None:
     pytest.importorskip("google.adk")
     pytest.importorskip("litellm")