From fa935877d9c24d221cc4e79c696a14adc97ddb74 Mon Sep 17 00:00:00 2001 From: Jeremy Alvis Date: Tue, 7 Jul 2026 13:03:44 -0700 Subject: [PATCH 1/2] Force flush spans before a2a response completes for substrate tracing Signed-off-by: Jeremy Alvis --- go/adk/pkg/a2a/executor.go | 8 +++- go/adk/pkg/telemetry/tracing.go | 20 +++++++++ go/adk/pkg/telemetry/tracing_test.go | 42 +++++++++++++++++++ .../src/kagent/adk/_agent_executor.py | 8 ++++ .../src/kagent/core/tracing/__init__.py | 4 +- .../src/kagent/core/tracing/_utils.py | 19 +++++++++ .../tests/test_tracing_configure.py | 27 ++++++++++++ 7 files changed, 125 insertions(+), 3 deletions(-) create mode 100644 go/adk/pkg/telemetry/tracing_test.go diff --git a/go/adk/pkg/a2a/executor.go b/go/adk/pkg/a2a/executor.go index e92fc7df1c..06254a85c5 100644 --- a/go/adk/pkg/a2a/executor.go +++ b/go/adk/pkg/a2a/executor.go @@ -165,7 +165,13 @@ func (e *KAgentExecutor) Execute(ctx context.Context, reqCtx *a2asrv.RequestCont } ctx = telemetry.SetKAgentSpanAttributes(ctx, spanAttributes) ctx, invocationSpan := telemetry.StartInvocationSpan(ctx) - defer invocationSpan.End() + defer func() { + invocationSpan.End() + // Flush after the root span ends so it is included: on Agent + // Substrate the actor is checkpointed as soon as the response + // closes, freezing any unexported spans into the snapshot. + telemetry.ForceFlush(ctx) + }() telemetry.SetMessageMetadataAttributes(ctx, reqCtx.Message.Metadata) diff --git a/go/adk/pkg/telemetry/tracing.go b/go/adk/pkg/telemetry/tracing.go index 695ec0711e..259c9c6788 100644 --- a/go/adk/pkg/telemetry/tracing.go +++ b/go/adk/pkg/telemetry/tracing.go @@ -38,6 +38,26 @@ func StartInvocationSpan(ctx context.Context) (context.Context, trace.Span) { return otel.Tracer("gcp.vertex.agent").Start(ctx, "invocation") } +// ForceFlush exports any spans still buffered in the tracer provider's batch +// processor. Call it before an A2A response completes when the process may be +// suspended right afterwards: Agent Substrate checkpoints the actor as soon as +// the response body closes, so unexported spans stay frozen in the snapshot +// until the session's next resume (or forever, for a session's last message). +// Uses its own detached timeout because the request context is typically +// already canceled by the time deferred cleanup runs. +func ForceFlush(ctx context.Context) { + type flusher interface{ ForceFlush(context.Context) error } + fp, ok := otel.GetTracerProvider().(flusher) + if !ok { + return + } + flushCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 3*time.Second) + defer cancel() + if err := fp.ForceFlush(flushCtx); err != nil { + otel.Handle(err) + } +} + // Init initializes OpenTelemetry providers for Go ADK, sets global providers and // propagators, and returns a shutdown function. func Init(ctx context.Context, serviceName string, serviceNamespace string) (shutdown func(context.Context) error, enabled bool, err error) { diff --git a/go/adk/pkg/telemetry/tracing_test.go b/go/adk/pkg/telemetry/tracing_test.go new file mode 100644 index 0000000000..9e04ef0e3c --- /dev/null +++ b/go/adk/pkg/telemetry/tracing_test.go @@ -0,0 +1,42 @@ +package telemetry + +import ( + "context" + "testing" + + "go.opentelemetry.io/otel" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + "go.opentelemetry.io/otel/sdk/trace/tracetest" + "go.opentelemetry.io/otel/trace/noop" +) + +// ForceFlush must drain spans buffered in a batch processor (which would +// otherwise wait out its schedule delay) and be a no-op for providers +// without ForceFlush support (e.g. the default global no-op provider). +func TestForceFlush(t *testing.T) { + exporter := tracetest.NewInMemoryExporter() + tp := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(sdktrace.NewBatchSpanProcessor(exporter))) + prev := otel.GetTracerProvider() + otel.SetTracerProvider(tp) + t.Cleanup(func() { + otel.SetTracerProvider(prev) + _ = tp.Shutdown(context.Background()) + }) + + _, span := otel.Tracer("test").Start(context.Background(), "buffered") + span.End() + if got := len(exporter.GetSpans()); got != 0 { + t.Fatalf("span exported before flush: %d", got) + } + + // A canceled request context must not prevent the flush. + ctx, cancel := context.WithCancel(context.Background()) + cancel() + ForceFlush(ctx) + if got := len(exporter.GetSpans()); got != 1 { + t.Fatalf("expected 1 span after flush, got %d", got) + } + + otel.SetTracerProvider(noop.NewTracerProvider()) + ForceFlush(context.Background()) // must not panic +} diff --git a/python/packages/kagent-adk/src/kagent/adk/_agent_executor.py b/python/packages/kagent-adk/src/kagent/adk/_agent_executor.py index 0ed10a3177..b8ced48860 100644 --- a/python/packages/kagent-adk/src/kagent/adk/_agent_executor.py +++ b/python/packages/kagent-adk/src/kagent/adk/_agent_executor.py @@ -44,6 +44,7 @@ extract_rejection_reasons_from_message, get_kagent_metadata_key, ) +from kagent.core.tracing import force_flush as force_flush_tracing from kagent.core.tracing._span_processor import ( clear_kagent_span_attributes, set_kagent_span_attributes, @@ -198,6 +199,13 @@ async def execute( event_queue, str(e) or "A2A request execution was cancelled.", ) + finally: + # Flush buffered spans before the response completes: on Agent + # Substrate the actor is checkpointed as soon as the response + # body closes, freezing any unexported spans into the snapshot. + # Run in a thread so the blocking gRPC export (bounded by the + # flush timeout) doesn't stall the event loop. + await asyncio.to_thread(force_flush_tracing) async def _execute_impl( self, diff --git a/python/packages/kagent-core/src/kagent/core/tracing/__init__.py b/python/packages/kagent-core/src/kagent/core/tracing/__init__.py index 2d9bb77fc9..826775371d 100644 --- a/python/packages/kagent-core/src/kagent/core/tracing/__init__.py +++ b/python/packages/kagent-core/src/kagent/core/tracing/__init__.py @@ -1,3 +1,3 @@ -from ._utils import configure +from ._utils import configure, force_flush -__all__ = ["configure"] +__all__ = ["configure", "force_flush"] diff --git a/python/packages/kagent-core/src/kagent/core/tracing/_utils.py b/python/packages/kagent-core/src/kagent/core/tracing/_utils.py index f56939e986..38cb196cfe 100644 --- a/python/packages/kagent-core/src/kagent/core/tracing/_utils.py +++ b/python/packages/kagent-core/src/kagent/core/tracing/_utils.py @@ -106,6 +106,25 @@ def _instrument_google_generativeai(): pass +def force_flush(timeout_millis: int = 3000) -> None: + """Export any spans still buffered in the tracer provider's batch processor. + + Call before a response completes when the process may be suspended right + afterwards: Agent Substrate checkpoints the actor as soon as the A2A + response body closes, so unexported spans stay frozen in the snapshot + until the session's next resume (or forever, for a session's last + message). No-op when the provider has no force_flush (tracing disabled). + """ + provider = trace.get_tracer_provider() + flush = getattr(provider, "force_flush", None) + if flush is None: + return + try: + flush(timeout_millis) + except Exception: + logging.warning("Failed to flush pending spans", exc_info=True) + + def configure(name: str = "kagent", namespace: str = "kagent", fastapi_app: FastAPI | None = None): """Configure OpenTelemetry tracing and logging for this service. diff --git a/python/packages/kagent-core/tests/test_tracing_configure.py b/python/packages/kagent-core/tests/test_tracing_configure.py index ac99319c6e..d37ee8730c 100644 --- a/python/packages/kagent-core/tests/test_tracing_configure.py +++ b/python/packages/kagent-core/tests/test_tracing_configure.py @@ -120,3 +120,30 @@ def test_resolve_otlp_timeout_seconds_uses_milliseconds(monkeypatch, signal, env monkeypatch.setenv(key, value) assert _utils._resolve_otlp_timeout_seconds(signal) == expected + + +def test_force_flush_calls_provider_force_flush(monkeypatch): + calls = [] + provider = SimpleNamespace(force_flush=lambda timeout: calls.append(timeout)) + monkeypatch.setattr(_utils.trace, "get_tracer_provider", lambda: provider) + + _utils.force_flush() + + assert calls == [3000] + + +def test_force_flush_noop_without_provider_support(monkeypatch): + # The default (no-op) provider has no force_flush; must not raise. + monkeypatch.setattr(_utils.trace, "get_tracer_provider", lambda: SimpleNamespace()) + + _utils.force_flush() + + +def test_force_flush_swallows_exporter_errors(monkeypatch): + def boom(timeout): + raise RuntimeError("collector down") + + provider = SimpleNamespace(force_flush=boom) + monkeypatch.setattr(_utils.trace, "get_tracer_provider", lambda: provider) + + _utils.force_flush() From b2684db665f36bfc402d99bf7c2bab09d5eafe2b Mon Sep 17 00:00:00 2001 From: Jeremy Alvis Date: Tue, 7 Jul 2026 13:04:18 -0700 Subject: [PATCH 2/2] Fix grafana-mcp and querydoc helm image helpers Signed-off-by: Jeremy Alvis --- helm/tools/grafana-mcp/templates/_helpers.tpl | 4 ++-- helm/tools/querydoc/templates/_helpers.tpl | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/helm/tools/grafana-mcp/templates/_helpers.tpl b/helm/tools/grafana-mcp/templates/_helpers.tpl index bd074a815c..af4da7f48e 100644 --- a/helm/tools/grafana-mcp/templates/_helpers.tpl +++ b/helm/tools/grafana-mcp/templates/_helpers.tpl @@ -69,10 +69,10 @@ Create the grafana server URL {{- end }} {{/* -Join registry/repository for grafana-mcp image, skipping empty segments, then append tag +Join registry/repository/name/tag for grafana-mcp image, skipping empty segments, then append tag */}} {{- define "grafana-mcp.image" -}} {{- $img := .Values.image -}} -{{- $parts := compact (list $img.registry $img.repository) -}} +{{- $parts := compact (list $img.registry $img.repository $img.name) -}} {{- printf "%s:%s" (join "/" $parts) $img.tag -}} {{- end -}} \ No newline at end of file diff --git a/helm/tools/querydoc/templates/_helpers.tpl b/helm/tools/querydoc/templates/_helpers.tpl index fd9e15da92..b886aaf3b1 100644 --- a/helm/tools/querydoc/templates/_helpers.tpl +++ b/helm/tools/querydoc/templates/_helpers.tpl @@ -69,10 +69,10 @@ Create the querydoc server URL {{- end }} {{/* -Constructs the full image reference from registry/repository/tag for querydoc image +Constructs the full image reference from registry/repository/name/tag for querydoc image */}} {{- define "querydoc.image" -}} {{- $img := .Values.image -}} -{{- $parts := compact (list $img.registry $img.repository) -}} +{{- $parts := compact (list $img.registry $img.repository $img.name) -}} {{- printf "%s:%s" (join "/" $parts) ($img.tag | default .Chart.AppVersion) -}} {{- end }} \ No newline at end of file