feat: support authorizations field in async_stream_query

tests/unit/vertex_adk/test_agent_engine_templates_adk.py

@@ -429,6 +429,118 @@ async def test_async_stream_query(
             events.append(event)
         assert len(events) == 1
 
+    @pytest.mark.asyncio
+    async def test_async_stream_query_with_authorizations(
+        self,
+        default_instrumentor_builder_mock: mock.Mock,
+        get_project_id_mock: mock.Mock,
+    ):
+        app = agent_engines.AdkApp(agent=_TEST_AGENT)
+        app.set_up()
+        mock_runner = mock.AsyncMock()
+
+        async def mock_run_async(**kwargs):
+            from google.adk.events import event
+
+            yield event.Event(
+                **{
+                    "author": "test_agent",
+                    "content": {"parts": [{"text": "ok"}], "role": "model"},
+                    "id": "test-id",
+                    "invocation_id": "e-test",
+                }
+            )
+
+        mock_runner.run_async = mock_run_async
+        app._tmpl_attrs["runner"] = mock_runner
+
+        events = []
+        async for event in app.async_stream_query(
+            user_id=_TEST_USER_ID,
+            message="test message",
+            authorizations={
+                "exa_mcp_access": {"access_token": "test-token-123"},
+                "other_service": {"accessToken": "other-token-456"},
+            },
+        ):
+            events.append(event)
+        assert len(events) == 1
+
+    @pytest.mark.asyncio
+    async def test_async_stream_query_authorizations_passed_as_state_delta(
+        self,
+        default_instrumentor_builder_mock: mock.Mock,
+        get_project_id_mock: mock.Mock,
+    ):
+        app = agent_engines.AdkApp(agent=_TEST_AGENT)
+        app.set_up()
+
+        captured_kwargs = {}
+
+        async def capturing_run_async(**kwargs):
+            captured_kwargs.update(kwargs)
+            from google.adk.events import event
+
+            yield event.Event(
+                **{
+                    "author": "test_agent",
+                    "content": {"parts": [{"text": "ok"}], "role": "model"},
+                    "id": "test-id",
+                    "invocation_id": "e-test",
+                }
+            )
+
+        app._tmpl_attrs["runner"].run_async = capturing_run_async
+
+        events = []
+        async for event in app.async_stream_query(
+            user_id=_TEST_USER_ID,
+            message="test message",
+            authorizations={
+                "exa_mcp_access": {"access_token": "test-token-123"},
+            },
+        ):
+            events.append(event)
+
+        assert captured_kwargs["state_delta"] == {
+            "exa_mcp_access": "test-token-123",
+        }
+
+    @pytest.mark.asyncio
+    async def test_async_stream_query_no_authorizations_passes_none_state_delta(
+        self,
+        default_instrumentor_builder_mock: mock.Mock,
+        get_project_id_mock: mock.Mock,
+    ):
+        app = agent_engines.AdkApp(agent=_TEST_AGENT)
+        app.set_up()
+
+        captured_kwargs = {}
+
+        async def capturing_run_async(**kwargs):
+            captured_kwargs.update(kwargs)
+            from google.adk.events import event
+
+            yield event.Event(
+                **{
+                    "author": "test_agent",
+                    "content": {"parts": [{"text": "ok"}], "role": "model"},
+                    "id": "test-id",
+                    "invocation_id": "e-test",
+                }
+            )
+
+        app._tmpl_attrs["runner"].run_async = capturing_run_async
+
+        events = []
+        async for event in app.async_stream_query(
+            user_id=_TEST_USER_ID,
+            message="test message",
+        ):
+            events.append(event)
+
+        assert captured_kwargs["state_delta"] is None
+
     def test_set_up_runner_auto_create_session_enabled(
         self,
         default_instrumentor_builder_mock: mock.Mock,

vertexai/agent_engines/templates/adk.py

@@ -1123,6 +1123,7 @@ async def async_stream_query(
         session_id: Optional[str] = None,
         session_events: Optional[List[Dict[str, Any]]] = None,
         run_config: Optional[Dict[str, Any]] = None,
+        authorizations: Optional[Dict[str, Any]] = None,
         **kwargs,
     ) -> AsyncIterable[Dict[str, Any]]:
         """Streams responses asynchronously from the ADK application.
@@ -1143,6 +1144,12 @@ async def async_stream_query(
                 Optional. The run config to use for the query. If you want to
                 pass in a `run_config` pydantic object, you can pass in a dict
                 representing it as `run_config.model_dump(mode="json")`.
+            authorizations (Optional[Dict[str, Any]]):
+                Optional. The authorizations of the user, keyed by
+                authorization ID. Each value should be a dictionary with an
+                ``access_token`` (or ``accessToken``) field. The tokens will
+                be written to the session state so that agent tools can access
+                them at runtime via ``ctx.state[auth_id]``.
             **kwargs (dict[str, Any]):
                 Optional. Additional keyword arguments to pass to the
                 runner.
@@ -1190,20 +1197,29 @@ async def async_stream_query(
                         event=event,
                     )
 
+        state_delta = None
+        if authorizations:
+            state_delta = {}
+            for auth_id, auth in authorizations.items():
+                auth_obj = _Authorization(**auth) if isinstance(auth, dict) else auth
+                state_delta[auth_id] = auth_obj.access_token
+
         run_config = _validate_run_config(run_config)
         if run_config:
             events_async = self._tmpl_attrs.get("runner").run_async(
                 user_id=user_id,
                 session_id=session_id,
                 new_message=content,
                 run_config=run_config,
+                state_delta=state_delta,
                 **kwargs,
             )
         else:
             events_async = self._tmpl_attrs.get("runner").run_async(
                 user_id=user_id,
                 session_id=session_id,
                 new_message=content,
+                state_delta=state_delta,
                 **kwargs,
             )