From 41203c8320f017f7590481e820ab6022ec469d22 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Fri, 24 Jul 2026 10:45:14 -0700 Subject: [PATCH] feat: Forward per-request labels to RunConfig in streaming_agent_run_with_events PiperOrigin-RevId: 953446344 --- agentplatform/agent_engines/templates/adk.py | 13 +++++ .../frameworks/test_frameworks_adk.py | 56 ++++++++++++++++++ .../test_agent_engine_templates_adk.py | 57 +++++++++++++++++++ vertexai/agent_engines/templates/adk.py | 13 +++++ 4 files changed, 139 insertions(+) diff --git a/agentplatform/agent_engines/templates/adk.py b/agentplatform/agent_engines/templates/adk.py index e7b1b60001..dfedd6785b 100644 --- a/agentplatform/agent_engines/templates/adk.py +++ b/agentplatform/agent_engines/templates/adk.py @@ -236,6 +236,11 @@ def __init__(self, **kwargs): ) # The session ID. + self.labels: Optional[Dict[str, str]] = kwargs.get("labels") + # Per-request user labels (e.g. for billing/attribution) to attach to + # the RunConfig for this invocation, so they are propagated onto the + # downstream Vertex API requests. + class _StreamingRunResponse: """Response object for `streaming_agent_run_with_events` method. @@ -1414,12 +1419,20 @@ async def streaming_agent_run_with_events(self, request_json: str): # Run the agent message_for_agent = types.Content(**request.message) + # Propagate per-request user labels (e.g. billing/attribution) onto a + # RunConfig so the ADK flow forwards them to downstream Vertex requests. + run_config = None + if request.labels: + from google.adk.agents.run_config import RunConfig + + run_config = RunConfig(labels=request.labels) try: async for event in runner.run_async( user_id=request.user_id, session_id=session.id, new_message=message_for_agent, state_delta=state_delta, + run_config=run_config, ): converted_event = await self._convert_response_events( user_id=request.user_id, diff --git a/tests/unit/agentplatform/frameworks/test_frameworks_adk.py b/tests/unit/agentplatform/frameworks/test_frameworks_adk.py index 7e3d5dba6f..b66ef2369e 100644 --- a/tests/unit/agentplatform/frameworks/test_frameworks_adk.py +++ b/tests/unit/agentplatform/frameworks/test_frameworks_adk.py @@ -655,6 +655,62 @@ async def test_streaming_agent_run_with_events( events.append(event) assert len(events) == 1 + @pytest.mark.asyncio + async def test_streaming_agent_run_with_events_propagates_labels( + self, + default_instrumentor_builder_mock: mock.Mock, + get_project_id_mock: mock.Mock, + ): + app = adk_template.AdkApp(agent=_TEST_AGENT) + app.set_up() + + # Pre-create a session in the real in-memory session service. + await app.async_create_session( + user_id=_TEST_USER_ID, session_id="test_session_id" + ) + + runner_mock = mock.Mock() + + async def mock_run_async(*args, **kwargs): + from google.adk.events import event + + yield event.Event( + **{ + "author": "currency_exchange_agent", + "content": {"parts": [{"text": "Sweden"}], "role": "model"}, + "id": "9aaItGK9", + "invocation_id": "e-6543c213-6417-484b-9551-b67915d1d5f7", + } + ) + + spy = mock.MagicMock(side_effect=mock_run_async) + runner_mock.run_async = spy + app._tmpl_attrs["runner"] = runner_mock + + labels = {"goog-originating-logical-product-id": "prod1"} + request_json = json.dumps( + { + "user_id": _TEST_USER_ID, + "session_id": "test_session_id", + "message": { + "parts": [{"text": "What is the exchange rate from USD to SEK?"}], + "role": "user", + }, + "labels": labels, + } + ) + + async for _ in app.streaming_agent_run_with_events( + request_json=request_json, + ): + pass + + # The labels are surfaced to run_async via a RunConfig. + spy.assert_called_once() + run_config = spy.call_args.kwargs["run_config"] + assert run_config is not None + assert run_config.labels == labels + @pytest.mark.asyncio @mock.patch.dict( os.environ, 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 044852b575..ae2d46efa8 100644 --- a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py +++ b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py @@ -572,8 +572,65 @@ async def mock_run_async(*args, **kwargs): session_id="test_session_id", new_message=mock.ANY, state_delta={"test_user_id1": "test_access_token"}, + run_config=None, ) + @pytest.mark.asyncio + async def test_streaming_agent_run_with_events_propagates_labels( + self, + default_instrumentor_builder_mock: mock.Mock, + get_project_id_mock: mock.Mock, + ): + app = agent_engines.AdkApp(agent=_TEST_AGENT) + app.set_up() + + # Pre-create a session in the real in-memory session service. + await app.async_create_session( + user_id=_TEST_USER_ID, session_id="test_session_id" + ) + + runner_mock = mock.Mock() + + async def mock_run_async(*args, **kwargs): + from google.adk.events import event + + yield event.Event( + **{ + "author": "currency_exchange_agent", + "content": {"parts": [{"text": "Sweden"}], "role": "model"}, + "id": "9aaItGK9", + "invocation_id": "e-6543c213-6417-484b-9551-b67915d1d5f7", + } + ) + + spy = mock.MagicMock(side_effect=mock_run_async) + runner_mock.run_async = spy + app._tmpl_attrs["runner"] = runner_mock + + labels = {"goog-originating-logical-product-id": "prod1"} + request_json = json.dumps( + { + "user_id": _TEST_USER_ID, + "session_id": "test_session_id", + "message": { + "parts": [{"text": "What is the exchange rate from USD to SEK?"}], + "role": "user", + }, + "labels": labels, + } + ) + + async for _ in app.streaming_agent_run_with_events( + request_json=request_json, + ): + pass + + # The labels are surfaced to run_async via a RunConfig. + spy.assert_called_once() + run_config = spy.call_args.kwargs["run_config"] + assert run_config is not None + assert run_config.labels == labels + @pytest.mark.asyncio @mock.patch.dict( os.environ, diff --git a/vertexai/agent_engines/templates/adk.py b/vertexai/agent_engines/templates/adk.py index a90140dc3c..a38771b590 100644 --- a/vertexai/agent_engines/templates/adk.py +++ b/vertexai/agent_engines/templates/adk.py @@ -220,6 +220,11 @@ def __init__(self, **kwargs): ) # The session ID. + self.labels: Optional[Dict[str, str]] = kwargs.get("labels") + # Per-request user labels (e.g. for billing/attribution) to attach to + # the RunConfig for this invocation, so they are propagated onto the + # downstream Vertex API requests. + class _StreamingRunResponse: """Response object for `streaming_agent_run_with_events` method. @@ -1374,12 +1379,20 @@ async def streaming_agent_run_with_events(self, request_json: str): # Run the agent message_for_agent = types.Content(**request.message) + # Propagate per-request user labels (e.g. billing/attribution) onto a + # RunConfig so the ADK flow forwards them to downstream Vertex requests. + run_config = None + if request.labels: + from google.adk.agents.run_config import RunConfig + + run_config = RunConfig(labels=request.labels) try: async for event in runner.run_async( user_id=request.user_id, session_id=session.id, new_message=message_for_agent, state_delta=state_delta, + run_config=run_config, ): converted_event = await self._convert_response_events( user_id=request.user_id,