feat: Forward per-request labels to RunConfig in streaming_agent_run_with_events

PiperOrigin-RevId: 953446344

Author: vertex-sdk-bot

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,

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,

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,

tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py

@@ -601,6 +601,37 @@ def test_streaming_agent_run_with_events(self):
         events = list(app.streaming_agent_run_with_events(request_json=request_json))
         assert len(events) == 1
 
+    def test_streaming_agent_run_with_events_propagates_labels(self):
+        captured = {}
+
+        class _LabelCapturingRunner(_MockRunner):
+            def run(self, *args, **kwargs):
+                captured["run_config"] = kwargs.get("run_config")
+                yield from super().run(*args, **kwargs)
+
+        app = reasoning_engines.AdkApp(
+            agent=Agent(name=_TEST_AGENT_NAME, model=_TEST_MODEL)
+        )
+        app.set_up()
+        app._tmpl_attrs["in_memory_runner"] = _LabelCapturingRunner()
+
+        labels = {"goog-originating-logical-product-id": "prod1"}
+        request_json = json.dumps(
+            {
+                "user_id": _TEST_USER_ID,
+                "message": {
+                    "parts": [{"text": "What is the exchange rate from USD to SEK?"}],
+                    "role": "user",
+                },
+                "labels": labels,
+            }
+        )
+        events = list(app.streaming_agent_run_with_events(request_json=request_json))
+
+        assert len(events) == 1
+        assert captured["run_config"] is not None
+        assert captured["run_config"].labels == labels
+
     @pytest.mark.asyncio
     @mock.patch.dict(
         os.environ,

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,

vertexai/preview/reasoning_engines/templates/adk.py

@@ -223,6 +223,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.
@@ -1196,11 +1201,20 @@ async def _invoke_agent_async():
                 raise RuntimeError("Session initialization failed.")
             # 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:
                 for event in runner.run(
                     user_id=request.user_id,
                     session_id=session.id,
                     new_message=message_for_agent,
+                    run_config=run_config,
                 ):
                     converted_event = await self._convert_response_events(
                         user_id=request.user_id,