diff --git a/src/google/adk/cli/__init__.py b/src/google/adk/cli/__init__.py index 1c4740c2e10..80aaa8abe70 100644 --- a/src/google/adk/cli/__init__.py +++ b/src/google/adk/cli/__init__.py @@ -19,11 +19,15 @@ from ..utils import _lazy if TYPE_CHECKING: + from ._sse import format_event_frames + from ._sse import serialize_event from .cli_tools_click import main _LAZY_MEMBERS: dict[str, str] = { + 'format_event_frames': '._sse', 'main': '.cli_tools_click', + 'serialize_event': '._sse', } -__all__ = ['main'] +__all__ = ['format_event_frames', 'main', 'serialize_event'] __getattr__, __dir__ = _lazy.accessors(globals(), _LAZY_MEMBERS) diff --git a/src/google/adk/cli/_sse.py b/src/google/adk/cli/_sse.py new file mode 100644 index 00000000000..1a6df87eaff --- /dev/null +++ b/src/google/adk/cli/_sse.py @@ -0,0 +1,90 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""SSE framing for the events streamed by ``/run_sse``. + +``/run_sse`` owns a wire format that ADK clients parse: internal metadata is +dropped, credential secrets are redacted, and an event carrying an +``artifactDelta`` is split into two frames. Applications that keep a ``Runner`` +outside the HTTP handler (or put their own transport in front of it) still need +to emit exactly that format, so the framing lives here instead of inside the +route's generator. ``run_agent_sse`` calls these functions directly, which keeps +the public helper and the built-in server from drifting apart. + +The module is private per the repo's visibility rules; ``serialize_event`` and +``format_event_frames`` are exported from ``google.adk.cli``. +""" + +from __future__ import annotations + +import pydantic_core + +from ..auth.auth_credential import _redact_credential_secrets +from ..events._internal_metadata import public_event +from ..events.event import Event + + +def serialize_event(event: Event) -> str: + """Serializes an event to the JSON text carried by a ``data:`` frame. + + This applies the same transformations as ``/run_sse``: internal metadata is + removed via ``public_event`` and credential secrets are redacted. + + Args: + event: The event to serialize. It is not modified. + + Returns: + The JSON text of the event, without the SSE ``data:`` prefix or trailing + blank line. + """ + return pydantic_core.to_json( + _redact_credential_secrets( + public_event(event).model_dump( + exclude_none=True, by_alias=True, mode="json" + ) + ) + ).decode("utf-8") + + +def format_event_frames( + event: Event, *, split_artifact_delta: bool = True +) -> list[str]: + """Formats the SSE frames that ``/run_sse`` writes for one runner event. + + Args: + event: The event to format. It is not modified. + split_artifact_delta: Whether an event that has both content and an + ``artifactDelta`` is split into a content frame and a content-less + ``artifactDelta`` frame. ADK Web renders artifacts from both content and + actions, so the split keeps an artifact from rendering twice. Pass + ``False`` to emit the event as a single frame, as ``/run_sse`` does for + the resume request that answers a long-running function call. + + Returns: + One or two ``data: \\n\\n`` strings, ready to be written to an SSE + stream in order. + """ + events_to_stream = [event] + if ( + split_artifact_delta + and event.actions.artifact_delta + and event.content + and event.content.parts + ): + content_event = event.model_copy(deep=True) + content_event.actions.artifact_delta = {} + artifact_event = event.model_copy(deep=True) + artifact_event.content = None + events_to_stream = [content_event, artifact_event] + return [f"data: {serialize_event(e)}\n\n" for e in events_to_stream] diff --git a/src/google/adk/cli/api_server.py b/src/google/adk/cli/api_server.py index 407b9892c56..60e43615b85 100644 --- a/src/google/adk/cli/api_server.py +++ b/src/google/adk/cli/api_server.py @@ -63,7 +63,6 @@ from opentelemetry.sdk.trace import TracerProvider from pydantic import Field from pydantic import ValidationError -import pydantic_core from starlette.types import Lifespan from starlette.types import Scope from typing_extensions import deprecated @@ -105,6 +104,8 @@ from ..utils.context_utils import Aclosing from ..utils.feature_decorator import experimental from ..version import __version__ +from ._sse import format_event_frames +from ._sse import serialize_event from .cli_eval import _LEGACY_EVAL_SESSION_ID_PREFIX from .cli_eval import EVAL_SESSION_ID_PREFIX from .utils import cleanup @@ -2273,36 +2274,15 @@ async def _produce_events() -> None: if event is None: break # ADK Web renders artifacts from `actions.artifactDelta` - # during part processing *and* during action processing - # 1) the original event with `artifactDelta` cleared (content) - # 2) a content-less "action-only" event carrying `artifactDelta` - events_to_stream = [event] - if ( - not req.function_call_event_id - and event.actions.artifact_delta - and event.content - and event.content.parts + # during part processing *and* during action processing, so the + # frame split keeps an artifact from rendering twice. The resume + # request that answers a long-running function call must not + # split. + for frame in format_event_frames( + event, split_artifact_delta=not req.function_call_event_id ): - content_event = event.model_copy(deep=True) - content_event.actions.artifact_delta = {} - artifact_event = event.model_copy(deep=True) - artifact_event.content = None - events_to_stream = [content_event, artifact_event] - - for event_to_stream in events_to_stream: - sse_event = pydantic_core.to_json( - _redact_credential_secrets( - public_event(event_to_stream).model_dump( - exclude_none=True, - by_alias=True, - mode="json", - ) - ) - ).decode("utf-8") - logger.debug( - "Generated event in agent run streaming: %s", sse_event - ) - yield f"data: {sse_event}\n\n" + logger.debug("Generated event in agent run stream: %s", frame) + yield frame next_step_event.set() except (GeneratorExit, asyncio.CancelledError) as e: is_closing = True @@ -2453,15 +2433,9 @@ async def forward_events() -> None: ) ) as agen: async for event in agen: - await websocket.send_text( - pydantic_core.to_json( - _redact_credential_secrets( - public_event(event).model_dump( - exclude_none=True, by_alias=True, mode="json" - ) - ) - ).decode("utf-8") - ) + # /run_live carries the same event JSON as a /run_sse data frame, + # minus the SSE framing, so it uses the one serializer too. + await websocket.send_text(serialize_event(event)) async def process_messages() -> None: try: diff --git a/tests/unittests/cli/test_fast_api.py b/tests/unittests/cli/test_fast_api.py index d754a35fa8c..f4e0d530e12 100755 --- a/tests/unittests/cli/test_fast_api.py +++ b/tests/unittests/cli/test_fast_api.py @@ -39,6 +39,8 @@ from google.adk.auth.auth_credential import _redact_credential_secrets from google.adk.cli import api_server as api_server_module from google.adk.cli import fast_api as fast_api_module +from google.adk.cli import format_event_frames +from google.adk.cli import serialize_event from google.adk.cli import service_registry as service_registry_module from google.adk.cli.api_server import RunAgentRequest from google.adk.cli.fast_api import get_fast_api_app @@ -6082,6 +6084,44 @@ async def setup_session(): assert data["customMetadata"] == {"keep": 1} +def test_run_live_frames_match_the_public_serializer( + test_app, mock_session_service, monkeypatch +): + """/run_live sends exactly `serialize_event(event)` as one unsplit frame.""" + yielded = [] + + async def run_live_with_artifact_delta(self, live_request_queue, **kwargs): + del live_request_queue, kwargs + event = Event( + author="dummy agent", + invocation_id="invocation_id", + content=types.Content(role="model", parts=[types.Part(text="hello")]), + actions=EventActions(artifact_delta={"artifact.txt": 0}), + ) + yielded.append(event) + yield event + + monkeypatch.setattr(Runner, "run_live", run_live_with_artifact_delta) + + async def setup_session(): + await mock_session_service.create_session( + app_name="test_app", user_id="user", session_id="session", state={} + ) + + asyncio.run(setup_session()) + + url = "/run_live?app_name=test_app&user_id=user&session_id=session&modalities=AUDIO" + with test_app.websocket_connect(url) as ws: + frame = ws.receive_text() + + # One frame, byte-identical to the public serializer: run_live must not + # split artifactDelta the way /run_sse does. + assert frame == serialize_event(yielded[0]) + payload = json.loads(frame) + assert payload["content"]["parts"][0]["text"] == "hello" + assert payload["actions"]["artifactDelta"] == {"artifact.txt": 0} + + def test_run_live_websocket_missing_app_name_raises_error( test_app, monkeypatch ): @@ -7788,5 +7828,54 @@ def test_url_prefix_propagated_to_docs_openapi_url( assert "/adk/openapi.json" in response.text +def test_agent_run_sse_frames_match_the_public_serializer( + test_app, create_test_session, monkeypatch +): + """/run_sse writes exactly the frames `format_event_frames` returns. + + The helper exists so a caller that streams a runner outside this route can + reproduce the wire format. Pinning the route's bytes to the helper keeps the + two from drifting: a framing change that is not made in the helper fails + here. + """ + info = create_test_session + streamed = Event( + author="agent", + invocation_id="invocation_id", + content=types.Content(role="model", parts=[types.Part(text="LLM reply")]), + actions=EventActions(artifact_delta={"artifact.txt": 0}), + ) + + async def run_async_with_artifact_delta( + self, + *, + user_id: str, + session_id: str, + invocation_id: Optional[str] = None, + new_message: Optional[types.Content] = None, + state_delta: Optional[dict[str, Any]] = None, + run_config: Optional[RunConfig] = None, + **kwargs, + ): + del user_id, session_id, invocation_id, new_message, state_delta + del run_config, kwargs + yield streamed + + monkeypatch.setattr(Runner, "run_async", run_async_with_artifact_delta) + + payload = { + "app_name": info["app_name"], + "user_id": info["user_id"], + "session_id": info["session_id"], + "new_message": {"role": "user", "parts": [{"text": "Hello agent"}]}, + "streaming": True, + } + + response = test_app.post("/run_sse", json=payload) + + assert response.status_code == 200 + assert response.text == "".join(format_event_frames(streamed)) + + if __name__ == "__main__": pytest.main(["-xvs", __file__]) diff --git a/tests/unittests/cli/test_sse.py b/tests/unittests/cli/test_sse.py new file mode 100644 index 00000000000..6d5ae82cb1f --- /dev/null +++ b/tests/unittests/cli/test_sse.py @@ -0,0 +1,168 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the public SSE framing helpers behind `/run_sse`.""" + +import json + +from google.adk.cli import format_event_frames +from google.adk.cli import serialize_event +from google.adk.events._internal_metadata import INTERNAL_METADATA_PREFIX +from google.adk.events.event import Event +from google.adk.events.event_actions import EventActions +from google.genai import types + + +def _text_event(text: str = "LLM reply") -> Event: + return Event( + author="agent", + invocation_id="invocation_id", + content=types.Content(role="model", parts=[types.Part(text=text)]), + ) + + +def _artifact_event() -> Event: + return Event( + author="agent", + invocation_id="invocation_id", + content=types.Content(role="model", parts=[types.Part(text="LLM reply")]), + actions=EventActions(artifact_delta={"artifact.txt": 0}), + ) + + +def test_serialize_event_uses_wire_aliases_and_drops_none_fields(): + """The payload matches the JSON `Event.model_dump` shape clients parse.""" + event = Event( + author="agent", + invocation_id="invocation_id", + content=types.Content( + role="model", + parts=[ + types.Part( + function_call=types.FunctionCall(name="tool", args={"a": 1}) + ) + ], + ), + ) + + payload = json.loads(serialize_event(event)) + + assert payload["content"]["parts"][0]["functionCall"]["name"] == "tool" + # snake_case aliases never reach the wire, and None fields are omitted. + assert "function_call" not in payload["content"]["parts"][0] + assert "errorCode" not in payload + + +def test_serialize_event_redacts_credential_secrets(): + """OAuth2 client secrets and PKCE verifiers never reach the frame.""" + auth_config = { + "authScheme": { + "type": "oauth2", + "flows": { + "authorizationCode": { + "authorizationUrl": "https://idp.example.com/auth", + "tokenUrl": "https://idp.example.com/token", + } + }, + }, + "rawAuthCredential": { + "authType": "oauth2", + "oauth2": { + "clientId": "public-client-id", + "clientSecret": "should-never-reach-the-client", + }, + }, + } + event = Event( + author="agent", + invocation_id="invocation_id", + actions=EventActions(requested_auth_configs={"fc-id": auth_config}), + ) + + frame = serialize_event(event) + + assert "should-never-reach-the-client" not in frame + payload = json.loads(frame) + raw_credential = payload["actions"]["requestedAuthConfigs"]["fc-id"][ + "rawAuthCredential" + ]["oauth2"] + assert raw_credential["clientId"] == "public-client-id" + assert "clientSecret" not in raw_credential + + +def test_serialize_event_drops_internal_metadata(): + """Internal-only custom metadata keys are not exposed to clients.""" + event = _text_event() + event.custom_metadata = { + f"{INTERNAL_METADATA_PREFIX}internal": "secret", + "public": "kept", + } + + payload = json.loads(serialize_event(event)) + + assert payload["customMetadata"] == {"public": "kept"} + + +def test_format_event_frames_splits_artifact_delta(): + """An event with content and an artifactDelta becomes two frames.""" + frames = format_event_frames(_artifact_event()) + + assert len(frames) == 2 + assert all(frame.startswith("data: ") for frame in frames) + assert all(frame.endswith("\n\n") for frame in frames) + + content_payload, artifact_payload = ( + json.loads(frame.removeprefix("data: ").removesuffix("\n\n")) + for frame in frames + ) + assert content_payload["content"]["parts"][0]["text"] == "LLM reply" + assert content_payload["actions"]["artifactDelta"] == {} + assert "content" not in artifact_payload + assert artifact_payload["actions"]["artifactDelta"] == {"artifact.txt": 0} + + +def test_format_event_frames_emits_one_frame_without_split(): + """Callers that must not split, e.g. a function-call resume, get one frame.""" + frames = format_event_frames(_artifact_event(), split_artifact_delta=False) + + assert len(frames) == 1 + payload = json.loads(frames[0].removeprefix("data: ").removesuffix("\n\n")) + assert payload["content"]["parts"][0]["text"] == "LLM reply" + assert payload["actions"]["artifactDelta"] == {"artifact.txt": 0} + + +def test_format_event_frames_does_not_split_contentless_events(): + """An artifactDelta without content stays in a single frame.""" + event = Event( + author="agent", + invocation_id="invocation_id", + actions=EventActions(artifact_delta={"artifact.txt": 0}), + ) + + frames = format_event_frames(event) + + assert len(frames) == 1 + payload = json.loads(frames[0].removeprefix("data: ").removesuffix("\n\n")) + assert "content" not in payload + assert payload["actions"]["artifactDelta"] == {"artifact.txt": 0} + + +def test_format_event_frames_does_not_modify_the_event(): + """Formatting leaves the caller's event untouched, including actions.""" + event = _artifact_event() + + format_event_frames(event) + + assert event.actions.artifact_delta == {"artifact.txt": 0} + assert event.content.parts[0].text == "LLM reply"