From 8536acf7b986c4e686f5715b95dde05514cfdb5f Mon Sep 17 00:00:00 2001 From: neville-m-exa Date: Wed, 29 Jul 2026 11:20:32 -0700 Subject: [PATCH] feat: support authorizations field in async_stream_query MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Currently, the `authorizations` field (for forwarding user credentials to agent tools via session state) is only supported by `streaming_agent_run_with_events`. When passed to `async_stream_query`, it flows through **kwargs to `Runner.run_async()` which rejects it: TypeError: Runner.run_async() got an unexpected keyword argument 'authorizations' This change adds explicit `authorizations` handling to `async_stream_query` — matching the behavior of `streaming_agent_run_with_events`: 1. Accept `authorizations` as a named parameter (Dict[str, Any]) 2. Convert each authorization's access_token into a state_delta entry 3. Pass state_delta to Runner.run_async() (which already supports it) This allows standard API callers to forward user credentials without switching to the AgentSpace-oriented method and its different request/response format. --- .../test_agent_engine_templates_adk.py | 112 ++++++++++++++++++ vertexai/agent_engines/templates/adk.py | 16 +++ 2 files changed, 128 insertions(+) diff --git a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py index c80e96a5b5..1534902f43 100644 --- a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py +++ b/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, diff --git a/vertexai/agent_engines/templates/adk.py b/vertexai/agent_engines/templates/adk.py index 25af45c2dd..5ea573215b 100644 --- a/vertexai/agent_engines/templates/adk.py +++ b/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,6 +1197,13 @@ 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( @@ -1197,6 +1211,7 @@ async def async_stream_query( session_id=session_id, new_message=content, run_config=run_config, + state_delta=state_delta, **kwargs, ) else: @@ -1204,6 +1219,7 @@ async def async_stream_query( user_id=user_id, session_id=session_id, new_message=content, + state_delta=state_delta, **kwargs, )