Skip to content

Commit 250e790

Browse files
pawbanapawbana
authored
feat: add client label to InterceptionCount, PromptCount and TokenUseCount metrics (#209)
1 parent b01456d commit 250e790

4 files changed

Lines changed: 96 additions & 25 deletions

File tree

‎bridge.go‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -190,6 +190,7 @@ func newInterceptionProcessor(p provider.Provider, cbs *circuitbreaker.ProviderC
190190
asyncRecorder.WithProvider(p.Name())
191191
asyncRecorder.WithModel(interceptor.Model())
192192
asyncRecorder.WithInitiatorID(actor.ID)
193+
asyncRecorder.WithClient(string(client))
193194
interceptor.Setup(logger, asyncRecorder, mcpProxy)
194195

195196
if err := rec.RecordInterception(ctx, &recorder.InterceptionRecord{
@@ -231,13 +232,13 @@ func newInterceptionProcessor(p provider.Provider, cbs *circuitbreaker.ProviderC
231232
return interceptor.ProcessRequest(rw, r)
232233
}); err != nil {
233234
if m != nil {
234-
m.InterceptionCount.WithLabelValues(p.Name(), interceptor.Model(), metrics.InterceptionCountStatusFailed, route, r.Method, actor.ID).Add(1)
235+
m.InterceptionCount.WithLabelValues(p.Name(), interceptor.Model(), metrics.InterceptionCountStatusFailed, route, r.Method, actor.ID, string(client)).Add(1)
235236
}
236237
span.SetStatus(codes.Error, fmt.Sprintf("interception failed: %v", err))
237238
log.Warn(ctx, "interception failed", slog.Error(err))
238239
} else {
239240
if m != nil {
240-
m.InterceptionCount.WithLabelValues(p.Name(), interceptor.Model(), metrics.InterceptionCountStatusCompleted, route, r.Method, actor.ID).Add(1)
241+
m.InterceptionCount.WithLabelValues(p.Name(), interceptor.Model(), metrics.InterceptionCountStatusCompleted, route, r.Method, actor.ID, string(client)).Add(1)
241242
}
242243
log.Debug(ctx, "interception ended")
243244
}

‎internal/integrationtest/metrics_test.go‎

Lines changed: 67 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -25,10 +25,12 @@ func TestMetrics_Interception(t *testing.T) {
2525
name string
2626
fixture []byte
2727
path string
28+
headers http.Header
2829
expectStatus string
2930
expectModel string
3031
expectRoute string
3132
expectProvider string
33+
expectClient aibridge.Client
3234
allowOverflow bool // error fixtures may cause retries
3335
}{
3436
{
@@ -39,72 +41,98 @@ func TestMetrics_Interception(t *testing.T) {
3941
expectModel: "claude-sonnet-4-0",
4042
expectRoute: "/v1/messages",
4143
expectProvider: config.ProviderAnthropic,
44+
expectClient: aibridge.ClientUnknown,
4245
},
4346
{
4447
name: "ant_error",
4548
fixture: fixtures.AntNonStreamError,
4649
path: pathAnthropicMessages,
50+
headers: http.Header{"User-Agent": []string{"kilo-code/1.2.3"}},
4751
expectStatus: metrics.InterceptionCountStatusFailed,
4852
expectModel: "claude-sonnet-4-0",
4953
expectRoute: "/v1/messages",
5054
expectProvider: config.ProviderAnthropic,
55+
expectClient: aibridge.ClientKilo,
5156
allowOverflow: true,
5257
},
58+
{
59+
name: "ant_simple_claude_code",
60+
fixture: fixtures.AntSimple,
61+
path: pathAnthropicMessages,
62+
headers: http.Header{"User-Agent": []string{"claude-code/1.0.0"}},
63+
expectStatus: metrics.InterceptionCountStatusCompleted,
64+
expectModel: "claude-sonnet-4-0",
65+
expectRoute: "/v1/messages",
66+
expectProvider: config.ProviderAnthropic,
67+
expectClient: aibridge.ClientClaudeCode,
68+
},
5369
{
5470
name: "oai_chat_simple",
5571
fixture: fixtures.OaiChatSimple,
5672
path: pathOpenAIChatCompletions,
73+
headers: http.Header{"User-Agent": []string{"copilot/1.0.0"}},
5774
expectStatus: metrics.InterceptionCountStatusCompleted,
5875
expectModel: "gpt-4.1",
5976
expectRoute: "/v1/chat/completions",
6077
expectProvider: config.ProviderOpenAI,
78+
expectClient: aibridge.ClientCopilotCLI,
6179
},
6280
{
6381
name: "oai_chat_error",
6482
fixture: fixtures.OaiChatNonStreamError,
6583
path: pathOpenAIChatCompletions,
84+
headers: http.Header{"User-Agent": []string{"githubcopilotchat/0.30.0"}},
6685
expectStatus: metrics.InterceptionCountStatusFailed,
6786
expectModel: "gpt-4.1",
6887
expectRoute: "/v1/chat/completions",
6988
expectProvider: config.ProviderOpenAI,
89+
expectClient: aibridge.ClientCopilotVSC,
7090
allowOverflow: true,
7191
},
7292
{
7393
name: "oai_responses_blocking_simple",
7494
fixture: fixtures.OaiResponsesBlockingSimple,
7595
path: pathOpenAIResponses,
96+
headers: http.Header{"X-Cursor-Client-Version": []string{"0.50.0"}},
7697
expectStatus: metrics.InterceptionCountStatusCompleted,
7798
expectModel: "gpt-4o-mini",
7899
expectRoute: "/v1/responses",
79100
expectProvider: config.ProviderOpenAI,
101+
expectClient: aibridge.ClientCursor,
80102
},
81103
{
82104
name: "oai_responses_blocking_error",
83105
fixture: fixtures.OaiResponsesBlockingHttpErr,
84106
path: pathOpenAIResponses,
107+
headers: http.Header{"User-Agent": []string{"codex/1.0.0"}},
85108
expectStatus: metrics.InterceptionCountStatusFailed,
86109
expectModel: "gpt-4o-mini",
87110
expectRoute: "/v1/responses",
88111
expectProvider: config.ProviderOpenAI,
112+
expectClient: aibridge.ClientCodex,
89113
allowOverflow: true,
90114
},
91115
{
92116
name: "oai_responses_streaming_simple",
93117
fixture: fixtures.OaiResponsesStreamingSimple,
94118
path: pathOpenAIResponses,
119+
headers: http.Header{"User-Agent": []string{"zed/0.200.0"}},
95120
expectStatus: metrics.InterceptionCountStatusCompleted,
96121
expectModel: "gpt-4o-mini",
97122
expectRoute: "/v1/responses",
98123
expectProvider: config.ProviderOpenAI,
124+
expectClient: aibridge.ClientZed,
99125
},
100126
{
101127
name: "oai_responses_streaming_error",
102128
fixture: fixtures.OaiResponsesStreamingHttpErr,
103129
path: pathOpenAIResponses,
130+
headers: http.Header{"Originator": []string{"roo-code"}},
104131
expectStatus: metrics.InterceptionCountStatusFailed,
105132
expectModel: "gpt-4o-mini",
106133
expectRoute: "/v1/responses",
107134
expectProvider: config.ProviderOpenAI,
135+
expectClient: aibridge.ClientRoo,
108136
allowOverflow: true,
109137
},
110138
}
@@ -125,12 +153,12 @@ func TestMetrics_Interception(t *testing.T) {
125153
withMetrics(m),
126154
)
127155

128-
resp := bridgeServer.makeRequest(t, http.MethodPost, tc.path, fix.Request())
156+
resp := bridgeServer.makeRequest(t, http.MethodPost, tc.path, fix.Request(), tc.headers)
129157
_, err := io.ReadAll(resp.Body)
130158
require.NoError(t, err)
131159

132160
count := promtest.ToFloat64(m.InterceptionCount.WithLabelValues(
133-
tc.expectProvider, tc.expectModel, tc.expectStatus, tc.expectRoute, "POST", defaultActorID))
161+
tc.expectProvider, tc.expectModel, tc.expectStatus, tc.expectRoute, "POST", defaultActorID, string(tc.expectClient)))
134162
require.Equal(t, 1.0, count)
135163
require.Equal(t, 1, promtest.CollectAndCount(m.InterceptionDuration))
136164
require.Equal(t, 1, promtest.CollectAndCount(m.InterceptionCount))
@@ -229,16 +257,51 @@ func TestMetrics_PromptCount(t *testing.T) {
229257
withMetrics(m),
230258
)
231259

232-
resp := bridgeServer.makeRequest(t, http.MethodPost, pathOpenAIChatCompletions, fix.Request())
260+
resp := bridgeServer.makeRequest(t, http.MethodPost, pathOpenAIChatCompletions, fix.Request(), http.Header{"User-Agent": []string{"claude-code/1.0.0"}})
233261
require.Equal(t, http.StatusOK, resp.StatusCode)
234262
_, err := io.ReadAll(resp.Body)
235263
require.NoError(t, err)
236264

237265
prompts := promtest.ToFloat64(m.PromptCount.WithLabelValues(
238-
config.ProviderOpenAI, "gpt-4.1", defaultActorID))
266+
config.ProviderOpenAI, "gpt-4.1", defaultActorID, string(aibridge.ClientClaudeCode)))
239267
require.Equal(t, 1.0, prompts)
240268
}
241269

