Add telemetry

2026-04-14 10:31:56 +02:00
parent 20049881ad
commit 9cc052c162
15 changed files with 580 additions and 62 deletions
@@ -9,9 +9,12 @@ import (
 	"github.com/gin-gonic/gin"
 	"github.com/rs/zerolog/log"
 	"github.com/tidwall/gjson"
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/metric"

 	"github.com/fujin/anthropic-proxy/internal/auth"
 	"github.com/fujin/anthropic-proxy/internal/logging"
+	"github.com/fujin/anthropic-proxy/internal/telemetry"
 )

 func HandleMessages(pool *auth.Pool, profile *SniffedProfile, getSanitizer func() *Sanitizer) gin.HandlerFunc {
@@ -55,10 +58,16 @@ func HandleMessages(pool *auth.Pool, profile *SniffedProfile, getSanitizer func(

 func handleNonStream(c *gin.Context, upstream *UpstreamClient, san *Sanitizer, pool *auth.Pool, cred *auth.Credential, body []byte, originalBody []byte) {
 	startTime := time.Now()
-	respBody, headers, statusCode, err := upstream.Execute(c.Request.Context(), cred, body)
+	model := gjson.GetBytes(body, "model").String()
+	ctx := c.Request.Context()
+
+	telemetry.RequestBodySize.Record(ctx, int64(len(body)),
+		metric.WithAttributes(attribute.String("model", model), attribute.Bool("stream", false)))
+
+	respBody, headers, statusCode, err := upstream.Execute(ctx, cred, body)
+	latencyMs := float64(time.Since(startTime).Milliseconds())
+
 	if err != nil {
-		latencyMs := float64(time.Since(startTime).Milliseconds())
-		model := gjson.GetBytes(body, "model").String()
 		log.Error().
 			Err(err).
 			Str("credential", cred.Email).
@@ -69,14 +78,39 @@ func handleNonStream(c *gin.Context, upstream *UpstreamClient, san *Sanitizer, p
 			Int("request_body_size", len(body)).
 			Float64("latency_ms", latencyMs).
 			Msg("upstream connection error")
+
+		telemetry.UpstreamErrors.Add(ctx, 1,
+			metric.WithAttributes(
+				attribute.String("error_type", "connection"),
+				attribute.String("credential", cred.Email),
+				attribute.Int("status_code", http.StatusBadGateway),
+			))
+		telemetry.RequestCounter.Add(ctx, 1,
+			metric.WithAttributes(
+				attribute.String("model", model),
+				attribute.Bool("stream", false),
+				attribute.Int("status_code", http.StatusBadGateway),
+			))
+		telemetry.RequestDuration.Record(ctx, latencyMs,
+			metric.WithAttributes(attribute.String("model", model), attribute.Bool("stream", false), attribute.Int("status_code", http.StatusBadGateway)))
+
 		c.JSON(http.StatusBadGateway, gin.H{"error": "upstream request failed"})
 		return
 	}

+	attrs := []attribute.KeyValue{
+		attribute.String("model", model),
+		attribute.Bool("stream", false),
+		attribute.Int("status_code", statusCode),
+	}
+
+	telemetry.RequestCounter.Add(ctx, 1, metric.WithAttributes(attrs...))
+	telemetry.RequestDuration.Record(ctx, latencyMs, metric.WithAttributes(attrs...))
+
 	if statusCode >= 400 {
 		pool.MarkFailure(cred, statusCode)
-		latencyMs := float64(time.Since(startTime).Milliseconds())
-		model := gjson.GetBytes(body, "model").String()
+		telemetry.CredentialCooldowns.Add(ctx, 1,
+			metric.WithAttributes(attribute.Int("status_code", statusCode)))
 		errorType := gjson.GetBytes(respBody, "error.type").String()
 		errorMessage := gjson.GetBytes(respBody, "error.message").String()
 		log.Error().
@@ -94,9 +128,33 @@ func handleNonStream(c *gin.Context, upstream *UpstreamClient, san *Sanitizer, p
 			Int("request_body_size", len(body)).
 			Str("request_headers", logging.RedactHeaders(c.Request.Header)).
 			Msg("upstream error")
+
+		telemetry.UpstreamErrors.Add(ctx, 1,
+			metric.WithAttributes(
+				attribute.Int("status_code", statusCode),
+				attribute.String("error_type", errorType),
+				attribute.String("credential", cred.Email),
+			))
 	} else {
 		pool.MarkSuccess(cred)
 		respBody = san.DesanitizeResponse(respBody)
+
+		inputTokens := gjson.GetBytes(respBody, "usage.input_tokens").Int()
+		outputTokens := gjson.GetBytes(respBody, "usage.output_tokens").Int()
+		tokenAttrs := metric.WithAttributes(
+			attribute.String("model", model),
+			attribute.String("credential", cred.Email),
+		)
+		telemetry.TokensInput.Add(ctx, inputTokens, tokenAttrs)
+		telemetry.TokensOutput.Add(ctx, outputTokens, tokenAttrs)
+
+		log.Info().
+			Int("status", statusCode).
+			Float64("latency_ms", latencyMs).
+			Str("model", model).
+			Int64("input_tokens", inputTokens).
+			Int64("output_tokens", outputTokens).
+			Msg("request completed")
 	}

 	for _, h := range []string{"Content-Type", "X-Request-Id"} {
@@ -110,10 +168,16 @@ func handleNonStream(c *gin.Context, upstream *UpstreamClient, san *Sanitizer, p

 func handleStream(c *gin.Context, upstream *UpstreamClient, san *Sanitizer, pool *auth.Pool, cred *auth.Credential, body []byte, originalBody []byte) {
 	startTime := time.Now()
-	resp, err := upstream.ExecuteStream(c.Request.Context(), cred, body)
+	model := gjson.GetBytes(body, "model").String()
+	ctx := c.Request.Context()
+
+	telemetry.StreamRequests.Add(ctx, 1, metric.WithAttributes(attribute.String("model", model)))
+	telemetry.RequestBodySize.Record(ctx, int64(len(body)),
+		metric.WithAttributes(attribute.String("model", model), attribute.Bool("stream", true)))
+
+	resp, err := upstream.ExecuteStream(ctx, cred, body)
 	if err != nil {
 		latencyMs := float64(time.Since(startTime).Milliseconds())
-		model := gjson.GetBytes(body, "model").String()
 		log.Error().
 			Err(err).
 			Str("credential", cred.Email).
@@ -124,6 +188,22 @@ func handleStream(c *gin.Context, upstream *UpstreamClient, san *Sanitizer, pool
 			Int("request_body_size", len(body)).
 			Float64("latency_ms", latencyMs).
 			Msg("upstream connection error")
+
+		telemetry.UpstreamErrors.Add(ctx, 1,
+			metric.WithAttributes(
+				attribute.String("error_type", "connection"),
+				attribute.String("credential", cred.Email),
+				attribute.Int("status_code", http.StatusBadGateway),
+			))
+		telemetry.RequestCounter.Add(ctx, 1,
+			metric.WithAttributes(
+				attribute.String("model", model),
+				attribute.Bool("stream", true),
+				attribute.Int("status_code", http.StatusBadGateway),
+			))
+		telemetry.RequestDuration.Record(ctx, latencyMs,
+			metric.WithAttributes(attribute.String("model", model), attribute.Bool("stream", true), attribute.Int("status_code", http.StatusBadGateway)))
+
 		c.JSON(http.StatusBadGateway, gin.H{"error": "upstream stream request failed"})
 		return
 	}
@@ -131,9 +211,10 @@ func handleStream(c *gin.Context, upstream *UpstreamClient, san *Sanitizer, pool

 	if resp.StatusCode >= 400 {
 		pool.MarkFailure(cred, resp.StatusCode)
+		telemetry.CredentialCooldowns.Add(ctx, 1,
+			metric.WithAttributes(attribute.Int("status_code", resp.StatusCode)))
 		respBody, _ := io.ReadAll(resp.Body)
 		latencyMs := float64(time.Since(startTime).Milliseconds())
-		model := gjson.GetBytes(body, "model").String()
 		errorType := gjson.GetBytes(respBody, "error.type").String()
 		errorMessage := gjson.GetBytes(respBody, "error.message").String()
 		log.Error().
@@ -151,6 +232,21 @@ func handleStream(c *gin.Context, upstream *UpstreamClient, san *Sanitizer, pool
 			Int("request_body_size", len(body)).
 			Str("request_headers", logging.RedactHeaders(c.Request.Header)).
 			Msg("upstream error")
+
+		attrs := []attribute.KeyValue{
+			attribute.String("model", model),
+			attribute.Bool("stream", true),
+			attribute.Int("status_code", resp.StatusCode),
+		}
+		telemetry.RequestCounter.Add(ctx, 1, metric.WithAttributes(attrs...))
+		telemetry.RequestDuration.Record(ctx, latencyMs, metric.WithAttributes(attrs...))
+		telemetry.UpstreamErrors.Add(ctx, 1,
+			metric.WithAttributes(
+				attribute.Int("status_code", resp.StatusCode),
+				attribute.String("error_type", errorType),
+				attribute.String("credential", cred.Email),
+			))
+
 		c.Data(resp.StatusCode, resp.Header.Get("Content-Type"), respBody)
 		return
 	}
@@ -169,14 +265,50 @@ func handleStream(c *gin.Context, upstream *UpstreamClient, san *Sanitizer, pool
 		return
 	}

+	var inputTokens, outputTokens int64
 	scanner := bufio.NewScanner(resp.Body)
 	scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
 	for scanner.Scan() {
 		line := san.DesanitizeStreamEvent(scanner.Text())
 		c.Writer.WriteString(line + "\n")
 		flusher.Flush()
+
+		// Extract token usage from message_delta event
+		if len(line) > 5 && line[:5] == "data:" {
+			data := line[5:]
+			if gjson.Get(data, "type").String() == "message_delta" {
+				inputTokens = gjson.Get(data, "usage.input_tokens").Int()
+				outputTokens = gjson.Get(data, "usage.output_tokens").Int()
+			}
+		}
 	}

+	latencyMs := float64(time.Since(startTime).Milliseconds())
+	attrs := []attribute.KeyValue{
+		attribute.String("model", model),
+		attribute.Bool("stream", true),
+		attribute.Int("status_code", http.StatusOK),
+	}
+	telemetry.RequestCounter.Add(ctx, 1, metric.WithAttributes(attrs...))
+	telemetry.RequestDuration.Record(ctx, latencyMs, metric.WithAttributes(attrs...))
+
+	if inputTokens > 0 || outputTokens > 0 {
+		tokenAttrs := metric.WithAttributes(
+			attribute.String("model", model),
+			attribute.String("credential", cred.Email),
+		)
+		telemetry.TokensInput.Add(ctx, inputTokens, tokenAttrs)
+		telemetry.TokensOutput.Add(ctx, outputTokens, tokenAttrs)
+	}
+
+	log.Info().
+		Float64("latency_ms", latencyMs).
+		Str("model", model).
+		Bool("stream", true).
+		Int64("input_tokens", inputTokens).
+		Int64("output_tokens", outputTokens).
+		Msg("stream completed")
+
 	if err := scanner.Err(); err != nil {
 		log.Error().Err(err).Msg("stream scan error")
 	}