fix(adk): carry donkey.last_call into after_model_callback on google-adk 2.13

python/src/donkey_kit/integrations/adk.py

@@ -28,22 +28,32 @@
   work. The default DDK proxies are
   ``Format=OpenAI``, so point ``base_url`` at a ``Format=Gemini`` proxy.
 
+Both factories return a subclass of the ADK class that carries
+``donkey.last_call`` back to the context that started the model call (#1207).
+google-adk 2.13.0 runs each step of the model's response generator in a copy of
+the flow's context, so without it the record stays in that copy and an
+``after_model_callback`` reads ``UNOBSERVED``.
+
 Class names / kwargs UNVERIFIED — docs/verified-apis.md §8.
 """
 
 from __future__ import annotations
 
 import functools
+from contextlib import aclosing
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Any
 
 from ..core import _verify
+from ..core.lastcall import LastCallBridge, open_last_call_bridge
 from ..core.masking import masked
 from ..core.refusals import capture_refusals
 from . import AdapterCapabilities
 from ._base import Adapter, default_adapter
 
 if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator
+
     from google.adk.models import Gemini
     from google.adk.models.lite_llm import LiteLlm, LiteLLMClient
 
@@ -116,7 +126,8 @@ def gemini_connection_kwargs(self, *, base_url: str | None = None) -> dict[str,
         yourself, bound to a ``Format=Gemini`` proxy (``base_url``, default the
         configured proxy URL).
 
-        ADK replaces its own ``http_options`` with ``client_kwargs``, so every
+        ADK builds its client from ``client_kwargs`` (from google-adk 2.13.0 it
+        adds its tracking headers and fills only the fields left unset), so every
         governed value rides there: the shared http client (header injection,
         retries, last_call), the consumer-auth headers, the ``api_key`` slot
         google-genai requires for the Gemini API backend, and the SDK timeout —
@@ -149,7 +160,11 @@ def model(self, model: str, **kw: Any) -> LiteLlm:
         Its ``llm_client`` is a ``LiteLLMClient`` subclass that lets a gateway
         refusal reach ``donkey.run()`` as its typed class (#969, see the module
         docstring). An ``llm_client`` you pass replaces it, and the refusal
-        then stays LiteLLM's error."""
+        then stays LiteLLM's error.
+
+        The object is a ``LiteLlm`` subclass that carries ``donkey.last_call``
+        back to the context that started each model call, so an
+        ``after_model_callback`` reads it (#1207)."""
         self._allow_endpoints(kw, "api_base", "base_url")
         with self._native_import():
             from google.adk.models.lite_llm import (  # VERIFY name/path: docs/verified-apis.md §8
@@ -163,13 +178,18 @@ def model(self, model: str, **kw: Any) -> LiteLlm:
             # LiteLLM sends to the client's own base URL, so rebuild it there.
             conn["client"] = self._proxy_openai_client(str(override))
         conn["llm_client"] = _governed_llm_client_class(LiteLLMClient)()
+        governed: type[LiteLlm] = _last_call_bridged_class(LiteLlm)
         # LiteLLM's OpenAI-compatible route needs the ``openai/`` prefix.
-        return LiteLlm(model=f"openai/{model}", **{**conn, **kw})
+        return governed(model=f"openai/{model}", **{**conn, **kw})
 
     def gemini(self, model: str, *, base_url: str | None = None, **kw: Any) -> Gemini:
         """Return ADK's native ``google.adk.models.Gemini`` bound to a
         ``Format=Gemini`` proxy, with the shared http client injected (#691).
-        Pass the bare model id (``"gemini-2.5-flash"``) — no provider prefix."""
+        Pass the bare model id (``"gemini-2.5-flash"``) — no provider prefix.
+
+        The object is a ``Gemini`` subclass that carries ``donkey.last_call``
+        back to the context that started each model call, so an
+        ``after_model_callback`` reads it (#1207, see the module docstring)."""
         conn = self.gemini_connection_kwargs(base_url=base_url)
         with self._native_import():
             from google.adk.models import Gemini  # VERIFY name/path: docs/verified-apis.md §8
@@ -189,7 +209,59 @@ def gemini(self, model: str, *, base_url: str | None = None, **kw: Any) -> Gemin
                 "older versions ADK drops the governed client without an error. "
                 "Upgrade with: pip install -U 'donkey-kit[adk]'"
             )
-        return Gemini(model=model, **{**conn, **kw})
+        governed: type[Gemini] = _last_call_bridged_class(Gemini)
+        return governed(model=model, **{**conn, **kw})
+
+
+@functools.cache
+def _last_call_bridged_class(base: Any) -> Any:
+    """The subclass of an ADK model class (``base``: ``Gemini`` or ``LiteLlm``)
+    that :meth:`ADKAdapter.gemini` and :meth:`ADKAdapter.model` build, defined
+    on first use so ``google.adk`` is imported lazily (§1.1).
+
+    Only the public ``BaseLlm.generate_content_async`` is overridden. ADK's flow
+    calls it in its own context and then iterates the generator it returns.
+    From google-adk 2.13.0 it iterates that generator in a copy of the context
+    (``_stream_in_caller_scope``), so the record the transport sets stays in the
+    copy, and the ``after_model_callback`` that ADK runs in the flow's context
+    reads ``UNOBSERVED`` (docs/verified-apis.md §8). The override is a plain
+    ``def``, so it opens a :class:`~donkey_kit.core.lastcall.LastCallBridge`
+    in the flow's context before ADK takes the copy. The copy inherits the
+    bridge, and the transport records through it.
+
+    The bridge serves one model call in one task, so hazard #2 holds: a
+    ``ParallelAgent`` branch runs in a task of its own and opens its own bridge.
+    It closes before the generator yields a complete (non-partial) response,
+    because ADK runs that response's function calls before it resumes the
+    generator, and a tool that ADK runs in a task of its own must not write
+    through the bridge. A partial response keeps it open, since ADK runs no
+    function calls for one."""
+
+    class LastCallBridged(base):
+        """An ADK model whose ``donkey.last_call`` reaches the context that
+        started the call (#1207)."""
+
+        def generate_content_async(self, *args: Any, **kwargs: Any) -> AsyncGenerator[Any, None]:
+            bridge = open_last_call_bridge()
+            return _bridged_responses(super().generate_content_async(*args, **kwargs), bridge)
+
+    LastCallBridged.__name__ = LastCallBridged.__qualname__ = f"Governed{base.__name__}"
+    return LastCallBridged
+
+
+async def _bridged_responses(
+    responses: AsyncGenerator[Any, None], bridge: LastCallBridge
+) -> AsyncGenerator[Any, None]:
+    """Yield ``responses`` and close ``bridge`` before the first complete
+    response, or when the call ends early (see :func:`_last_call_bridged_class`)."""
+    try:
+        async with aclosing(responses):
+            async for response in responses:
+                if not getattr(response, "partial", False):
+                    bridge.close()
+                yield response
+    finally:
+        bridge.close()
 
 
 @functools.cache

python/tests/unit/test_adapter_async_last_call.py

@@ -14,20 +14,27 @@
 * The OpenAI Agents SDK ``Runner`` sends a run's first turn from a task of its
   own, and ADK's ``Runner`` runs the whole agent in one. Neither runs anything
   the adapter controls in the caller's context, so there is nowhere to open a
-  :class:`~donkey_kit.core.lastcall.LastCallBridge`. Bridging the whole run
-  instead would break hazard #2, since a run can fan out (agents as tools, ADK's
-  ``ParallelAgent``). The gap is documented on the framework pages, with the
-  in-task hook that reads the record (``RunHooks.on_llm_end``,
+  :class:`~donkey_kit.core.lastcall.LastCallBridge` for the caller. Bridging the
+  whole run instead would break hazard #2, since a run can fan out (agents as
+  tools, ADK's ``ParallelAgent``). The gap is documented on the framework pages,
+  with the in-task hook that reads the record (``RunHooks.on_llm_end``,
   ``after_model_callback``). The ``UNOBSERVED`` assertions below pin that
   documented behaviour: if one starts failing, the framework changed where it
   sends the call, and the page needs updating.
+* From google-adk 2.13.0, ADK's flow iterates the model's response generator in
+  a copy of its own context, so both ADK factories return a model subclass that
+  bridges each model call back to the flow, where ``after_model_callback`` runs
+  (#1207). The callback tests below pin that for both factories. The bridge
+  closes before a complete response, so a task started on one (as ADK starts
+  tool calls) records in its own context (hazard #2).
 
 Framework-free at module level (the base-only job); each test skips unless its
 framework is installed. CI runs this file in the jobs that install them.
 """
 
 from __future__ import annotations
 
+import asyncio
 import json
 from pathlib import Path
 from typing import Any
@@ -252,3 +259,112 @@ def after_model(callback_context: Any, llm_response: Any) -> None:
         # Documented gap: ADK's Runner runs the agent in a task of its own.
         assert donkey.last_call.status is LastCallStatus.UNOBSERVED
     assert in_callback == [LastCallStatus.OBSERVED]
+
+
+async def test_adk_model_runner_is_unobserved_but_after_model_callback_sees_it() -> None:
+    pytest.importorskip("google.adk")
+    pytest.importorskip("litellm")
+    from google.adk.agents import LlmAgent
+    from google.adk.runners import InMemoryRunner
+
+    donkey = _donkey()
+    in_callback: list[str | None] = []
+
+    def after_model(callback_context: Any, llm_response: Any) -> None:
+        in_callback.append(donkey.last_call.request_id)
+
+    async with donkey:
+        agent = LlmAgent(
+            name="a",
+            model=donkey.adk.model("gpt-4o"),
+            after_model_callback=after_model,
+        )
+        runner = InMemoryRunner(agent=agent, app_name="app")
+        session = await runner.session_service.create_session(app_name="app", user_id="u")
+        async for _ in runner.run_async(
+            user_id="u", session_id=session.id, new_message=_gemini_message()
+        ):
+            pass
+        # Documented gap: ADK's Runner runs the agent in a task of its own.
+        assert donkey.last_call.status is LastCallStatus.UNOBSERVED
+    assert in_callback == ["rid-answer"]
+
+
+def _gemini_stream(request: httpx.Request) -> httpx.Response:
+    return httpx.Response(
+        200,
+        content=(GEMINI_FIXTURES / "responses.stream.body.txt").read_bytes(),
+        headers=parse_headers((GEMINI_FIXTURES / "responses.stream.headers.txt").read_text()),
+    )
+
+
+async def test_adk_streamed_run_after_model_callback_sees_every_response() -> None:
+    """The bridge stays open across partial responses and its record survives
+    the close, so each callback, the final complete one included, reads it."""
+    pytest.importorskip("google.adk")
+    from google.adk.agents import LlmAgent
+    from google.adk.agents.run_config import RunConfig, StreamingMode
+    from google.adk.runners import InMemoryRunner
+
+    donkey = _donkey(_gemini_stream)
+    in_callback: list[tuple[bool, LastCallStatus]] = []
+
+    def after_model(callback_context: Any, llm_response: Any) -> None:
+        in_callback.append((bool(llm_response.partial), donkey.last_call.status))
+
+    async with donkey:
+        agent = LlmAgent(
+            name="a",
+            model=donkey.adk.gemini("gemini-2.5-flash"),
+            after_model_callback=after_model,
+        )
+        runner = InMemoryRunner(agent=agent, app_name="app")
+        session = await runner.session_service.create_session(app_name="app", user_id="u")
+        async for _ in runner.run_async(
+            user_id="u",
+            session_id=session.id,
+            new_message=_gemini_message(),
+            run_config=RunConfig(streaming_mode=StreamingMode.SSE),
+        ):
+            pass
+    observed = LastCallStatus.OBSERVED
+    assert in_callback == [(True, observed)] * 3 + [(False, observed)]
+
+
+async def test_adk_factories_return_the_adk_classes() -> None:
+    pytest.importorskip("google.adk")
+    pytest.importorskip("litellm")
+    from google.adk.models import Gemini
+    from google.adk.models.lite_llm import LiteLlm
+
+    donkey = _donkey()
+    async with donkey:
+        assert isinstance(donkey.adk.gemini("gemini-2.5-flash"), Gemini)
+        assert isinstance(donkey.adk.model("gpt-4o"), LiteLlm)
+
+
+async def test_adk_bridge_closes_before_a_complete_response() -> None:
+    """A task started while ADK handles a complete response, as ADK starts a tool
+    call, records in its own context and not through the model call's bridge."""
+    pytest.importorskip("google.adk")
+    from google.adk.models.llm_request import LlmRequest
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        if request.url.path.endswith("/chat/completions"):
+            return _answer(request)
+        return _gemini(request)
+
+    donkey = _donkey(handler)
+
+    async def sibling() -> None:
+        await donkey.llm.client().chat.completions.create(
+            model="gpt-4o", messages=[{"role": "user", "content": "hi"}]
+        )
+
+    async with donkey:
+        request = LlmRequest(model="gemini-2.5-flash", contents=[_gemini_message()])
+        model = donkey.adk.gemini("gemini-2.5-flash")
+        async for response in model.generate_content_async(request):
+            if not response.partial:
+                await asyncio.create_task(sibling())
+        assert donkey.last_call.requested_model == "gemini-2.5-flash"

website/content/frameworks/adk.mdx

@@ -50,7 +50,9 @@ from donkey_kit.integrations.adk import model
 llm = model("gpt-4o")
 ```
 
-`llm` is a real `google.adk.models.lite_llm.LiteLlm` instance. The model string
+`llm` is a real `google.adk.models.lite_llm.LiteLlm` instance (a subclass that
+carries `donkey.last_call` to your callbacks; see
+[Native Gemini](#native-gemini)). The model string
 is prefixed with `openai/` before it reaches LiteLLM (`openai/gpt-4o`), which
 is the prefix LiteLLM's OpenAI-compatible route expects — you don't add it
 yourself.
@@ -190,8 +192,10 @@ llm = Gemini(
 )
 ```
 
-`client_kwargs` replaces ADK's default HTTP options wholesale, so all five
-`http_options` keys are passed together. `google-genai` requires an API key and
+All five `http_options` keys are passed together, because ADK builds its client
+from `client_kwargs`. The `base_url`, `api_version` and `headers` you set are
+kept, and ADK adds its own `x-goog-api-client` and `user-agent` tracking values
+to the headers. `google-genai` requires an API key and
 always sends it as `x-goog-api-key`; the gateway authenticates on the
 `client_id`/`client_secret` pair and ignores it.
 
@@ -206,7 +210,10 @@ populated.
 
 `last_call` is contextvar-scoped, and ADK's `Runner` makes the model call in a
 task of its own, so the caller of `runner.run_async(...)` reads `UNOBSERVED`.
-Read it in an `after_model_callback`, which runs in the same task as the call:
+Read it in an `after_model_callback`. Both `gemini()` and `model()` return a
+subclass of the ADK class that carries the record from the model call to the
+callback, including on `google-adk` 2.13 and later, which runs the model call in
+a copy of the agent's context:
 
 ```python
 from google.adk.agents import LlmAgent
@@ -252,7 +259,10 @@ run in `asyncio.run(...)` can build the `Donkey` (or call the module-level
   correlation ID bound by `donkey.run()` reaches every request.
 - **`donkey.last_call` is set in the context that made the call.** ADK's
   `Runner` calls the model in a task of its own, so read it in an
-  `after_model_callback` (see [Native Gemini](#native-gemini)). A cold read on
+  `after_model_callback` (see [Native Gemini](#native-gemini)). The factories'
+  models carry the record to that callback. A `LiteLlm` or `Gemini` you build
+  from `connection_kwargs()` or `gemini_connection_kwargs()` does not, so on
+  `google-adk` 2.13 and later its callback reads `UNOBSERVED`. A cold read on
   a `Donkey` that resolved only ADK reports `UNOBSERVED`, and `OBSERVED` after
   a call, for `model()` as well as `gemini()`.
 - **Refusals on the `model()` path are typed inside `donkey.run()` and