chore(adk): address code review feedback on token usage metric

Signed-off-by: MartinForReal <fanshangxiang@gmail.com>

Author: MartinForReal

docs/genai-token-metrics.md

@@ -13,20 +13,22 @@ A Prometheus histogram, served at **`/metrics`** on the agent's HTTP port:
 | --- | --- | --- |
 | `gen_ai_client_token_usage` | `gen_ai.client.token.usage` | histogram, semconv-recommended buckets |
 
-Labels (semconv attributes, dots → underscores), aligned with what the upstream Google ADK Python
-runtime emits for the same instrument so a single dashboard works across both runtimes:
+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` | `chat` |
+| `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` | agent that produced the tokens (the kagent app name) |
+| `gen_ai_agent_name` | ADK agent that made the model call (not the application name) |
 | `error_type` | set on failed requests; empty otherwise |
 
-One observation is recorded per LLM call (streaming partial chunks are not double-counted).
+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
 
@@ -89,14 +91,14 @@ chunks, and the semconv labels including `gen_ai.agent.name`, response model, an
 ```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="chat",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="chat",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="chat",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="chat",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="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.
+  environments that push to an OTLP collector rather than scrape.

go/adk/cmd/main.go

@@ -20,7 +20,6 @@ import (
 	runnerpkg "github.com/kagent-dev/kagent/go/adk/pkg/runner"
 	"github.com/kagent-dev/kagent/go/adk/pkg/session"
 	"github.com/kagent-dev/kagent/go/adk/pkg/telemetry"
-	"github.com/kagent-dev/kagent/go/api/adk"
 	"github.com/kagent-dev/kagent/go/core/pkg/env"
 	"github.com/kagent-dev/kagent/go/pkg/logging"
 	"github.com/kagent-dev/kagent/go/pkg/tracing"
@@ -196,16 +195,13 @@ func main() {
 	}
 
 	stream := agentConfig.GetStream()
-	modelName, providerName := resolveModelLabels(agentConfig)
 	executor, err := a2a.NewKAgentExecutor(a2a.KAgentExecutorConfig{
 		RunnerConfig:   runnerConfig,
 		SessionService: sessionService,
 		Stream:         stream,
 		AppName:        appName,
 		Logger:         logger,
 		Output:         agentConfig.Output,
-		ModelName:      modelName,
-		ProviderName:   providerName,
 	})
 	if err != nil {
 		logger.Error("failed to create A2A executor", "error", err)
@@ -249,16 +245,6 @@ func main() {
 	}
 }
 
-// resolveModelLabels resolves the gen_ai.request.model / gen_ai.provider.name
-// labels for token-usage metrics from the agent config. Returns empty strings
-// when no model is configured; the metric simply omits those attributes.
-func resolveModelLabels(agentConfig *adk.AgentConfig) (model, provider string) {
-	if agentConfig == nil || agentConfig.Model == nil {
-		return "", ""
-	}
-	return config.ModelName(agentConfig.Model), telemetry.SemconvProviderName(agentConfig.Model.GetType())
-}
-
 func deriveAppName(kagentName, kagentNamespace string, agentCard *a2atype.AgentCard, logger *slog.Logger) string {
 	if kagentNamespace != "" && kagentName != "" {
 		namespace := strings.ReplaceAll(kagentNamespace, "-", "_")

go/adk/pkg/a2a/executor.go

@@ -41,11 +41,6 @@ type KAgentExecutorConfig struct {
 	AppName        string
 	Logger         *slog.Logger
 	Output         *apiadk.OutputConfig
-	// ModelName and ProviderName label GenAI token-usage metrics
-	// (gen_ai.request.model / gen_ai.provider.name). Both may be empty, in
-	// which case the corresponding metric attributes are omitted.
-	ModelName    string
-	ProviderName string
 }
 
 // KAgentExecutor keeps kagent's request/session glue around the upstream ADK
@@ -106,7 +101,6 @@ func NewKAgentExecutor(cfg KAgentExecutorConfig) (*KAgentExecutor, error) {
 				}
 				processed.Artifact.SetMeta(apia2a.TimelinePositionMetadataKey, position.UTC().Format(time.RFC3339Nano))
 			}
-			recordTokenUsage(cfg.ModelName, cfg.ProviderName, cfg.AppName, event)
 			return transformStructuredOutput(output, rootName, event, processed)
 		},
 		OutputMode: adka2a.OutputArtifactPerEvent,
@@ -208,30 +202,6 @@ func resolveStructuredOutput(config *apiadk.OutputConfig) (*structuredOutput, er
 	return &structuredOutput{schema: resolved, sha256: config.SHA256}, nil
 }
 
-// recordTokenUsage records GenAI token usage for a single ADK event on the
-// gen_ai.client.token.usage histogram. Partial (streaming) events are skipped:
-// a streamed LLM call emits many Partial chunks but usage is reported once on
-// the aggregated non-partial event, so this counts one observation per LLM
-// call, not per stream chunk. Output combines candidate + reasoning tokens.
-func recordTokenUsage(modelName, providerName, agentName string, adkEvent *adksession.Event) {
-	if adkEvent == nil {
-		return
-	}
-	um := adkEvent.UsageMetadata
-	if um == nil || adkEvent.Partial {
-		return
-	}
-	telemetry.RecordTokenUsage(telemetry.TokenUsage{
-		RequestModel:  modelName,
-		ResponseModel: adkEvent.ModelVersion,
-		Provider:      providerName,
-		AgentName:     agentName,
-		ErrorType:     adkEvent.ErrorCode,
-		InputTokens:   int64(um.PromptTokenCount),
-		OutputTokens:  int64(um.CandidatesTokenCount) + int64(um.ThoughtsTokenCount),
-	})
-}
-
 // UserIDCallInterceptor returns an a2asrv.CallInterceptor that extracts the
 // x-user-id HTTP header from the incoming request metadata and sets it as the
 // authenticated user on the CallContext.

go/adk/pkg/a2a/executor_metrics_test.go

@@ -24,6 +24,7 @@ import (
 	"google.golang.org/grpc/health"
 	grpc_health_v1 "google.golang.org/grpc/health/grpc_health_v1"
 
+	"github.com/kagent-dev/kagent/go/adk/pkg/telemetry"
 	"github.com/kagent-dev/kagent/go/pkg/tracing"
 )
 

go/adk/pkg/a2a/server/server.go

@@ -18,6 +18,7 @@ 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"
+	"github.com/stretchr/testify/require"
 	"go.opentelemetry.io/otel"
 	sdktrace "go.opentelemetry.io/otel/sdk/trace"
 	"go.opentelemetry.io/otel/sdk/trace/tracetest"
@@ -28,6 +29,44 @@ import (
 	"github.com/kagent-dev/kagent/go/adk/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)
+			telemetry.RecordTokenUsage(telemetry.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 mimics KAgentExecutor's telemetry: it starts the
 // invocation span from the request-derived context. It does not flush —
 // exporting everything (including the otelhttp server span, still open

go/adk/pkg/a2a/server/server_test.go

@@ -76,7 +76,7 @@ func GetAgentConfigSummary(config *adk.AgentConfig) string {
 
 	summary := "AgentConfig:\n"
 	if config.Model != nil {
-		summary += fmt.Sprintf("  Model: %s (%s)\n", config.Model.GetType(), getModelName(config.Model))
+		summary += fmt.Sprintf("  Model: %s (%s)\n", config.Model.GetType(), GetModelName(config.Model))
 	} else {
 		summary += "  Model: (nil)\n"
 	}
@@ -91,30 +91,32 @@ func GetAgentConfigSummary(config *adk.AgentConfig) string {
 	return summary
 }
 
-func getModelName(m adk.Model) string {
-	switch m := m.(type) {
+// GetModelName returns the configured model identifier, or "unknown" for an unrecognized model type.
+func GetModelName(model adk.Model) string {
+	switch model := model.(type) {
 	case *adk.OpenAI:
-		return m.Model
+		return model.Model
 	case *adk.AzureOpenAI:
-		return m.Model
+		return model.Model
 	case *adk.Anthropic:
-		return m.Model
+		return model.Model
 	case *adk.GeminiVertexAI:
-		return m.Model
+		return model.Model
 	case *adk.GeminiAnthropic:
-		return m.Model
+		return model.Model
 	case *adk.Ollama:
-		return m.Model
+		return model.Model
 	case *adk.Gemini:
-		return m.Model
+		return model.Model
+	case *adk.Bedrock:
+		return model.Model
+	case *adk.SAPAICore:
+		return model.Model
+	case *adk.Foundry:
+		return model.Model
+	case *adk.GenericModel:
+		return model.Model
 	default:
 		return "unknown"
 	}
 }
-
-// ModelName returns the configured model identifier (e.g. "gpt-4o",
-// "claude-3-5-sonnet"), or "unknown" for an unrecognized model type. It is the
-// exported form of getModelName, used to label GenAI telemetry.
-func ModelName(m adk.Model) string {
-	return getModelName(m)
-}

go/adk/pkg/config/config_usage.go

@@ -5,8 +5,36 @@ import (
 	"testing"
 
 	"github.com/kagent-dev/kagent/go/api/adk"
+	"github.com/stretchr/testify/require"
 )
 
+func TestGetModelName(t *testing.T) {
+	base := adk.BaseModel{Model: "configured-model"}
+	for _, test := range []struct {
+		name  string
+		input adk.Model
+		want  string
+	}{
+		{name: "openai", input: &adk.OpenAI{BaseModel: base}, want: base.Model},
+		{name: "azure", input: &adk.AzureOpenAI{BaseModel: base}, want: base.Model},
+		{name: "anthropic", input: &adk.Anthropic{BaseModel: base}, want: base.Model},
+		{name: "gemini", input: &adk.Gemini{BaseModel: base}, want: base.Model},
+		{name: "vertex", input: &adk.GeminiVertexAI{BaseModel: base}, want: base.Model},
+		{name: "vertex anthropic", input: &adk.GeminiAnthropic{BaseModel: base}, want: base.Model},
+		{name: "ollama", input: &adk.Ollama{BaseModel: base}, want: base.Model},
+		{name: "bedrock", input: &adk.Bedrock{BaseModel: base}, want: base.Model},
+		{name: "sap", input: &adk.SAPAICore{BaseModel: base}, want: base.Model},
+		{name: "foundry", input: &adk.Foundry{BaseModel: base}, want: base.Model},
+		{name: "generic", input: &adk.GenericModel{BaseModel: base}, want: base.Model},
+		{name: "nil", want: "unknown"},
+		{name: "empty", input: &adk.OpenAI{}},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			require.Equal(t, test.want, GetModelName(test.input))
+		})
+	}
+}
+
 func TestValidateAgentConfigUsage_NilConfig(t *testing.T) {
 	err := ValidateAgentConfigUsage(nil)
 	if err == nil {

go/adk/pkg/config/config_usage_test.go

@@ -82,6 +82,11 @@ func CreateRunnerConfig(
 	}
 
 	var adkPlugins []*adkplugin.Plugin
+	metricsPlugin, err := newTokenUsagePlugin(agentConfig.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 {