From dc2d4898707d7eb582888264bb05ca047f4d0a9c Mon Sep 17 00:00:00 2001 From: "Shangxiang Fan (from Dev Box)" Date: Thu, 8 Oct 2026 16:36:15 +0800 Subject: [PATCH 1/3] feat(adk): expose gen_ai token usage Prometheus metric Record the GenAI-semconv gen_ai.client.token.usage histogram from the Go ADK runtime with the native Prometheus client, served at /metrics behind the default-OFF OTEL_METRICS_ENABLED gate. Labels follow the upstream Google ADK Python runtime so one dashboard works across both runtimes (gen_ai.token.type, gen_ai.request.model, gen_ai.response.model, gen_ai.provider.name, gen_ai.agent.name, gen_ai.operation.name, error.type), and token counts are clamped to the int32 range the metric field carries. Signed-off-by: Shangxiang Fan (from Dev Box) --- docs/genai-token-metrics.md | 104 +++++++++++++ go/adk/pkg/a2a/server/server.go | 8 + go/adk/pkg/a2a/server/server_test.go | 40 +++++ go/adk/pkg/runner/adapter.go | 5 + go/adk/pkg/runner/metrics.go | 40 +++++ go/adk/pkg/runner/metrics_test.go | 184 +++++++++++++++++++++++ go/adk/pkg/telemetry/metrics.go | 149 +++++++++++++++++++ go/adk/pkg/telemetry/metrics_test.go | 209 +++++++++++++++++++++++++++ 8 files changed, 739 insertions(+) create mode 100644 docs/genai-token-metrics.md create mode 100644 go/adk/pkg/runner/metrics.go create mode 100644 go/adk/pkg/runner/metrics_test.go create mode 100644 go/adk/pkg/telemetry/metrics.go create mode 100644 go/adk/pkg/telemetry/metrics_test.go diff --git a/docs/genai-token-metrics.md b/docs/genai-token-metrics.md new file mode 100644 index 0000000000..bd4a055922 --- /dev/null +++ b/docs/genai-token-metrics.md @@ -0,0 +1,104 @@ +# GenAI token-usage metrics + +kagent's **Go ADK** agent runtime records the OpenTelemetry GenAI-semconv metric +[`gen_ai.client.token.usage`](https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-metrics/#metric-gen_aiclienttokenusage) +using the native Prometheus client library and exposes it for scraping. It lets you graph and alert +on token spend per model / provider without parsing traces. + +## What is emitted + +A Prometheus histogram, served at **`/metrics`** on the agent's HTTP port: + +| Prometheus name | OTel semconv | Notes | +| --- | --- | --- | +| `gen_ai_client_token_usage` | `gen_ai.client.token.usage` | histogram, semconv-recommended buckets | + +Labels use semconv attribute names with dots mapped to underscores: + +| Label | Values | +| --- | --- | +| `gen_ai_token_type` | `input`, `output` (output = candidate + reasoning tokens) | +| `gen_ai_operation_name` | `generate_content` for Gemini; `chat` for chat-completion providers | +| `gen_ai_provider_name` | well-known value, e.g. `openai`, `anthropic`, `gcp.vertex_ai`, `aws.bedrock`, `azure.ai.openai` | +| `gen_ai_request_model` | configured model, e.g. `gpt-4o` | +| `gen_ai_response_model` | model the provider served (falls back to request model) | +| `gen_ai_agent_name` | ADK agent that made the model call (not the application name) | +| `error_type` | set on failed requests; empty otherwise | + +For each positive input/output count, one observation is recorded per LLM call. +Streaming partial chunks are not double-counted. Recording runs in the model callback, +so usage-only responses are counted even when they do not produce an A2A artifact. +Missing optional labels are exported as empty strings. + +## Configuration + +- **Runtime**: available on Declarative agents with `runtime: go`. (The Python runtime + records the same metric from upstream Google ADK.) +- **Gate**: recording and the `/metrics` endpoint are **default-OFF**, matching kagent's other + observability gates. Set `OTEL_METRICS_ENABLED=true` to turn them on: + + ```yaml + env: + - name: OTEL_METRICS_ENABLED + value: "true" + ``` + +- To have Prometheus scrape the endpoint, annotate the Go-runtime agent pods: + + ```yaml + metadata: + annotations: + prometheus.io/scrape: "true" + prometheus.io/port: "" + prometheus.io/path: "/metrics" + ``` + +### Scraping it + +Any Prometheus-compatible scraper that honors pod annotations will pick agents up. With an +OpenTelemetry Collector, add a `prometheus` receiver job with pod discovery: + +```yaml +receivers: + prometheus: + config: + scrape_configs: + - job_name: kagent-agents + kubernetes_sd_configs: [{ role: pod }] + relabel_configs: + - source_labels: [__meta_kubernetes_pod_annotation_prometheus_io_scrape] + regex: "true" + action: keep + - source_labels: [__meta_kubernetes_pod_annotation_prometheus_io_path] + target_label: __metrics_path__ + - source_labels: [__address__, __meta_kubernetes_pod_annotation_prometheus_io_port] + regex: ([^:]+)(?::\d+)?;(\d+) + replacement: $$1:$$2 + target_label: __address__ +``` + +## Verifying + +```bash +# exec into a Go-runtime agent pod and curl its metrics endpoint +kubectl exec -- wget -qO- localhost:/metrics | grep gen_ai_client_token_usage +``` + +Typical output after two chat requests (note `_count` equals the number of LLM calls, not stream +chunks, and the semconv labels including `gen_ai.agent.name`, response model, and an empty +`error.type`): + +```text +# HELP gen_ai_client_token_usage Measures the number of input and output tokens used by GenAI requests. +# TYPE gen_ai_client_token_usage histogram +gen_ai_client_token_usage_sum{error_type="",gen_ai_agent_name="my_agent",gen_ai_operation_name="generate_content",gen_ai_provider_name="gcp.vertex_ai",gen_ai_request_model="gemini-2.5-flash",gen_ai_response_model="gemini-2.5-flash",gen_ai_token_type="input"} 137 +gen_ai_client_token_usage_count{error_type="",gen_ai_agent_name="my_agent",gen_ai_operation_name="generate_content",gen_ai_provider_name="gcp.vertex_ai",gen_ai_request_model="gemini-2.5-flash",gen_ai_response_model="gemini-2.5-flash",gen_ai_token_type="input"} 2 +gen_ai_client_token_usage_sum{error_type="",gen_ai_agent_name="my_agent",gen_ai_operation_name="generate_content",gen_ai_provider_name="gcp.vertex_ai",gen_ai_request_model="gemini-2.5-flash",gen_ai_response_model="gemini-2.5-flash",gen_ai_token_type="output"} 91 +gen_ai_client_token_usage_count{error_type="",gen_ai_agent_name="my_agent",gen_ai_operation_name="generate_content",gen_ai_provider_name="gcp.vertex_ai",gen_ai_request_model="gemini-2.5-flash",gen_ai_response_model="gemini-2.5-flash",gen_ai_token_type="output"} 2 +``` + +## Follow-ups + +- Optional OTLP **push** (in addition to the scrape endpoint) via the OpenTelemetry + [Prometheus→OTLP bridge](https://pkg.go.dev/go.opentelemetry.io/contrib/bridges/prometheus), for + environments that push to an OTLP collector rather than scrape. diff --git a/go/adk/pkg/a2a/server/server.go b/go/adk/pkg/a2a/server/server.go index 5678033cf6..720ba3a1e0 100644 --- a/go/adk/pkg/a2a/server/server.go +++ b/go/adk/pkg/a2a/server/server.go @@ -25,6 +25,7 @@ import ( grpc_health_v1 "google.golang.org/grpc/health/grpc_health_v1" "google.golang.org/grpc/stats" + adktelemetry "github.com/kagent-dev/kagent/go/adk/pkg/telemetry" "github.com/kagent-dev/kagent/go/api/client" "github.com/kagent-dev/kagent/go/pkg/telemetry" "github.com/kagent-dev/kagent/go/pkg/tracing" @@ -96,6 +97,11 @@ func NewA2AServer(agentCard a2atype.AgentCard, executor a2asrv.AgentExecutor, lo mux := http.NewServeMux() registerHealthEndpoints(mux, healthPaths, healthHandler) mux.Handle(a2asrv.WellKnownAgentCardPath, a2asrv.NewStaticAgentCardHandler(&agentCard)) + // Serve Prometheus metrics for scraping when the metrics gate is on. This + // endpoint is excluded from request tracing and span flushing below. + if adktelemetry.MetricsEnabled() { + mux.Handle("/metrics", adktelemetry.MetricsHandler()) + } mux.Handle("/", jsonrpcHandler) grpcServer := grpc.NewServer( @@ -118,6 +124,8 @@ func NewA2AServer(agentCard a2atype.AgentCard, executor a2asrv.AgentExecutor, lo return false case r.URL.Path == a2asrv.WellKnownAgentCardPath, slices.Contains(healthPaths, r.URL.Path): return false + case r.URL.Path == "/metrics": + return false default: return true } diff --git a/go/adk/pkg/a2a/server/server_test.go b/go/adk/pkg/a2a/server/server_test.go index f185674df3..27d012192e 100644 --- a/go/adk/pkg/a2a/server/server_test.go +++ b/go/adk/pkg/a2a/server/server_test.go @@ -18,6 +18,8 @@ import ( a2apb "github.com/a2aproject/a2a-go/v2/a2apb/v1" "github.com/a2aproject/a2a-go/v2/a2apb/v1/pbconv" "github.com/a2aproject/a2a-go/v2/a2asrv" + adktelemetry "github.com/kagent-dev/kagent/go/adk/pkg/telemetry" + "github.com/stretchr/testify/require" "go.opentelemetry.io/otel" sdktrace "go.opentelemetry.io/otel/sdk/trace" "go.opentelemetry.io/otel/sdk/trace/tracetest" @@ -29,6 +31,44 @@ import ( "github.com/kagent-dev/kagent/go/pkg/telemetry" ) +func TestMetricsEndpoint(t *testing.T) { + for _, test := range []struct { + name string + input string + want bool + }{ + {name: "disabled"}, + {name: "enabled", input: "true", want: true}, + } { + t.Run(test.name, func(t *testing.T) { + t.Setenv("OTEL_METRICS_ENABLED", test.input) + t.Setenv("KAGENT_PRE_RESPONSE_TRACE_FLUSH", "true") + exporter := tracetest.NewInMemoryExporter() + provider := sdktrace.NewTracerProvider(sdktrace.WithSyncer(exporter)) + previous := otel.GetTracerProvider() + otel.SetTracerProvider(provider) + t.Cleanup(func() { + otel.SetTracerProvider(previous) + require.NoError(t, provider.Shutdown(context.Background())) + }) + testServer, _ := startTestServer(t) + adktelemetry.RecordTokenUsage(adktelemetry.TokenUsage{RequestModel: "test-model", Provider: "openai", InputTokens: 10}) + response, err := testServer.Client().Get(testServer.URL + "/metrics") + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, response.Body.Close()) }) + body, err := io.ReadAll(response.Body) + require.NoError(t, err) + if test.want { + require.Equal(t, http.StatusOK, response.StatusCode) + require.Contains(t, string(body), "gen_ai_client_token_usage_count") + } else { + require.NotContains(t, string(body), "gen_ai_client_token_usage") + } + require.Empty(t, exporter.GetSpans()) + }) + } +} + // substrateExecutor stands in for an ADK turn: it opens the invoke_agent span // ADK emits beneath the request span. It does not flush; exporting everything, // including the otelhttp server span that stays open until the mux handler diff --git a/go/adk/pkg/runner/adapter.go b/go/adk/pkg/runner/adapter.go index 6dc8d335da..929c42f077 100644 --- a/go/adk/pkg/runner/adapter.go +++ b/go/adk/pkg/runner/adapter.go @@ -82,6 +82,11 @@ func CreateRunnerConfig( } var adkPlugins []*adkplugin.Plugin + metricsPlugin, err := newTokenUsagePlugin(config.Agent.Model) + if err != nil { + return runner.Config{}, fmt.Errorf("failed to create token usage plugin: %w", err) + } + adkPlugins = append(adkPlugins, metricsPlugin) if stsPlugin != nil { p, err := stsPlugin.ADKPlugin() if err != nil { diff --git a/go/adk/pkg/runner/metrics.go b/go/adk/pkg/runner/metrics.go new file mode 100644 index 0000000000..9c92ba06ef --- /dev/null +++ b/go/adk/pkg/runner/metrics.go @@ -0,0 +1,40 @@ +package runner + +import ( + "github.com/kagent-dev/kagent/go/adk/pkg/config" + "github.com/kagent-dev/kagent/go/adk/pkg/telemetry" + configv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/config/v1" + adkagent "google.golang.org/adk/v2/agent" + "google.golang.org/adk/v2/model" + "google.golang.org/adk/v2/plugin" +) + +func newTokenUsagePlugin(modelConfig *configv1.ADKModelConfig) (*plugin.Plugin, error) { + modelName := modelConfig.GetModel() + modelType := config.ModelType(modelConfig) + providerName := telemetry.SemconvProviderName(modelType) + operationName := "chat" + if modelType == "gemini" || modelType == "gemini_vertex_ai" { + operationName = "generate_content" + } + return plugin.New(plugin.Config{ + Name: "kagent_token_usage", + AfterModelCallback: func(ctx adkagent.Context, response *model.LLMResponse, _ error) (*model.LLMResponse, error) { + if response == nil || response.Partial || response.UsageMetadata == nil { + return nil, nil + } + usage := response.UsageMetadata + telemetry.RecordTokenUsage(telemetry.TokenUsage{ + OperationName: operationName, + RequestModel: modelName, + ResponseModel: response.ModelVersion, + Provider: providerName, + AgentName: ctx.AgentName(), + ErrorType: response.ErrorCode, + InputTokens: int64(usage.PromptTokenCount), + OutputTokens: int64(usage.CandidatesTokenCount) + int64(usage.ThoughtsTokenCount), + }) + return nil, nil + }, + }) +} diff --git a/go/adk/pkg/runner/metrics_test.go b/go/adk/pkg/runner/metrics_test.go new file mode 100644 index 0000000000..cd3f4d0033 --- /dev/null +++ b/go/adk/pkg/runner/metrics_test.go @@ -0,0 +1,184 @@ +package runner + +import ( + "context" + "iter" + "strings" + "testing" + + "github.com/google/uuid" + configv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/config/v1" + "github.com/prometheus/client_golang/prometheus" + "github.com/stretchr/testify/require" + adkagent "google.golang.org/adk/v2/agent" + "google.golang.org/adk/v2/agent/llmagent" + "google.golang.org/adk/v2/model" + adkrunner "google.golang.org/adk/v2/runner" + "google.golang.org/adk/v2/session" + "google.golang.org/genai" +) + +type metricsContext struct { + adkagent.Context + agentName string +} + +var _ adkagent.Context = (*metricsContext)(nil) + +func (metrics *metricsContext) AgentName() string { return metrics.agentName } + +type tokenHistogram struct { + labels map[string]string + count uint64 + sum float64 +} + +func tokenHistograms(t *testing.T, agentName string) map[string]tokenHistogram { + t.Helper() + families, err := prometheus.DefaultGatherer.Gather() + require.NoError(t, err) + histograms := make(map[string]tokenHistogram) + for _, family := range families { + if family.GetName() != "gen_ai_client_token_usage" { + continue + } + for _, metric := range family.Metric { + labels := make(map[string]string) + for _, label := range metric.Label { + labels[label.GetName()] = label.GetValue() + } + if labels["gen_ai_agent_name"] == agentName { + tokenType := labels["gen_ai_token_type"] + _, duplicate := histograms[tokenType] + require.False(t, duplicate, "unexpected duplicate token type: %s", tokenType) + histograms[tokenType] = tokenHistogram{labels: labels, count: metric.Histogram.GetSampleCount(), sum: metric.Histogram.GetSampleSum()} + } + } + } + return histograms +} + +func TestTokenUsagePlugin(t *testing.T) { + const modelName = "configured-model" + for _, test := range []struct { + name string + input *configv1.ADKModelConfig + provider string + operation string + }{ + {name: "openai", input: &configv1.ADKModelConfig{Model: modelName, Provider: &configv1.ADKModelConfig_Openai{Openai: &configv1.ADKOpenAI{}}}, provider: "openai", operation: "chat"}, + {name: "gemini", input: &configv1.ADKModelConfig{Model: modelName, Provider: &configv1.ADKModelConfig_Gemini{Gemini: &configv1.ADKGemini{}}}, provider: "gcp.gemini", operation: "generate_content"}, + {name: "vertex", input: &configv1.ADKModelConfig{Model: modelName, Provider: &configv1.ADKModelConfig_GeminiVertexAi{GeminiVertexAi: &configv1.ADKGeminiVertexAI{}}}, provider: "gcp.vertex_ai", operation: "generate_content"}, + {name: "vertex anthropic", input: &configv1.ADKModelConfig{Model: modelName, Provider: &configv1.ADKModelConfig_GeminiAnthropic{GeminiAnthropic: &configv1.ADKGeminiAnthropic{}}}, provider: "gcp.vertex_ai", operation: "chat"}, + {name: "bedrock", input: &configv1.ADKModelConfig{Model: modelName, Provider: &configv1.ADKModelConfig_Bedrock{Bedrock: &configv1.ADKBedrock{}}}, provider: "aws.bedrock", operation: "chat"}, + {name: "sap", input: &configv1.ADKModelConfig{Model: modelName, Provider: &configv1.ADKModelConfig_SapAiCore{SapAiCore: &configv1.ADKSAPAICore{}}}, provider: "sap_ai_core", operation: "chat"}, + {name: "foundry", input: &configv1.ADKModelConfig{Model: modelName, Provider: &configv1.ADKModelConfig_Foundry{Foundry: &configv1.ADKFoundry{}}}, provider: "foundry", operation: "chat"}, + } { + t.Run(test.name, func(t *testing.T) { + t.Setenv("OTEL_METRICS_ENABLED", "true") + metricsPlugin, err := newTokenUsagePlugin(test.input) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, metricsPlugin.Close()) }) + agentName := "metrics_" + strings.ReplaceAll(uuid.NewString(), "-", "_") + ctx := &metricsContext{agentName: agentName} + callback := metricsPlugin.AfterModelCallback() + for _, response := range []*model.LLMResponse{ + nil, + {}, + {Partial: true, UsageMetadata: &genai.GenerateContentResponseUsageMetadata{PromptTokenCount: 100}}, + } { + replacement, err := callback(ctx, response, nil) + require.NoError(t, err) + require.Nil(t, replacement) + } + require.Empty(t, tokenHistograms(t, agentName)) + response := &model.LLMResponse{ + ModelVersion: "served-model", + ErrorCode: "overloaded_error", + UsageMetadata: &genai.GenerateContentResponseUsageMetadata{ + PromptTokenCount: 100, CandidatesTokenCount: 40, ThoughtsTokenCount: 2, + }, + } + for range 2 { + replacement, err := callback(ctx, response, nil) + require.NoError(t, err) + require.Nil(t, replacement) + } + histograms := tokenHistograms(t, agentName) + require.Len(t, histograms, 2) + for tokenType, want := range map[string]float64{"input": 200, "output": 84} { + histogram := histograms[tokenType] + require.Equal(t, uint64(2), histogram.count) + require.Equal(t, want, histogram.sum) + require.Equal(t, map[string]string{ + "gen_ai_token_type": tokenType, "gen_ai_operation_name": test.operation, + "gen_ai_provider_name": test.provider, "gen_ai_request_model": modelName, + "gen_ai_response_model": "served-model", "gen_ai_agent_name": agentName, + "error_type": "overloaded_error", + }, histogram.labels) + } + }) + } +} + +type metricsModel struct { + responses []*model.LLMResponse +} + +var _ model.LLM = (*metricsModel)(nil) + +func (metrics *metricsModel) Name() string { return "configured-model" } + +func (metrics *metricsModel) GenerateContent(context.Context, *model.LLMRequest, bool) iter.Seq2[*model.LLMResponse, error] { + return func(yield func(*model.LLMResponse, error) bool) { + for _, response := range metrics.responses { + if !yield(response, nil) { + return + } + } + } +} + +func TestRunnerRecordsTokenUsage(t *testing.T) { + t.Setenv("OTEL_METRICS_ENABLED", "true") + t.Setenv("OPENAI_API_KEY", "test-key") + t.Setenv("KAGENT_PROPAGATE_TOKEN", "false") + t.Setenv("STS_WELL_KNOWN_URI", "") + t.Setenv("KAGENT_SKILLS_FOLDER", "") + usage := &genai.GenerateContentResponseUsageMetadata{PromptTokenCount: 100, CandidatesTokenCount: 40, ThoughtsTokenCount: 2} + for _, test := range []struct { + name string + input []*model.LLMResponse + }{ + {name: "usage without content", input: []*model.LLMResponse{{UsageMetadata: usage}}}, + {name: "streamed response", input: []*model.LLMResponse{ + {Partial: true, Content: genai.NewContentFromText("partial", genai.RoleModel), UsageMetadata: usage}, + {Content: genai.NewContentFromText("final", genai.RoleModel), UsageMetadata: usage}, + }}, + } { + t.Run(test.name, func(t *testing.T) { + agentConfig := &configv1.ADKConfig{Agent: &configv1.ADKAgentConfig{Model: &configv1.ADKModelConfig{Model: "configured-model", Provider: &configv1.ADKModelConfig_Openai{Openai: &configv1.ADKOpenAI{}}}}} + runnerConfig, err := CreateRunnerConfig(t.Context(), agentConfig, session.InMemoryService(), "metrics-app", nil, nil) + require.NoError(t, err) + require.Len(t, runnerConfig.PluginConfig.Plugins, 1) + t.Cleanup(func() { require.NoError(t, runnerConfig.PluginConfig.Plugins[0].Close()) }) + agentName := "metrics_" + strings.ReplaceAll(uuid.NewString(), "-", "_") + runnerConfig.Agent, err = llmagent.New(llmagent.Config{ + Name: agentName, Model: &metricsModel{responses: test.input}, + }) + require.NoError(t, err) + runnerConfig.AutoCreateSession = true + runtime, err := adkrunner.New(runnerConfig) + require.NoError(t, err) + for _, err := range runtime.Run(t.Context(), "user", "session", genai.NewContentFromText("hello", genai.RoleUser), adkagent.RunConfig{StreamingMode: adkagent.StreamingModeSSE}) { + require.NoError(t, err) + } + histograms := tokenHistograms(t, agentName) + require.Len(t, histograms, 2) + require.Equal(t, uint64(1), histograms["input"].count) + require.Equal(t, float64(100), histograms["input"].sum) + require.Equal(t, uint64(1), histograms["output"].count) + require.Equal(t, float64(42), histograms["output"].sum) + }) + } +} diff --git a/go/adk/pkg/telemetry/metrics.go b/go/adk/pkg/telemetry/metrics.go new file mode 100644 index 0000000000..3ca5fcb9a7 --- /dev/null +++ b/go/adk/pkg/telemetry/metrics.go @@ -0,0 +1,149 @@ +package telemetry + +import ( + "net/http" + "os" + "strings" + + "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/client_golang/prometheus/promauto" + "github.com/prometheus/client_golang/prometheus/promhttp" +) + +// GenAI token-usage instrumentation using the native Prometheus client library. +// +// Metric and attribute names follow the OpenTelemetry GenAI semantic +// conventions (semconv 1.40.0), mapped to Prometheus naming (dots -> underscores): +// - metric gen_ai.client.token.usage -> gen_ai_client_token_usage +// - attrs gen_ai.token.type, gen_ai.request.model, gen_ai.response.model, +// gen_ai.provider.name, gen_ai.agent.name, gen_ai.operation.name, error.type +// +// https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-metrics/ +const ( + metricGenAIClientTokenUsage = "gen_ai_client_token_usage" + + labelGenAITokenType = "gen_ai_token_type" + labelGenAIOperationName = "gen_ai_operation_name" + labelGenAIProviderName = "gen_ai_provider_name" + labelGenAIRequestModel = "gen_ai_request_model" + labelGenAIResponseModel = "gen_ai_response_model" + labelGenAIAgentName = "gen_ai_agent_name" + labelErrorType = "error_type" + + tokenTypeInput = "input" + tokenTypeOutput = "output" + + // operationChat is the default gen_ai.operation.name for chat-completion calls. + operationChat = "chat" + + // metricsEnabledEnvVar gates metrics recording and the /metrics endpoint. + metricsEnabledEnvVar = "OTEL_METRICS_ENABLED" +) + +// tokenUsageBuckets is the explicit bucket layout recommended by the GenAI +// metrics semantic conventions for gen_ai.client.token.usage. +// https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-metrics/#metric-gen_aiclienttokenusage +var tokenUsageBuckets = []float64{1, 4, 16, 64, 256, 1024, 4096, 16384, 65536, 262144, 1048576, 4194304, 16777216, 67108864} + +// tokenUsage is the gen_ai.client.token.usage histogram, registered on the +// default Prometheus registry so it is served by MetricsHandler alongside the +// standard Go/process collectors. It is always registered; MetricsEnabled +// gates whether observations are recorded, keeping the scrape endpoint stable. +var tokenUsage = promauto.NewHistogramVec( + prometheus.HistogramOpts{ + Name: metricGenAIClientTokenUsage, + Help: "Measures the number of input and output tokens used by GenAI requests.", + Buckets: tokenUsageBuckets, + }, + []string{ + labelGenAITokenType, + labelGenAIOperationName, + labelGenAIProviderName, + labelGenAIRequestModel, + labelGenAIResponseModel, + labelGenAIAgentName, + labelErrorType, + }, +) + +// MetricsEnabled reports whether the GenAI token metrics pipeline is on, i.e. +// the OTEL_METRICS_ENABLED gate used consistently across kagent's runtimes. +func MetricsEnabled() bool { + return strings.EqualFold(strings.TrimSpace(os.Getenv(metricsEnabledEnvVar)), "true") +} + +// MetricsHandler returns an http.Handler that serves the agent's Prometheus +// metrics, intended to be mounted at /metrics for scraping. +func MetricsHandler() http.Handler { + return promhttp.Handler() +} + +// SemconvProviderName maps a kagent model type (adk.Model.GetType()) to the +// OpenTelemetry GenAI well-known gen_ai.provider.name value. Types without a +// well-known mapping (e.g. ollama, sap_ai_core, custom) pass through unchanged. +// https://opentelemetry.io/docs/specs/semconv/registry/attributes/gen-ai/#gen-ai-provider-name +func SemconvProviderName(modelType string) string { + switch modelType { + case "openai": + return "openai" + case "azure_openai": + return "azure.ai.openai" + case "anthropic": + return "anthropic" + case "gemini": + return "gcp.gemini" + case "gemini_vertex_ai", "gemini_anthropic": + return "gcp.vertex_ai" + case "bedrock": + return "aws.bedrock" + default: + return modelType + } +} + +// TokenUsage carries the per-request labels and counts for one recording on the +// gen_ai.client.token.usage histogram. +type TokenUsage struct { + OperationName string + // RequestModel is gen_ai.request.model (the configured model). + RequestModel string + // ResponseModel is gen_ai.response.model (the model the provider actually + // served). Falls back to RequestModel when empty. + ResponseModel string + // Provider is gen_ai.provider.name (a semconv well-known value). + Provider string + // AgentName is gen_ai.agent.name (the agent that produced the tokens). + AgentName string + // ErrorType is error.type; empty for successful requests. + ErrorType string + // InputTokens / OutputTokens are the token counts (output = candidate + + // reasoning tokens). Non-positive counts are skipped. + InputTokens int64 + OutputTokens int64 +} + +// RecordTokenUsage records input/output token counts on the +// gen_ai.client.token.usage histogram. If the metric pipeline is disabled +// (OTEL_METRICS_ENABLED unset or not "true"), it is a no-op. Zero/negative +// counts are skipped. +func RecordTokenUsage(usage TokenUsage) { + if !MetricsEnabled() { + return + } + operationName := usage.OperationName + if operationName == "" { + operationName = operationChat + } + responseModel := usage.ResponseModel + if responseModel == "" { + responseModel = usage.RequestModel + } + if usage.InputTokens > 0 { + tokenUsage.WithLabelValues(tokenTypeInput, operationName, usage.Provider, usage.RequestModel, responseModel, usage.AgentName, usage.ErrorType). + Observe(float64(usage.InputTokens)) + } + if usage.OutputTokens > 0 { + tokenUsage.WithLabelValues(tokenTypeOutput, operationName, usage.Provider, usage.RequestModel, responseModel, usage.AgentName, usage.ErrorType). + Observe(float64(usage.OutputTokens)) + } +} diff --git a/go/adk/pkg/telemetry/metrics_test.go b/go/adk/pkg/telemetry/metrics_test.go new file mode 100644 index 0000000000..61b76257e1 --- /dev/null +++ b/go/adk/pkg/telemetry/metrics_test.go @@ -0,0 +1,209 @@ +package telemetry + +import ( + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + + "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/client_golang/prometheus/testutil" + dto "github.com/prometheus/client_model/go" + "github.com/stretchr/testify/require" +) + +// TestRecordTokenUsage_RecordsHistogram verifies input/output token counts are +// recorded as two separate series on the gen_ai_client_token_usage histogram. +func TestRecordTokenUsage_RecordsHistogram(t *testing.T) { + t.Setenv(metricsEnabledEnvVar, "true") + tokenUsage.Reset() + + RecordTokenUsage(TokenUsage{ + RequestModel: "gpt-4o", Provider: "openai", InputTokens: 100, OutputTokens: 42, + }) + + // One series per token type (input, output). + if got := testutil.CollectAndCount(tokenUsage); got != 2 { + t.Fatalf("expected 2 histogram series (input+output), got %d", got) + } +} + +func TestRecordTokenUsage_SkipsZero(t *testing.T) { + t.Setenv(metricsEnabledEnvVar, "true") + tokenUsage.Reset() + + RecordTokenUsage(TokenUsage{RequestModel: "gpt-4o", Provider: "openai"}) + + if got := testutil.CollectAndCount(tokenUsage); got != 0 { + t.Fatalf("expected no series for zero token counts, got %d", got) + } +} + +func TestRecordTokenUsage_Disabled(t *testing.T) { + t.Setenv(metricsEnabledEnvVar, "false") + tokenUsage.Reset() + + RecordTokenUsage(TokenUsage{RequestModel: "gpt-4o", Provider: "openai", InputTokens: 100}) + + if got := testutil.CollectAndCount(tokenUsage); got != 0 { + t.Fatalf("expected no series when metrics are disabled, got %d", got) + } +} + +func TestRecordTokenUsage_ResponseModelFallback(t *testing.T) { + t.Setenv(metricsEnabledEnvVar, "true") + tokenUsage.Reset() + RecordTokenUsage(TokenUsage{RequestModel: "gemini-2.5-flash", Provider: "gcp.vertex_ai", InputTokens: 3}) + RecordTokenUsage(TokenUsage{RequestModel: "gemini-2.5-flash", ResponseModel: "gemini-2.5-flash-002", Provider: "gcp.vertex_ai", InputTokens: 3}) + + body := serveMetrics(t) + if !strings.Contains(body, "gen_ai_response_model=\"gemini-2.5-flash\"") { + t.Errorf("expected response model to fall back to request model") + } + if !strings.Contains(body, "gen_ai_response_model=\"gemini-2.5-flash-002\"") { + t.Errorf("expected explicit response model to be used") + } +} + +func TestMetricsHandler_ServesLabels(t *testing.T) { + t.Setenv(metricsEnabledEnvVar, "true") + tokenUsage.Reset() + RecordTokenUsage(TokenUsage{ + RequestModel: "claude-3-5-sonnet", ResponseModel: "claude-3-5-sonnet-20241022", + Provider: "anthropic", AgentName: "my-agent", ErrorType: "overloaded_error", + InputTokens: 10, OutputTokens: 5, + }) + + body := serveMetrics(t) + for _, want := range []string{ + "gen_ai_client_token_usage_count", + "gen_ai_token_type=\"input\"", + "gen_ai_token_type=\"output\"", + "gen_ai_operation_name=\"chat\"", + "gen_ai_provider_name=\"anthropic\"", + "gen_ai_request_model=\"claude-3-5-sonnet\"", + "gen_ai_response_model=\"claude-3-5-sonnet-20241022\"", + "gen_ai_agent_name=\"my-agent\"", + "error_type=\"overloaded_error\"", + } { + if !strings.Contains(body, want) { + t.Errorf("metrics output missing %q", want) + } + } +} + +func TestSemconvProviderName(t *testing.T) { + cases := map[string]string{ + "openai": "openai", + "azure_openai": "azure.ai.openai", + "anthropic": "anthropic", + "gemini": "gcp.gemini", + "gemini_vertex_ai": "gcp.vertex_ai", + "gemini_anthropic": "gcp.vertex_ai", + "bedrock": "aws.bedrock", + "ollama": "ollama", + "sap_ai_core": "sap_ai_core", + "some-custom": "some-custom", + } + for input, want := range cases { + t.Run(input, func(t *testing.T) { + require.Equal(t, want, SemconvProviderName(input)) + }) + } +} + +func TestRecordTokenUsageValues(t *testing.T) { + for _, test := range []struct { + name string + input TokenUsage + want map[string]float64 + }{ + {name: "input and output", input: TokenUsage{InputTokens: 100, OutputTokens: 42}, want: map[string]float64{"input": 100, "output": 42}}, + {name: "input only", input: TokenUsage{InputTokens: 100}, want: map[string]float64{"input": 100}}, + {name: "output only", input: TokenUsage{OutputTokens: 42}, want: map[string]float64{"output": 42}}, + {name: "zero", input: TokenUsage{}}, + {name: "negative", input: TokenUsage{InputTokens: -1, OutputTokens: -2}}, + {name: "negative input", input: TokenUsage{InputTokens: -1, OutputTokens: 42}, want: map[string]float64{"output": 42}}, + } { + t.Run(test.name, func(t *testing.T) { + t.Setenv(metricsEnabledEnvVar, "true") + tokenUsage.Reset() + t.Cleanup(tokenUsage.Reset) + RecordTokenUsage(test.input) + RecordTokenUsage(test.input) + require.Equal(t, len(test.want), testutil.CollectAndCount(tokenUsage)) + for tokenType, want := range test.want { + metric := &dto.Metric{} + histogram := tokenUsage.WithLabelValues(tokenType, operationChat, "", "", "", "", "") + require.NoError(t, histogram.(prometheus.Metric).Write(metric)) + require.Equal(t, uint64(2), metric.Histogram.GetSampleCount()) + require.Equal(t, 2*want, metric.Histogram.GetSampleSum()) + require.Len(t, metric.Histogram.Bucket, len(tokenUsageBuckets)) + for index, bucket := range metric.Histogram.Bucket { + require.Equal(t, tokenUsageBuckets[index], bucket.GetUpperBound()) + var count uint64 + if want <= bucket.GetUpperBound() { + count = 2 + } + require.Equal(t, count, bucket.GetCumulativeCount()) + } + } + }) + } +} + +func TestRecordTokenUsageConcurrent(t *testing.T) { + t.Setenv(metricsEnabledEnvVar, "true") + tokenUsage.Reset() + t.Cleanup(tokenUsage.Reset) + var workers sync.WaitGroup + for range 16 { + workers.Go(func() { + for range 64 { + RecordTokenUsage(TokenUsage{InputTokens: 3, OutputTokens: 7}) + } + }) + } + for range 16 { + require.NotEmpty(t, serveMetrics(t)) + } + workers.Wait() + for tokenType, want := range map[string]float64{"input": 3, "output": 7} { + metric := &dto.Metric{} + histogram := tokenUsage.WithLabelValues(tokenType, operationChat, "", "", "", "", "") + require.NoError(t, histogram.(prometheus.Metric).Write(metric)) + require.Equal(t, uint64(16*64), metric.Histogram.GetSampleCount()) + require.Equal(t, 16*64*want, metric.Histogram.GetSampleSum()) + } +} + +func TestMetricsEnabled(t *testing.T) { + for _, test := range []struct { + name string + input string + want bool + }{ + {name: "unset"}, + {name: "enabled", input: "true", want: true}, + {name: "case and whitespace", input: " TRUE ", want: true}, + {name: "disabled", input: "false"}, + {name: "invalid", input: "1"}, + } { + t.Run(test.name, func(t *testing.T) { + t.Setenv(metricsEnabledEnvVar, test.input) + require.Equal(t, test.want, MetricsEnabled()) + }) + } +} + +func serveMetrics(t *testing.T) string { + t.Helper() + req := httptest.NewRequest(http.MethodGet, "/metrics", nil) + rec := httptest.NewRecorder() + MetricsHandler().ServeHTTP(rec, req) + if rec.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", rec.Code) + } + return rec.Body.String() +} From ee861cae4481b3cdb0b8d9766ec09ea6ff2e5b9e Mon Sep 17 00:00:00 2001 From: "Shangxiang Fan (from Dev Box)" Date: Thu, 8 Oct 2026 16:36:57 +0800 Subject: [PATCH 2/3] feat(adk): emit cached token series on gen_ai token usage metric Adds a cached-token series to the gen_ai token usage histogram so prompt-cache hits are observable per provider. The token-type series are additive: the input observation excludes the cached portion, so input + cached always equals the prompt count the model reported and summing the series never double counts cache hits. Counts are clamped to the int32 range the metric field carries. Signed-off-by: Shangxiang Fan (from Dev Box) --- docs/genai-token-metrics.md | 24 +++-- go/adk/pkg/runner/metrics.go | 1 + go/adk/pkg/telemetry/metrics.go | 32 +++++- go/adk/pkg/telemetry/metrics_test.go | 142 ++++++++++++++++++++++++++- 4 files changed, 184 insertions(+), 15 deletions(-) diff --git a/docs/genai-token-metrics.md b/docs/genai-token-metrics.md index bd4a055922..85468164f0 100644 --- a/docs/genai-token-metrics.md +++ b/docs/genai-token-metrics.md @@ -17,7 +17,7 @@ Labels use semconv attribute names with dots mapped to underscores: | Label | Values | | --- | --- | -| `gen_ai_token_type` | `input`, `output` (output = candidate + reasoning tokens) | +| `gen_ai_token_type` | `input` (prompt tokens not served from cache), `output` (candidate + reasoning tokens), `cached` (prompt tokens served from cache; kagent extension) | | `gen_ai_operation_name` | `generate_content` for Gemini; `chat` for chat-completion providers | | `gen_ai_provider_name` | well-known value, e.g. `openai`, `anthropic`, `gcp.vertex_ai`, `aws.bedrock`, `azure.ai.openai` | | `gen_ai_request_model` | configured model, e.g. `gpt-4o` | @@ -25,10 +25,18 @@ Labels use semconv attribute names with dots mapped to underscores: | `gen_ai_agent_name` | ADK agent that made the model call (not the application name) | | `error_type` | set on failed requests; empty otherwise | -For each positive input/output count, one observation is recorded per LLM call. -Streaming partial chunks are not double-counted. Recording runs in the model callback, -so usage-only responses are counted even when they do not produce an A2A artifact. -Missing optional labels are exported as empty strings. +One observation is recorded per LLM call for each token type with a positive count +(streaming partial chunks are not double-counted). Zero or negative counts emit no observation. +Recording runs in the model callback, so usage-only responses are counted even when they +do not produce an A2A artifact. Missing optional labels are exported as empty strings. + +The token-type series are additive, so you can sum them without double counting: the +`input` series excludes prompt tokens served from cache, the `cached` series carries +exactly those tokens, and therefore `input + cached` always equals the prompt token +count the model reported (`PromptTokenCount`). Total token spend per call is +`input + cached + output`. The `cached` series reuses the same model, provider, agent, +operation, and error labels as the other series, and its `_count` includes only calls +with positive cached usage. ## Configuration @@ -89,12 +97,14 @@ chunks, and the semconv labels including `gen_ai.agent.name`, response model, an `error.type`): ```text -# HELP gen_ai_client_token_usage Measures the number of input and output tokens used by GenAI requests. +# HELP gen_ai_client_token_usage Measures the number of input, output, and cached tokens used by GenAI requests. # TYPE gen_ai_client_token_usage histogram -gen_ai_client_token_usage_sum{error_type="",gen_ai_agent_name="my_agent",gen_ai_operation_name="generate_content",gen_ai_provider_name="gcp.vertex_ai",gen_ai_request_model="gemini-2.5-flash",gen_ai_response_model="gemini-2.5-flash",gen_ai_token_type="input"} 137 +gen_ai_client_token_usage_sum{error_type="",gen_ai_agent_name="my_agent",gen_ai_operation_name="generate_content",gen_ai_provider_name="gcp.vertex_ai",gen_ai_request_model="gemini-2.5-flash",gen_ai_response_model="gemini-2.5-flash",gen_ai_token_type="input"} 107 gen_ai_client_token_usage_count{error_type="",gen_ai_agent_name="my_agent",gen_ai_operation_name="generate_content",gen_ai_provider_name="gcp.vertex_ai",gen_ai_request_model="gemini-2.5-flash",gen_ai_response_model="gemini-2.5-flash",gen_ai_token_type="input"} 2 gen_ai_client_token_usage_sum{error_type="",gen_ai_agent_name="my_agent",gen_ai_operation_name="generate_content",gen_ai_provider_name="gcp.vertex_ai",gen_ai_request_model="gemini-2.5-flash",gen_ai_response_model="gemini-2.5-flash",gen_ai_token_type="output"} 91 gen_ai_client_token_usage_count{error_type="",gen_ai_agent_name="my_agent",gen_ai_operation_name="generate_content",gen_ai_provider_name="gcp.vertex_ai",gen_ai_request_model="gemini-2.5-flash",gen_ai_response_model="gemini-2.5-flash",gen_ai_token_type="output"} 2 +gen_ai_client_token_usage_sum{error_type="",gen_ai_agent_name="my_agent",gen_ai_operation_name="generate_content",gen_ai_provider_name="gcp.vertex_ai",gen_ai_request_model="gemini-2.5-flash",gen_ai_response_model="gemini-2.5-flash",gen_ai_token_type="cached"} 30 +gen_ai_client_token_usage_count{error_type="",gen_ai_agent_name="my_agent",gen_ai_operation_name="generate_content",gen_ai_provider_name="gcp.vertex_ai",gen_ai_request_model="gemini-2.5-flash",gen_ai_response_model="gemini-2.5-flash",gen_ai_token_type="cached"} 1 ``` ## Follow-ups diff --git a/go/adk/pkg/runner/metrics.go b/go/adk/pkg/runner/metrics.go index 9c92ba06ef..9fc3ee2de6 100644 --- a/go/adk/pkg/runner/metrics.go +++ b/go/adk/pkg/runner/metrics.go @@ -33,6 +33,7 @@ func newTokenUsagePlugin(modelConfig *configv1.ADKModelConfig) (*plugin.Plugin, ErrorType: response.ErrorCode, InputTokens: int64(usage.PromptTokenCount), OutputTokens: int64(usage.CandidatesTokenCount) + int64(usage.ThoughtsTokenCount), + CachedTokens: int64(usage.CachedContentTokenCount), }) return nil, nil }, diff --git a/go/adk/pkg/telemetry/metrics.go b/go/adk/pkg/telemetry/metrics.go index 3ca5fcb9a7..60ab2712f6 100644 --- a/go/adk/pkg/telemetry/metrics.go +++ b/go/adk/pkg/telemetry/metrics.go @@ -32,6 +32,7 @@ const ( tokenTypeInput = "input" tokenTypeOutput = "output" + tokenTypeCached = "cached" // operationChat is the default gen_ai.operation.name for chat-completion calls. operationChat = "chat" @@ -52,7 +53,7 @@ var tokenUsageBuckets = []float64{1, 4, 16, 64, 256, 1024, 4096, 16384, 65536, 2 var tokenUsage = promauto.NewHistogramVec( prometheus.HistogramOpts{ Name: metricGenAIClientTokenUsage, - Help: "Measures the number of input and output tokens used by GenAI requests.", + Help: "Measures the number of input, output, and cached tokens used by GenAI requests.", Buckets: tokenUsageBuckets, }, []string{ @@ -117,15 +118,22 @@ type TokenUsage struct { // ErrorType is error.type; empty for successful requests. ErrorType string // InputTokens / OutputTokens are the token counts (output = candidate + - // reasoning tokens). Non-positive counts are skipped. + // reasoning tokens). CachedTokens are prompt tokens served from the prompt + // cache. Providers report prompt counts inconsistently, so RecordTokenUsage + // normalizes the emitted series to stay additive. Non-positive counts are skipped. InputTokens int64 OutputTokens int64 + CachedTokens int64 } -// RecordTokenUsage records input/output token counts on the +// RecordTokenUsage records input/output/cached token counts on the // gen_ai.client.token.usage histogram. If the metric pipeline is disabled // (OTEL_METRICS_ENABLED unset or not "true"), it is a no-op. Zero/negative // counts are skipped. +// +// The token-type series are additive: the input observation excludes the +// cached portion, so input + cached always equals the prompt token count the +// model reported and summing the series never double counts cache hits. func RecordTokenUsage(usage TokenUsage) { if !MetricsEnabled() { return @@ -138,12 +146,26 @@ func RecordTokenUsage(usage TokenUsage) { if responseModel == "" { responseModel = usage.RequestModel } - if usage.InputTokens > 0 { + inputTokens := usage.InputTokens + cachedTokens := usage.CachedTokens + if cachedTokens > 0 && cachedTokens <= inputTokens { + // Providers include cache hits in the prompt count (Gemini, OpenAI); + // emitting both unchanged would double count them in aggregates. + inputTokens -= cachedTokens + } + if cachedTokens < 0 { + cachedTokens = 0 + } + if inputTokens > 0 { tokenUsage.WithLabelValues(tokenTypeInput, operationName, usage.Provider, usage.RequestModel, responseModel, usage.AgentName, usage.ErrorType). - Observe(float64(usage.InputTokens)) + Observe(float64(inputTokens)) } if usage.OutputTokens > 0 { tokenUsage.WithLabelValues(tokenTypeOutput, operationName, usage.Provider, usage.RequestModel, responseModel, usage.AgentName, usage.ErrorType). Observe(float64(usage.OutputTokens)) } + if cachedTokens > 0 { + tokenUsage.WithLabelValues(tokenTypeCached, operationName, usage.Provider, usage.RequestModel, responseModel, usage.AgentName, usage.ErrorType). + Observe(float64(cachedTokens)) + } } diff --git a/go/adk/pkg/telemetry/metrics_test.go b/go/adk/pkg/telemetry/metrics_test.go index 61b76257e1..35ccf3fcc7 100644 --- a/go/adk/pkg/telemetry/metrics_test.go +++ b/go/adk/pkg/telemetry/metrics_test.go @@ -3,6 +3,7 @@ package telemetry import ( "net/http" "net/http/httptest" + "regexp" "strings" "sync" "testing" @@ -13,8 +14,143 @@ import ( "github.com/stretchr/testify/require" ) -// TestRecordTokenUsage_RecordsHistogram verifies input/output token counts are -// recorded as two separate series on the gen_ai_client_token_usage histogram. +func TestRecordTokenUsage_RecordsCachedSeries(t *testing.T) { + t.Setenv(metricsEnabledEnvVar, "true") + cases := []struct { + name string + input TokenUsage + wantResponseModel string + }{ + { + name: "response model fallback", + input: TokenUsage{ + RequestModel: "gpt-4o", Provider: "openai", AgentName: "openai-agent", + InputTokens: 100, OutputTokens: 42, CachedTokens: 30, + }, + wantResponseModel: "gpt-4o", + }, + { + name: "explicit response model and error", + input: TokenUsage{ + RequestModel: "claude-sonnet-4", ResponseModel: "claude-sonnet-4-20250514", + Provider: "anthropic", AgentName: "anthropic-agent", ErrorType: "overloaded_error", + InputTokens: 200, OutputTokens: 12, CachedTokens: 75, + }, + wantResponseModel: "claude-sonnet-4-20250514", + }, + { + name: "cached only", + input: TokenUsage{ + RequestModel: "gemini-2.5-flash", Provider: "gcp.vertex_ai", AgentName: "gemini-agent", + CachedTokens: 10, + }, + wantResponseModel: "gemini-2.5-flash", + }, + } + for _, testCase := range cases { + t.Run(testCase.name, func(t *testing.T) { + tokenUsage.Reset() + t.Cleanup(tokenUsage.Reset) + registry := prometheus.NewPedanticRegistry() + registry.MustRegister(tokenUsage) + + const workers = 8 + const callsPerWorker = 10 + var workersDone sync.WaitGroup + for range workers { + workersDone.Go(func() { + for range callsPerWorker { + RecordTokenUsage(testCase.input) + } + }) + } + workersDone.Wait() + + families, err := registry.Gather() + require.NoError(t, err) + require.Len(t, families, 1) + require.Equal(t, metricGenAIClientTokenUsage, families[0].GetName()) + // The emitted series are additive: input excludes the cached portion + // when the prompt count includes it, so input + cached equals the + // prompt tokens the model reported. + wantInput := testCase.input.InputTokens + wantCached := testCase.input.CachedTokens + if wantCached > 0 && wantCached <= wantInput { + wantInput -= wantCached + } + wantTokens := map[string]int64{} + if wantInput > 0 { + wantTokens[tokenTypeInput] = wantInput + } + if wantCached > 0 { + wantTokens[tokenTypeCached] = wantCached + } + if testCase.input.OutputTokens > 0 { + wantTokens[tokenTypeOutput] = testCase.input.OutputTokens + } + require.Len(t, families[0].GetMetric(), len(wantTokens)) + for _, series := range families[0].GetMetric() { + labels := make(map[string]string) + for _, label := range series.GetLabel() { + labels[label.GetName()] = label.GetValue() + } + tokenType := labels[labelGenAITokenType] + require.Contains(t, wantTokens, tokenType) + require.Equal(t, map[string]string{ + labelGenAITokenType: tokenType, labelGenAIOperationName: operationChat, + labelGenAIProviderName: testCase.input.Provider, labelGenAIRequestModel: testCase.input.RequestModel, + labelGenAIResponseModel: testCase.wantResponseModel, labelGenAIAgentName: testCase.input.AgentName, + labelErrorType: testCase.input.ErrorType, + }, labels) + require.NotNil(t, series.Histogram) + require.Equal(t, uint64(workers*callsPerWorker), series.GetHistogram().GetSampleCount()) + require.Equal(t, float64(workers*callsPerWorker*wantTokens[tokenType]), series.GetHistogram().GetSampleSum()) + delete(wantTokens, tokenType) + } + require.Empty(t, wantTokens) + }) + } +} + +// TestRecordTokenUsage_CachedSeriesAdditive ensures summing the token-type +// series yields the prompt total instead of double counting cache hits. +func TestRecordTokenUsage_CachedSeriesAdditive(t *testing.T) { + t.Setenv(metricsEnabledEnvVar, "true") + tokenUsage.Reset() + t.Cleanup(tokenUsage.Reset) + + // Prompt count includes cached tokens: the input series must exclude them + // so input + cached equals the 100 prompt tokens the model reported. + RecordTokenUsage(TokenUsage{RequestModel: "gpt-4o", Provider: "openai", InputTokens: 100, CachedTokens: 30}) + body := serveMetrics(t) + want := map[string]string{tokenTypeInput: "70", tokenTypeCached: "30"} + for tokenType, sum := range want { + pattern := `gen_ai_client_token_usage_sum[{][^}]*gen_ai_token_type="` + tokenType + `"[}] ([0-9]+)` + m := regexp.MustCompile(pattern).FindStringSubmatch(body) + if m == nil { + t.Fatalf("no %s series in metrics output:\n%s", tokenType, body) + } + if m[1] != sum { + t.Fatalf("%s series sum = %s, want %s\n%s", tokenType, m[1], sum, body) + } + } +} + +func TestRecordTokenUsage_SkipsNonPositiveCached(t *testing.T) { + t.Setenv(metricsEnabledEnvVar, "true") + tokenUsage.Reset() + + // Zero cached tokens emits no cached series. + RecordTokenUsage(TokenUsage{RequestModel: "gpt-4o", Provider: "openai", CachedTokens: 0, InputTokens: 10}) + // Negative cached tokens also emit no cached series. + RecordTokenUsage(TokenUsage{RequestModel: "gpt-4o", Provider: "openai", CachedTokens: -5, InputTokens: 10}) + + body := serveMetrics(t) + if strings.Contains(body, "gen_ai_token_type=\"cached\"") { + t.Fatalf("expected no cached series for non-positive CachedTokens") + } +} + func TestRecordTokenUsage_RecordsHistogram(t *testing.T) { t.Setenv(metricsEnabledEnvVar, "true") tokenUsage.Reset() @@ -44,7 +180,7 @@ func TestRecordTokenUsage_Disabled(t *testing.T) { t.Setenv(metricsEnabledEnvVar, "false") tokenUsage.Reset() - RecordTokenUsage(TokenUsage{RequestModel: "gpt-4o", Provider: "openai", InputTokens: 100}) + RecordTokenUsage(TokenUsage{RequestModel: "gpt-4o", Provider: "openai", InputTokens: 100, OutputTokens: 42, CachedTokens: 30}) if got := testutil.CollectAndCount(tokenUsage); got != 0 { t.Fatalf("expected no series when metrics are disabled, got %d", got) From 4a33f1e49c1f1d25a955986b912f673195b1dda8 Mon Sep 17 00:00:00 2001 From: "Shangxiang Fan (from Dev Box)" Date: Mon, 19 Oct 2026 12:09:15 +0800 Subject: [PATCH 3/3] fix(adk): keep cached token metrics semconv compliant Keep gen_ai.token.type limited to input and output while exposing cache reads through a separate kagent histogram. Signed-off-by: Shangxiang Fan (from Dev Box) --- docs/genai-token-metrics.md | 22 +++---- go/adk/pkg/telemetry/metrics.go | 69 ++++++++++++---------- go/adk/pkg/telemetry/metrics_test.go | 85 ++++++++++++++-------------- 3 files changed, 91 insertions(+), 85 deletions(-) diff --git a/docs/genai-token-metrics.md b/docs/genai-token-metrics.md index 85468164f0..71b71da11e 100644 --- a/docs/genai-token-metrics.md +++ b/docs/genai-token-metrics.md @@ -12,12 +12,13 @@ A Prometheus histogram, served at **`/metrics`** on the agent's HTTP port: | Prometheus name | OTel semconv | Notes | | --- | --- | --- | | `gen_ai_client_token_usage` | `gen_ai.client.token.usage` | histogram, semconv-recommended buckets | +| `kagent_gen_ai_cached_token_usage` | kagent extension | histogram for prompt-cache reads; no `gen_ai_token_type` label | Labels use semconv attribute names with dots mapped to underscores: | Label | Values | | --- | --- | -| `gen_ai_token_type` | `input` (prompt tokens not served from cache), `output` (candidate + reasoning tokens), `cached` (prompt tokens served from cache; kagent extension) | +| `gen_ai_token_type` | `input` (the full prompt count, including cache hits), `output` (candidate + reasoning tokens) | | `gen_ai_operation_name` | `generate_content` for Gemini; `chat` for chat-completion providers | | `gen_ai_provider_name` | well-known value, e.g. `openai`, `anthropic`, `gcp.vertex_ai`, `aws.bedrock`, `azure.ai.openai` | | `gen_ai_request_model` | configured model, e.g. `gpt-4o` | @@ -30,13 +31,12 @@ One observation is recorded per LLM call for each token type with a positive cou Recording runs in the model callback, so usage-only responses are counted even when they do not produce an A2A artifact. Missing optional labels are exported as empty strings. -The token-type series are additive, so you can sum them without double counting: the -`input` series excludes prompt tokens served from cache, the `cached` series carries -exactly those tokens, and therefore `input + cached` always equals the prompt token -count the model reported (`PromptTokenCount`). Total token spend per call is -`input + cached + output`. The `cached` series reuses the same model, provider, agent, -operation, and error labels as the other series, and its `_count` includes only calls -with positive cached usage. +The semconv metric keeps `gen_ai.token.type` to its allowed values, `input` and +`output`. The `input` observation is the full prompt count, including tokens served +from a provider-managed cache. `kagent_gen_ai_cached_token_usage` is a breakdown of +that input and reuses the model, provider, agent, operation, and error labels; its +`_count` includes only calls with positive cache reads. Do not add the cached metric +to `input` when calculating token spend because cache hits are already included there. ## Configuration @@ -97,14 +97,14 @@ chunks, and the semconv labels including `gen_ai.agent.name`, response model, an `error.type`): ```text -# HELP gen_ai_client_token_usage Measures the number of input, output, and cached tokens used by GenAI requests. +# HELP gen_ai_client_token_usage Measures the number of input and output tokens used by GenAI requests. # TYPE gen_ai_client_token_usage histogram gen_ai_client_token_usage_sum{error_type="",gen_ai_agent_name="my_agent",gen_ai_operation_name="generate_content",gen_ai_provider_name="gcp.vertex_ai",gen_ai_request_model="gemini-2.5-flash",gen_ai_response_model="gemini-2.5-flash",gen_ai_token_type="input"} 107 gen_ai_client_token_usage_count{error_type="",gen_ai_agent_name="my_agent",gen_ai_operation_name="generate_content",gen_ai_provider_name="gcp.vertex_ai",gen_ai_request_model="gemini-2.5-flash",gen_ai_response_model="gemini-2.5-flash",gen_ai_token_type="input"} 2 gen_ai_client_token_usage_sum{error_type="",gen_ai_agent_name="my_agent",gen_ai_operation_name="generate_content",gen_ai_provider_name="gcp.vertex_ai",gen_ai_request_model="gemini-2.5-flash",gen_ai_response_model="gemini-2.5-flash",gen_ai_token_type="output"} 91 gen_ai_client_token_usage_count{error_type="",gen_ai_agent_name="my_agent",gen_ai_operation_name="generate_content",gen_ai_provider_name="gcp.vertex_ai",gen_ai_request_model="gemini-2.5-flash",gen_ai_response_model="gemini-2.5-flash",gen_ai_token_type="output"} 2 -gen_ai_client_token_usage_sum{error_type="",gen_ai_agent_name="my_agent",gen_ai_operation_name="generate_content",gen_ai_provider_name="gcp.vertex_ai",gen_ai_request_model="gemini-2.5-flash",gen_ai_response_model="gemini-2.5-flash",gen_ai_token_type="cached"} 30 -gen_ai_client_token_usage_count{error_type="",gen_ai_agent_name="my_agent",gen_ai_operation_name="generate_content",gen_ai_provider_name="gcp.vertex_ai",gen_ai_request_model="gemini-2.5-flash",gen_ai_response_model="gemini-2.5-flash",gen_ai_token_type="cached"} 1 +kagent_gen_ai_cached_token_usage_sum{error_type="",gen_ai_agent_name="my_agent",gen_ai_operation_name="generate_content",gen_ai_provider_name="gcp.vertex_ai",gen_ai_request_model="gemini-2.5-flash",gen_ai_response_model="gemini-2.5-flash"} 30 +kagent_gen_ai_cached_token_usage_count{error_type="",gen_ai_agent_name="my_agent",gen_ai_operation_name="generate_content",gen_ai_provider_name="gcp.vertex_ai",gen_ai_request_model="gemini-2.5-flash",gen_ai_response_model="gemini-2.5-flash"} 1 ``` ## Follow-ups diff --git a/go/adk/pkg/telemetry/metrics.go b/go/adk/pkg/telemetry/metrics.go index 60ab2712f6..27c98570d8 100644 --- a/go/adk/pkg/telemetry/metrics.go +++ b/go/adk/pkg/telemetry/metrics.go @@ -12,15 +12,19 @@ import ( // GenAI token-usage instrumentation using the native Prometheus client library. // -// Metric and attribute names follow the OpenTelemetry GenAI semantic -// conventions (semconv 1.40.0), mapped to Prometheus naming (dots -> underscores): +// gen_ai.client.token.usage follows OpenTelemetry GenAI semconv 1.40.0, mapped +// to Prometheus naming (dots -> underscores): // - metric gen_ai.client.token.usage -> gen_ai_client_token_usage // - attrs gen_ai.token.type, gen_ai.request.model, gen_ai.response.model, // gen_ai.provider.name, gen_ai.agent.name, gen_ai.operation.name, error.type // +// gen_ai.token.type only permits input and output. Cache hits are a breakdown of +// the full input count, so they are exposed separately as a kagent extension. +// // https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-metrics/ const ( metricGenAIClientTokenUsage = "gen_ai_client_token_usage" + metricGenAICachedTokenUsage = "kagent_gen_ai_cached_token_usage" labelGenAITokenType = "gen_ai_token_type" labelGenAIOperationName = "gen_ai_operation_name" @@ -32,7 +36,6 @@ const ( tokenTypeInput = "input" tokenTypeOutput = "output" - tokenTypeCached = "cached" // operationChat is the default gen_ai.operation.name for chat-completion calls. operationChat = "chat" @@ -53,7 +56,7 @@ var tokenUsageBuckets = []float64{1, 4, 16, 64, 256, 1024, 4096, 16384, 65536, 2 var tokenUsage = promauto.NewHistogramVec( prometheus.HistogramOpts{ Name: metricGenAIClientTokenUsage, - Help: "Measures the number of input, output, and cached tokens used by GenAI requests.", + Help: "Measures the number of input and output tokens used by GenAI requests.", Buckets: tokenUsageBuckets, }, []string{ @@ -67,6 +70,25 @@ var tokenUsage = promauto.NewHistogramVec( }, ) +// cachedTokenUsage exposes provider cache hits as a kagent extension metric. +// Cached tokens are already included in the semconv input observation above; +// this histogram is a breakdown and must not be added to input totals. +var cachedTokenUsage = promauto.NewHistogramVec( + prometheus.HistogramOpts{ + Name: metricGenAICachedTokenUsage, + Help: "Measures prompt tokens served from provider-managed caches by GenAI requests.", + Buckets: tokenUsageBuckets, + }, + []string{ + labelGenAIOperationName, + labelGenAIProviderName, + labelGenAIRequestModel, + labelGenAIResponseModel, + labelGenAIAgentName, + labelErrorType, + }, +) + // MetricsEnabled reports whether the GenAI token metrics pipeline is on, i.e. // the OTEL_METRICS_ENABLED gate used consistently across kagent's runtimes. func MetricsEnabled() bool { @@ -117,23 +139,18 @@ type TokenUsage struct { AgentName string // ErrorType is error.type; empty for successful requests. ErrorType string - // InputTokens / OutputTokens are the token counts (output = candidate + - // reasoning tokens). CachedTokens are prompt tokens served from the prompt - // cache. Providers report prompt counts inconsistently, so RecordTokenUsage - // normalizes the emitted series to stay additive. Non-positive counts are skipped. + // InputTokens is the full prompt count (including cache hits); OutputTokens + // is candidate + reasoning tokens. CachedTokens is a breakdown of InputTokens + // served from the provider-managed prompt cache. Non-positive counts are skipped. InputTokens int64 OutputTokens int64 CachedTokens int64 } -// RecordTokenUsage records input/output/cached token counts on the -// gen_ai.client.token.usage histogram. If the metric pipeline is disabled -// (OTEL_METRICS_ENABLED unset or not "true"), it is a no-op. Zero/negative -// counts are skipped. -// -// The token-type series are additive: the input observation excludes the -// cached portion, so input + cached always equals the prompt token count the -// model reported and summing the series never double counts cache hits. +// RecordTokenUsage records input/output counts on the semconv token-usage +// histogram and positive cache hits on the kagent cached-token breakdown +// histogram. If the metric pipeline is disabled (OTEL_METRICS_ENABLED unset or +// not "true"), it is a no-op. Zero/negative counts are skipped. func RecordTokenUsage(usage TokenUsage) { if !MetricsEnabled() { return @@ -146,26 +163,16 @@ func RecordTokenUsage(usage TokenUsage) { if responseModel == "" { responseModel = usage.RequestModel } - inputTokens := usage.InputTokens - cachedTokens := usage.CachedTokens - if cachedTokens > 0 && cachedTokens <= inputTokens { - // Providers include cache hits in the prompt count (Gemini, OpenAI); - // emitting both unchanged would double count them in aggregates. - inputTokens -= cachedTokens - } - if cachedTokens < 0 { - cachedTokens = 0 - } - if inputTokens > 0 { + if usage.InputTokens > 0 { tokenUsage.WithLabelValues(tokenTypeInput, operationName, usage.Provider, usage.RequestModel, responseModel, usage.AgentName, usage.ErrorType). - Observe(float64(inputTokens)) + Observe(float64(usage.InputTokens)) } if usage.OutputTokens > 0 { tokenUsage.WithLabelValues(tokenTypeOutput, operationName, usage.Provider, usage.RequestModel, responseModel, usage.AgentName, usage.ErrorType). Observe(float64(usage.OutputTokens)) } - if cachedTokens > 0 { - tokenUsage.WithLabelValues(tokenTypeCached, operationName, usage.Provider, usage.RequestModel, responseModel, usage.AgentName, usage.ErrorType). - Observe(float64(cachedTokens)) + if usage.CachedTokens > 0 { + cachedTokenUsage.WithLabelValues(operationName, usage.Provider, usage.RequestModel, responseModel, usage.AgentName, usage.ErrorType). + Observe(float64(usage.CachedTokens)) } } diff --git a/go/adk/pkg/telemetry/metrics_test.go b/go/adk/pkg/telemetry/metrics_test.go index 35ccf3fcc7..619fd27862 100644 --- a/go/adk/pkg/telemetry/metrics_test.go +++ b/go/adk/pkg/telemetry/metrics_test.go @@ -14,7 +14,12 @@ import ( "github.com/stretchr/testify/require" ) -func TestRecordTokenUsage_RecordsCachedSeries(t *testing.T) { +func resetTokenMetrics() { + tokenUsage.Reset() + cachedTokenUsage.Reset() +} + +func TestRecordTokenUsage_RecordsInputOutputSeries(t *testing.T) { t.Setenv(metricsEnabledEnvVar, "true") cases := []struct { name string @@ -39,18 +44,18 @@ func TestRecordTokenUsage_RecordsCachedSeries(t *testing.T) { wantResponseModel: "claude-sonnet-4-20250514", }, { - name: "cached only", + name: "cache hit without output", input: TokenUsage{ RequestModel: "gemini-2.5-flash", Provider: "gcp.vertex_ai", AgentName: "gemini-agent", - CachedTokens: 10, + InputTokens: 10, CachedTokens: 10, }, wantResponseModel: "gemini-2.5-flash", }, } for _, testCase := range cases { t.Run(testCase.name, func(t *testing.T) { - tokenUsage.Reset() - t.Cleanup(tokenUsage.Reset) + resetTokenMetrics() + t.Cleanup(resetTokenMetrics) registry := prometheus.NewPedanticRegistry() registry.MustRegister(tokenUsage) @@ -70,20 +75,9 @@ func TestRecordTokenUsage_RecordsCachedSeries(t *testing.T) { require.NoError(t, err) require.Len(t, families, 1) require.Equal(t, metricGenAIClientTokenUsage, families[0].GetName()) - // The emitted series are additive: input excludes the cached portion - // when the prompt count includes it, so input + cached equals the - // prompt tokens the model reported. - wantInput := testCase.input.InputTokens - wantCached := testCase.input.CachedTokens - if wantCached > 0 && wantCached <= wantInput { - wantInput -= wantCached - } wantTokens := map[string]int64{} - if wantInput > 0 { - wantTokens[tokenTypeInput] = wantInput - } - if wantCached > 0 { - wantTokens[tokenTypeCached] = wantCached + if testCase.input.InputTokens > 0 { + wantTokens[tokenTypeInput] = testCase.input.InputTokens } if testCase.input.OutputTokens > 0 { wantTokens[tokenTypeOutput] = testCase.input.OutputTokens @@ -112,33 +106,37 @@ func TestRecordTokenUsage_RecordsCachedSeries(t *testing.T) { } } -// TestRecordTokenUsage_CachedSeriesAdditive ensures summing the token-type -// series yields the prompt total instead of double counting cache hits. -func TestRecordTokenUsage_CachedSeriesAdditive(t *testing.T) { +// TestRecordTokenUsage_CachedBreakdownMetric keeps the semconv input series as +// the full prompt total and exposes cache hits through a separate kagent metric. +func TestRecordTokenUsage_CachedBreakdownMetric(t *testing.T) { t.Setenv(metricsEnabledEnvVar, "true") - tokenUsage.Reset() - t.Cleanup(tokenUsage.Reset) + resetTokenMetrics() + t.Cleanup(resetTokenMetrics) - // Prompt count includes cached tokens: the input series must exclude them - // so input + cached equals the 100 prompt tokens the model reported. RecordTokenUsage(TokenUsage{RequestModel: "gpt-4o", Provider: "openai", InputTokens: 100, CachedTokens: 30}) body := serveMetrics(t) - want := map[string]string{tokenTypeInput: "70", tokenTypeCached: "30"} - for tokenType, sum := range want { - pattern := `gen_ai_client_token_usage_sum[{][^}]*gen_ai_token_type="` + tokenType + `"[}] ([0-9]+)` + for name, pattern := range map[string]string{ + "input": `gen_ai_client_token_usage_sum[{][^}]*gen_ai_token_type="input"[}] ([0-9]+)`, + "cached": `kagent_gen_ai_cached_token_usage_sum\{([^}]*)\} ([0-9]+)`, + } { m := regexp.MustCompile(pattern).FindStringSubmatch(body) if m == nil { - t.Fatalf("no %s series in metrics output:\n%s", tokenType, body) + t.Fatalf("no %s series in metrics output:\n%s", name, body) } - if m[1] != sum { - t.Fatalf("%s series sum = %s, want %s\n%s", tokenType, m[1], sum, body) + got := m[len(m)-1] + want := map[string]string{"input": "100", "cached": "30"}[name] + if got != want { + t.Fatalf("%s series sum = %s, want %s\n%s", name, got, want, body) } } + if strings.Contains(body, `gen_ai_token_type="cached"`) { + t.Fatalf("gen_ai.token.type must only contain input/output:\n%s", body) + } } func TestRecordTokenUsage_SkipsNonPositiveCached(t *testing.T) { t.Setenv(metricsEnabledEnvVar, "true") - tokenUsage.Reset() + resetTokenMetrics() // Zero cached tokens emits no cached series. RecordTokenUsage(TokenUsage{RequestModel: "gpt-4o", Provider: "openai", CachedTokens: 0, InputTokens: 10}) @@ -146,14 +144,14 @@ func TestRecordTokenUsage_SkipsNonPositiveCached(t *testing.T) { RecordTokenUsage(TokenUsage{RequestModel: "gpt-4o", Provider: "openai", CachedTokens: -5, InputTokens: 10}) body := serveMetrics(t) - if strings.Contains(body, "gen_ai_token_type=\"cached\"") { + if strings.Contains(body, "gen_ai_token_type=\"cached\"") || strings.Contains(body, "kagent_gen_ai_cached_token_usage") { t.Fatalf("expected no cached series for non-positive CachedTokens") } } func TestRecordTokenUsage_RecordsHistogram(t *testing.T) { t.Setenv(metricsEnabledEnvVar, "true") - tokenUsage.Reset() + resetTokenMetrics() RecordTokenUsage(TokenUsage{ RequestModel: "gpt-4o", Provider: "openai", InputTokens: 100, OutputTokens: 42, @@ -167,7 +165,7 @@ func TestRecordTokenUsage_RecordsHistogram(t *testing.T) { func TestRecordTokenUsage_SkipsZero(t *testing.T) { t.Setenv(metricsEnabledEnvVar, "true") - tokenUsage.Reset() + resetTokenMetrics() RecordTokenUsage(TokenUsage{RequestModel: "gpt-4o", Provider: "openai"}) @@ -178,7 +176,7 @@ func TestRecordTokenUsage_SkipsZero(t *testing.T) { func TestRecordTokenUsage_Disabled(t *testing.T) { t.Setenv(metricsEnabledEnvVar, "false") - tokenUsage.Reset() + resetTokenMetrics() RecordTokenUsage(TokenUsage{RequestModel: "gpt-4o", Provider: "openai", InputTokens: 100, OutputTokens: 42, CachedTokens: 30}) @@ -189,7 +187,7 @@ func TestRecordTokenUsage_Disabled(t *testing.T) { func TestRecordTokenUsage_ResponseModelFallback(t *testing.T) { t.Setenv(metricsEnabledEnvVar, "true") - tokenUsage.Reset() + resetTokenMetrics() RecordTokenUsage(TokenUsage{RequestModel: "gemini-2.5-flash", Provider: "gcp.vertex_ai", InputTokens: 3}) RecordTokenUsage(TokenUsage{RequestModel: "gemini-2.5-flash", ResponseModel: "gemini-2.5-flash-002", Provider: "gcp.vertex_ai", InputTokens: 3}) @@ -204,11 +202,11 @@ func TestRecordTokenUsage_ResponseModelFallback(t *testing.T) { func TestMetricsHandler_ServesLabels(t *testing.T) { t.Setenv(metricsEnabledEnvVar, "true") - tokenUsage.Reset() + resetTokenMetrics() RecordTokenUsage(TokenUsage{ RequestModel: "claude-3-5-sonnet", ResponseModel: "claude-3-5-sonnet-20241022", Provider: "anthropic", AgentName: "my-agent", ErrorType: "overloaded_error", - InputTokens: 10, OutputTokens: 5, + InputTokens: 10, OutputTokens: 5, CachedTokens: 3, }) body := serveMetrics(t) @@ -222,6 +220,7 @@ func TestMetricsHandler_ServesLabels(t *testing.T) { "gen_ai_response_model=\"claude-3-5-sonnet-20241022\"", "gen_ai_agent_name=\"my-agent\"", "error_type=\"overloaded_error\"", + "kagent_gen_ai_cached_token_usage_count", } { if !strings.Contains(body, want) { t.Errorf("metrics output missing %q", want) @@ -264,8 +263,8 @@ func TestRecordTokenUsageValues(t *testing.T) { } { t.Run(test.name, func(t *testing.T) { t.Setenv(metricsEnabledEnvVar, "true") - tokenUsage.Reset() - t.Cleanup(tokenUsage.Reset) + resetTokenMetrics() + t.Cleanup(resetTokenMetrics) RecordTokenUsage(test.input) RecordTokenUsage(test.input) require.Equal(t, len(test.want), testutil.CollectAndCount(tokenUsage)) @@ -291,8 +290,8 @@ func TestRecordTokenUsageValues(t *testing.T) { func TestRecordTokenUsageConcurrent(t *testing.T) { t.Setenv(metricsEnabledEnvVar, "true") - tokenUsage.Reset() - t.Cleanup(tokenUsage.Reset) + resetTokenMetrics() + t.Cleanup(resetTokenMetrics) var workers sync.WaitGroup for range 16 { workers.Go(func() {