270+
func TestMetrics_TokenUseCount(t *testing.T) {
271+
t.Parallel()
272+
273+
ctx, cancel := context.WithTimeout(t.Context(), time.Second*30)
274+
t.Cleanup(cancel)
275+
276+
fix := fixtures.Parse(t, fixtures.OaiResponsesBlockingCachedInputTokens)
277+
upstream := newMockUpstream(t, ctx, newFixtureResponse(fix))
278+
279+
m := aibridge.NewMetrics(prometheus.NewRegistry())
280+
bridgeServer := newBridgeTestServer(t, ctx, upstream.URL,
281+
withMetrics(m),
282+
)
283+
284+
resp := bridgeServer.makeRequest(t, http.MethodPost, pathOpenAIResponses, fix.Request(),
285+
http.Header{"User-Agent": []string{"claude-code/1.0.0"}})
286+
require.Equal(t, http.StatusOK, resp.StatusCode)
287+
_, _ = io.ReadAll(resp.Body)
288+
289+
clientLabel := string(aibridge.ClientClaudeCode)
290+
// Token metrics are recorded asynchronously; wait for them to appear.
291+
require.Eventually(t, func() bool {
292+
return promtest.ToFloat64(m.TokenUseCount.WithLabelValues(
293+
config.ProviderOpenAI, "gpt-4.1", "input", defaultActorID, clientLabel)) > 0
294+
}, time.Second*10, time.Millisecond*50)
295+
296+
require.Equal(t, 129.0, promtest.ToFloat64(m.TokenUseCount.WithLabelValues(config.ProviderOpenAI, "gpt-4.1", "input", defaultActorID, clientLabel))) // 12033 - 11904 (cached)
297+
require.Equal(t, 44.0, promtest.ToFloat64(m.TokenUseCount.WithLabelValues(config.ProviderOpenAI, "gpt-4.1", "output", defaultActorID, clientLabel)))
298+
299+
// ExtraTokenTypes
300+
require.Equal(t, 11904.0, promtest.ToFloat64(m.TokenUseCount.WithLabelValues(config.ProviderOpenAI, "gpt-4.1", "input_cached", defaultActorID, clientLabel)))
301+
require.Equal(t, 0.0, promtest.ToFloat64(m.TokenUseCount.WithLabelValues(config.ProviderOpenAI, "gpt-4.1", "output_reasoning", defaultActorID, clientLabel)))
302+
require.Equal(t, 12077.0, promtest.ToFloat64(m.TokenUseCount.WithLabelValues(config.ProviderOpenAI, "gpt-4.1", "total_tokens", defaultActorID, clientLabel)))
303+
}
304+
242305
func TestMetrics_NonInjectedToolUseCount(t *testing.T) {
243306
t.Parallel()
244307

‎metrics/metrics.go‎

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -42,20 +42,20 @@ func NewMetrics(reg prometheus.Registerer) *Metrics {
4242
return &Metrics{
4343
// Interception-related metrics.
4444

45-
// Pessimistic cardinality: 2 providers, 5 models, 2 statuses, 2 routes, 3 methods = up to 120 PER INITIATOR.
45+
// Pessimistic cardinality: 3 providers, 5 models, 2 statuses, 3 routes, 3 methods, 10 clients = up to 2700 PER INITIATOR.
4646
InterceptionCount: promauto.With(reg).NewCounterVec(prometheus.CounterOpts{
4747
Subsystem: "interceptions",
4848
Name: "total",
4949
Help: "The count of intercepted requests.",
50-
}, append(baseLabels, "status", "route", "method", "initiator_id")),
51-
// Pessimistic cardinality: 2 providers, 5 models, 2 routes = up to 20.
50+
}, append(baseLabels, "status", "route", "method", "initiator_id", "client")),
51+
// Pessimistic cardinality: 3 providers, 5 models, 3 routes = up to 45.
5252
// NOTE: route is not unbounded because this is only for intercepted routes.
5353
InterceptionsInflight: promauto.With(reg).NewGaugeVec(prometheus.GaugeOpts{
5454
Subsystem: "interceptions",
5555
Name: "inflight",
5656
Help: "The number of intercepted requests which are being processed.",
5757
}, append(baseLabels, "route")),
58-
// Pessimistic cardinality: 2 providers, 5 models, 7 buckets + 3 extra series (count, sum, +Inf) = up to 100.
58+
// Pessimistic cardinality: 3 providers, 5 models, 7 buckets + 3 extra series (count, sum, +Inf) = up to 150.
5959
InterceptionDuration: promauto.With(reg).NewHistogramVec(prometheus.HistogramOpts{
6060
Subsystem: "interceptions",
6161
Name: "duration_seconds",
@@ -67,7 +67,7 @@ func NewMetrics(reg prometheus.Registerer) *Metrics {
6767
Buckets: []float64{0.5, 2, 5, 15, 30, 60, 120},
6868
}, baseLabels),
6969

70-
// Pessimistic cardinality: 2 providers, 10 routes, 3 methods = up to 60.
70+
// Pessimistic cardinality: 3 providers, 10 routes, 3 methods = up to 90.
7171
// NOTE: route is not unbounded because PassthroughRoutes (see provider.go) is a static list.
7272
PassthroughCount: promauto.With(reg).NewCounterVec(prometheus.CounterOpts{
7373
Subsystem: "passthrough",
@@ -77,31 +77,31 @@ func NewMetrics(reg prometheus.Registerer) *Metrics {
7777

7878
// Prompt-related metrics.
7979

80-
// Pessimistic cardinality: 2 providers, 5 models = up to 10 PER INITIATOR.
80+
// Pessimistic cardinality: 3 providers, 5 models, 10 clients = up to 150 PER INITIATOR.
8181
PromptCount: promauto.With(reg).NewCounterVec(prometheus.CounterOpts{
8282
Subsystem: "prompts",
8383
Name: "total",
8484
Help: "The number of prompts issued by users (initiators).",
85-
}, append(baseLabels, "initiator_id")),
85+
}, append(baseLabels, "initiator_id", "client")),
8686

8787
// Token-related metrics.
8888

89-
// Pessimistic cardinality: 2 providers, 5 models, 10 types = up to 100 PER INITIATOR.
89+
// Pessimistic cardinality: 3 providers, 5 models, 10 types, 10 clients = up to 1500 PER INITIATOR.
9090
TokenUseCount: promauto.With(reg).NewCounterVec(prometheus.CounterOpts{
9191
Subsystem: "tokens",
9292
Name: "total",
9393
Help: "The number of tokens used by intercepted requests.",
94-
}, append(baseLabels, "type", "initiator_id")),
94+
}, append(baseLabels, "type", "initiator_id", "client")),
9595

9696
// Tool-related metrics.
9797

98-
// Pessimistic cardinality: 2 providers, 5 models, 3 servers, 30 tools = up to 900.
98+
// Pessimistic cardinality: 3 providers, 5 models, 3 servers, 30 tools = up to 1350.
9999
InjectedToolUseCount: promauto.With(reg).NewCounterVec(prometheus.CounterOpts{
100100
Subsystem: "injected_tool_invocations",
101101
Name: "total",
102102
Help: "The number of times an injected MCP tool was invoked by aibridge.",
103103
}, append(baseLabels, "server", "name")),
104-
// Pessimistic cardinality: 2 providers, 5 models, 30 tools = up to 300.
104+
// Pessimistic cardinality: 3 providers, 5 models, 30 tools = up to 450.
105105
NonInjectedToolUseCount: promauto.With(reg).NewCounterVec(prometheus.CounterOpts{
106106
Subsystem: "non_injected_tool_selections",
107107
Name: "total",
@@ -110,19 +110,19 @@ func NewMetrics(reg prometheus.Registerer) *Metrics {
110110

111111
// Circuit breaker metrics.
112112

113-
// Pessimistic cardinality: 2 providers, 2 endpoints, 5 models = up to 20.
113+
// Pessimistic cardinality: 3 providers, 2 endpoints, 5 models = up to 30.
114114
CircuitBreakerState: promauto.With(reg).NewGaugeVec(prometheus.GaugeOpts{
115115
Subsystem: "circuit_breaker",
116116
Name: "state",
117117
Help: "Current state of the circuit breaker (0=closed, 0.5=half-open, 1=open).",
118118
}, []string{"provider", "endpoint", "model"}),
119-
// Pessimistic cardinality: 2 providers, 2 endpoints, 5 models = up to 20.
119+
// Pessimistic cardinality: 3 providers, 2 endpoints, 5 models = up to 30.
120120
CircuitBreakerTrips: promauto.With(reg).NewCounterVec(prometheus.CounterOpts{
121121
Subsystem: "circuit_breaker",
122122
Name: "trips_total",
123123
Help: "Total number of times the circuit breaker transitioned to open state.",
124124
}, []string{"provider", "endpoint", "model"}),
125-
// Pessimistic cardinality: 2 providers, 2 endpoints, 5 models = up to 20.
125+
// Pessimistic cardinality: 3 providers, 2 endpoints, 5 models = up to 30.
126126
CircuitBreakerRejects: promauto.With(reg).NewCounterVec(prometheus.CounterOpts{
127127
Subsystem: "circuit_breaker",
128128
Name: "rejects_total",

‎recorder/recorder.go‎

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -131,7 +131,10 @@ type AsyncRecorder struct {
131131
timeout time.Duration
132132
metrics *metrics.Metrics
133133

134-
provider, model, initiatorID string
134+
provider string
135+
model string
136+
initiatorID string
137+
client string
135138

136139
wg sync.WaitGroup
137140
}
@@ -158,6 +161,10 @@ func (a *AsyncRecorder) WithInitiatorID(initiatorID string) {
158161
a.initiatorID = initiatorID
159162
}
160163

164+
func (a *AsyncRecorder) WithClient(client string) {
165+
a.client = client
166+
}
167+
161168
// RecordInterception must NOT be called asynchronously.
162169
// If an interception cannot be recorded, the whole request should fail.
163170
func (a *AsyncRecorder) RecordInterception(ctx context.Context, req *InterceptionRecord) error {
@@ -193,7 +200,7 @@ func (a *AsyncRecorder) RecordPromptUsage(ctx context.Context, req *PromptUsageR
193200
}
194201

195202
if a.metrics != nil && req.Prompt != "" { // TODO: will be irrelevant once https://github.com/coder/aibridge/issues/55 is fixed.
196-
a.metrics.PromptCount.WithLabelValues(a.provider, a.model, a.initiatorID).Add(1)
203+
a.metrics.PromptCount.WithLabelValues(a.provider, a.model, a.initiatorID, a.client).Add(1)
197204
}
198205
}()
199206

@@ -213,10 +220,10 @@ func (a *AsyncRecorder) RecordTokenUsage(ctx context.Context, req *TokenUsageRec
213220
}
214221

215222
if a.metrics != nil {
216-
a.metrics.TokenUseCount.WithLabelValues(a.provider, a.model, "input", a.initiatorID).Add(float64(req.Input))
217-
a.metrics.TokenUseCount.WithLabelValues(a.provider, a.model, "output", a.initiatorID).Add(float64(req.Output))
223+
a.metrics.TokenUseCount.WithLabelValues(a.provider, a.model, "input", a.initiatorID, a.client).Add(float64(req.Input))
224+
a.metrics.TokenUseCount.WithLabelValues(a.provider, a.model, "output", a.initiatorID, a.client).Add(float64(req.Output))
218225
for k, v := range req.ExtraTokenTypes {
219-
a.metrics.TokenUseCount.WithLabelValues(a.provider, a.model, k, a.initiatorID).Add(float64(v))
226+
a.metrics.TokenUseCount.WithLabelValues(a.provider, a.model, k, a.initiatorID, a.client).Add(float64(v))
220227
}
221228
}
222229
}()

0 commit comments

Comments
 (0)