From cc28b4beb106a4db805dbf03bb1c70725867fe68 Mon Sep 17 00:00:00 2001 From: Aarav Mittal <137450929+a2105z@users.noreply.github.com> Date: Wed, 26 Aug 2026 15:18:24 -0500 Subject: [PATCH] fix(cli): stop unbounded debug trace retention on the api server get_fast_api_app(web=False) still installed ApiServerSpanExporter and InMemoryExporter even though the debug endpoints that read them are only mounted with web=True. Those buffers retain every span (including full llm_request / llm_response) for the life of the process and OOM long-lived Cloud Run / GKE instances. Register the exporters only on DevServer, and FIFO-cap the event dict, span deque, and session index so adk web cannot grow without bound either. Closes #6915 --- src/google/adk/cli/api_server.py | 49 +++++++++-- src/google/adk/cli/dev_server.py | 1 + tests/unittests/cli/test_fast_api.py | 122 +++++++++++++++++++++++++++ 3 files changed, 165 insertions(+), 7 deletions(-) diff --git a/src/google/adk/cli/api_server.py b/src/google/adk/cli/api_server.py index 8285af05a98..9b41998cf9e 100644 --- a/src/google/adk/cli/api_server.py +++ b/src/google/adk/cli/api_server.py @@ -21,6 +21,7 @@ import asyncio import base64 import binascii +from collections import deque from contextlib import asynccontextmanager import importlib import json @@ -455,10 +456,25 @@ async def __call__( await self._app(scope, receive, send) +# Debug-UI buffers. Production (`web=False`) does not install these exporters. +# Caps keep `adk web` from retaining every span for the life of the process. +# See https://github.com/google/adk-python/issues/6915 +_DEBUG_TRACE_EVENT_LIMIT = 512 +_DEBUG_SPAN_LIMIT = 4096 +_DEBUG_SESSION_TRACE_LIMIT = 512 + + +def _evict_oldest(store: dict[Any, Any], max_size: int) -> None: + """Drop insertion-oldest keys until store fits in max_size.""" + while len(store) > max_size: + store.pop(next(iter(store))) + + class ApiServerSpanExporter(export_lib.SpanExporter): - def __init__(self, trace_dict): + def __init__(self, trace_dict, *, max_events: int = _DEBUG_TRACE_EVENT_LIMIT): self.trace_dict = trace_dict + self._max_events = max_events def export( self, spans: typing.Sequence[ReadableSpan] @@ -474,6 +490,7 @@ def export( attributes["span_id"] = span.get_span_context().span_id if attributes.get("gcp.vertex.agent.event_id", None): self.trace_dict[attributes["gcp.vertex.agent.event_id"]] = attributes + _evict_oldest(self.trace_dict, self._max_events) return export_lib.SpanExportResult.SUCCESS def force_flush(self, timeout_millis: int = 30000) -> bool: @@ -482,10 +499,17 @@ def force_flush(self, timeout_millis: int = 30000) -> bool: class InMemoryExporter(export_lib.SpanExporter): - def __init__(self, trace_dict): + def __init__( + self, + trace_dict, + *, + max_spans: int = _DEBUG_SPAN_LIMIT, + max_sessions: int = _DEBUG_SESSION_TRACE_LIMIT, + ): super().__init__() - self._spans = [] + self._spans = deque(maxlen=max_spans) self.trace_dict = trace_dict + self._max_sessions = max_sessions @override def export( @@ -501,6 +525,7 @@ def export( trace_ids = self.trace_dict.setdefault(session_id, []) if trace_id not in trace_ids: trace_ids.append(trace_id) + _evict_oldest(self.trace_dict, self._max_sessions) self._spans.extend(spans) return export_lib.SpanExportResult.SUCCESS @@ -821,6 +846,11 @@ class ApiServer: _allow_special_agents: bool = False + # Only DevServer reads the debug trace data these exporters accumulate + # (`/dev/apps/.../debug/trace`). Registering them on ApiServer retains + # every span for the life of the process with no consumer (#6915). + _registers_debug_trace_exporters: bool = False + def __init__( self, *, @@ -1170,12 +1200,17 @@ async def internal_lifespan(app: FastAPI): memory_exporter = InMemoryExporter(session_trace_dict) self._memory_exporter = memory_exporter - _setup_telemetry( - otel_to_cloud=otel_to_cloud, - internal_exporters=[ + debug_trace_exporters = ( + [ export_lib.SimpleSpanProcessor(ApiServerSpanExporter(trace_dict)), export_lib.SimpleSpanProcessor(memory_exporter), - ], + ] + if self._registers_debug_trace_exporters + else [] + ) + _setup_telemetry( + otel_to_cloud=otel_to_cloud, + internal_exporters=debug_trace_exporters, ) if web_assets_dir: self._setup_runtime_config(web_assets_dir) diff --git a/src/google/adk/cli/dev_server.py b/src/google/adk/cli/dev_server.py index 67c2e09768c..7b8b6971c1a 100644 --- a/src/google/adk/cli/dev_server.py +++ b/src/google/adk/cli/dev_server.py @@ -425,6 +425,7 @@ class DevServer(ApiServer): """ _allow_special_agents: bool = True + _registers_debug_trace_exporters: bool = True def _get_agent_dir(self, app_name: str) -> str: """Resolves the agent directory and validates the app name to prevent path traversal.""" diff --git a/tests/unittests/cli/test_fast_api.py b/tests/unittests/cli/test_fast_api.py index a3fe28d35a3..f9c7f74bb81 100644 --- a/tests/unittests/cli/test_fast_api.py +++ b/tests/unittests/cli/test_fast_api.py @@ -4524,6 +4524,88 @@ def test_in_memory_exporter_clear_drops_spans_but_keeps_session_index(): assert session_trace_dict == {"session-a": [505]} +def test_api_server_span_exporter_evicts_oldest_events(): + """Event attributes are FIFO-capped so a long-lived web UI cannot OOM.""" + from google.adk.cli.api_server import ApiServerSpanExporter + + trace_dict = {} + exporter = ApiServerSpanExporter(trace_dict, max_events=2) + + exporter.export([ + _readable_span( + "call_llm", + trace_id=1, + attributes={"gcp.vertex.agent.event_id": "event-a"}, + ), + _readable_span( + "call_llm", + trace_id=2, + attributes={"gcp.vertex.agent.event_id": "event-b"}, + ), + _readable_span( + "call_llm", + trace_id=3, + attributes={"gcp.vertex.agent.event_id": "event-c"}, + ), + ]) + + assert list(trace_dict) == ["event-b", "event-c"] + + +def test_in_memory_exporter_evicts_oldest_spans(): + """Finished spans are deque-capped; lookups only see what still fits.""" + from google.adk.cli.api_server import InMemoryExporter + + session_trace_dict = {} + exporter = InMemoryExporter(session_trace_dict, max_spans=2) + + span_1 = _readable_span( + "call_llm", + trace_id=1, + attributes={"gcp.vertex.agent.session_id": "session-a"}, + ) + span_2 = _readable_span( + "call_llm", + trace_id=2, + attributes={"gcp.vertex.agent.session_id": "session-a"}, + ) + span_3 = _readable_span( + "call_llm", + trace_id=3, + attributes={"gcp.vertex.agent.session_id": "session-a"}, + ) + exporter.export([span_1, span_2, span_3]) + + assert exporter.get_finished_spans("session-a") == [span_2, span_3] + + +def test_in_memory_exporter_evicts_oldest_sessions(): + """The session -> trace-id index is capped the same way as the span buffer.""" + from google.adk.cli.api_server import InMemoryExporter + + session_trace_dict = {} + exporter = InMemoryExporter(session_trace_dict, max_sessions=2) + exporter.export([ + _readable_span( + "call_llm", + trace_id=1, + attributes={"gcp.vertex.agent.session_id": "session-a"}, + ), + _readable_span( + "call_llm", + trace_id=2, + attributes={"gcp.vertex.agent.session_id": "session-b"}, + ), + _readable_span( + "call_llm", + trace_id=3, + attributes={"gcp.vertex.agent.session_id": "session-c"}, + ), + ]) + + assert list(session_trace_dict) == ["session-b", "session-c"] + + ################################################# # Request-body plumbing tests ################################################# @@ -4622,6 +4704,46 @@ def test_dev_only_endpoints_absent_when_web_disabled( assert client.get("/list-apps").status_code == 200 +def test_debug_trace_exporters_only_registered_for_dev_server( + mock_session_service, + mock_artifact_service, + mock_memory_service, + mock_agent_loader, + mock_eval_sets_manager, + mock_eval_set_results_manager, +): + """web=False has no reader for debug traces, so it must not retain them.""" + from google.adk.cli import api_server as api_server_module + + with patch.object( + api_server_module, "_setup_telemetry", autospec=True + ) as mock_setup_telemetry: + _create_test_client( + mock_session_service, + mock_artifact_service, + mock_memory_service, + mock_agent_loader, + mock_eval_sets_manager, + mock_eval_set_results_manager, + web=False, + ) + assert mock_setup_telemetry.call_args.kwargs["internal_exporters"] == [] + + with patch.object( + api_server_module, "_setup_telemetry", autospec=True + ) as mock_setup_telemetry: + _create_test_client( + mock_session_service, + mock_artifact_service, + mock_memory_service, + mock_agent_loader, + mock_eval_sets_manager, + mock_eval_set_results_manager, + web=True, + ) + assert len(mock_setup_telemetry.call_args.kwargs["internal_exporters"]) == 2 + + def test_app_info_rejects_special_agent_only_in_api_server_mode( test_app, mock_session_service,