Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Commitc6631e1

Browse files
authored
feat: exposeaibridged metrics (#20865)
Upgrades `coder/aibridge` to v0.2.0 which includescoder/aibridge#62.Creates a `prometheus.Registerer` with a prefix `coder_aibridged_` andpasses that along to coder/aibridge which actually exposes the metrics.Also includes a side-effect of a change described incoder/aibridge#62 (comment).---------Signed-off-by: Danny Kopping <danny@coder.com>
1 parent6882c43 commitc6631e1

File tree

12 files changed

+295
-124
lines changed

12 files changed

+295
-124
lines changed

‎docs/admin/integrations/prometheus.md‎

Lines changed: 91 additions & 84 deletions
Large diffs are not rendered by default.

‎enterprise/aibridged/aibridged.go‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -55,13 +55,14 @@ func New(ctx context.Context, pool Pooler, rpcDialer Dialer, logger slog.Logger)
5555

5656
ctx,cancel:=context.WithCancel(ctx)
5757
daemon:=&Server{
58-
logger:logger,
59-
clientDialer:rpcDialer,
58+
logger:logger,
59+
clientDialer:rpcDialer,
60+
clientCh:make(chanDRPCClient),
61+
lifecycleCtx:ctx,
62+
cancelFn:cancel,
63+
initConnectionCh:make(chanstruct{}),
64+
6065
requestBridgePool:pool,
61-
clientCh:make(chanDRPCClient),
62-
lifecycleCtx:ctx,
63-
cancelFn:cancel,
64-
initConnectionCh:make(chanstruct{}),
6566
}
6667

6768
daemon.wg.Add(1)

‎enterprise/aibridged/aibridged_integration_test.go‎

Lines changed: 109 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@ import (
99
"testing"
1010
"time"
1111

12+
"github.com/prometheus/client_golang/prometheus"
13+
promtest"github.com/prometheus/client_golang/prometheus/testutil"
1214
"github.com/stretchr/testify/require"
1315

1416
"github.com/coder/aibridge"
@@ -166,7 +168,7 @@ func TestIntegration(t *testing.T) {
166168

167169
logger:=testutil.Logger(t)
168170
providers:= []aibridge.Provider{aibridge.NewOpenAIProvider(aibridge.OpenAIConfig{BaseURL:mockOpenAI.URL})}
169-
pool,err:=aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions,providers,logger)
171+
pool,err:=aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions,providers,nil,logger)
170172
require.NoError(t,err)
171173

172174
// Given: aibridged is started.
@@ -253,3 +255,109 @@ func TestIntegration(t *testing.T) {
253255
// Then: the MCP server was initialized.
254256
require.Contains(t,mcpTokenReceived,authLink.OAuthAccessToken,"mock MCP server not requested")
255257
}
258+
259+
// TestIntegrationWithMetrics validates that Prometheus metrics are correctly incremented
260+
// when requests are processed through aibridged.
261+
funcTestIntegrationWithMetrics(t*testing.T) {
262+
t.Parallel()
263+
264+
ctx:=testutil.Context(t,testutil.WaitLong)
265+
266+
// Create prometheus registry and metrics.
267+
registry:=prometheus.NewRegistry()
268+
metrics:=aibridge.NewMetrics(registry)
269+
270+
// Set up mock OpenAI server.
271+
mockOpenAI:=httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter,r*http.Request) {
272+
w.Header().Set("Content-Type","application/json")
273+
w.WriteHeader(http.StatusOK)
274+
_,_=w.Write([]byte(`{
275+
"id": "chatcmpl-test",
276+
"object": "chat.completion",
277+
"created": 1753343279,
278+
"model": "gpt-4.1",
279+
"choices": [
280+
{
281+
"index": 0,
282+
"message": {
283+
"role": "assistant",
284+
"content": "test response"
285+
},
286+
"finish_reason": "stop"
287+
}
288+
],
289+
"usage": {
290+
"prompt_tokens": 10,
291+
"completion_tokens": 5,
292+
"total_tokens": 15
293+
}
294+
}`))
295+
}))
296+
t.Cleanup(mockOpenAI.Close)
297+
298+
// Database and coderd setup.
299+
db,ps:=dbtestutil.NewDB(t)
300+
client,_,api,firstUser:=coderdenttest.NewWithAPI(t,&coderdenttest.Options{
301+
Options:&coderdtest.Options{
302+
Database:db,
303+
Pubsub:ps,
304+
},
305+
})
306+
307+
userClient,_:=coderdtest.CreateAnotherUser(t,client,firstUser.OrganizationID)
308+
309+
// Create an API token for the user.
310+
apiKey,err:=userClient.CreateToken(ctx,"me", codersdk.CreateTokenRequest{
311+
TokenName:fmt.Sprintf("test-key-%d",time.Now().UnixNano()),
312+
Lifetime:time.Hour,
313+
Scope:codersdk.APIKeyScopeCoderAll,
314+
})
315+
require.NoError(t,err)
316+
317+
// Create aibridge client.
318+
aiBridgeClient,err:=api.CreateInMemoryAIBridgeServer(ctx)
319+
require.NoError(t,err)
320+
321+
logger:=testutil.Logger(t)
322+
providers:= []aibridge.Provider{aibridge.NewOpenAIProvider(aibridge.OpenAIConfig{BaseURL:mockOpenAI.URL})}
323+
324+
// Create pool with metrics.
325+
pool,err:=aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions,providers,metrics,logger)
326+
require.NoError(t,err)
327+
328+
// Given: aibridged is started.
329+
srv,err:=aibridged.New(ctx,pool,func(ctx context.Context) (aibridged.DRPCClient,error) {
330+
returnaiBridgeClient,nil
331+
},logger)
332+
require.NoError(t,err,"create new aibridged")
333+
t.Cleanup(func() {
334+
_=srv.Shutdown(ctx)
335+
})
336+
337+
// When: a request is made to aibridged.
338+
req,err:=http.NewRequestWithContext(ctx,http.MethodPost,"/openai/v1/chat/completions",bytes.NewBufferString(`{
339+
"messages": [
340+
{
341+
"role": "user",
342+
"content": "test message"
343+
}
344+
],
345+
"model": "gpt-4.1"
346+
}`))
347+
require.NoError(t,err,"make request to test server")
348+
req.Header.Add("Authorization","Bearer "+apiKey.Key)
349+
req.Header.Add("Accept","application/json")
350+
351+
// When: aibridged handles the request.
352+
rec:=httptest.NewRecorder()
353+
srv.ServeHTTP(rec,req)
354+
require.Equal(t,http.StatusOK,rec.Code)
355+
356+
// Then: the interceptions metric should increase to 1.
357+
// This is not exhaustively checking the available metrics; just an indicative one to prove
358+
// the plumbing is working.
359+
require.Eventually(t,func()bool {
360+
count:=promtest.ToFloat64(metrics.InterceptionCount)
361+
returncount==1
362+
},testutil.WaitShort,testutil.IntervalFast,"interceptions_total metric should be 1")
363+
}

‎enterprise/aibridged/aibridged_test.go‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -41,8 +41,7 @@ func newTestServer(t *testing.T) (*aibridged.Server, *mock.MockDRPCClient, *mock
4141
pool,
4242
func(ctx context.Context) (aibridged.DRPCClient,error) {
4343
returnclient,nil
44-
},
45-
logger)
44+
},logger)
4645
require.NoError(t,err,"create new aibridged")
4746
t.Cleanup(func() {
4847
srv.Shutdown(context.Background())
@@ -291,7 +290,7 @@ func TestRouting(t *testing.T) {
291290
aibridge.NewOpenAIProvider(aibridge.OpenAIConfig{BaseURL:openaiSrv.URL}),
292291
aibridge.NewAnthropicProvider(aibridge.AnthropicConfig{BaseURL:antSrv.URL},nil),
293292
}
294-
pool,err:=aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions,providers,logger)
293+
pool,err:=aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions,providers,nil,logger)
295294
require.NoError(t,err)
296295
conn:=&mockDRPCConn{}
297296
client.EXPECT().DRPCConn().AnyTimes().Return(conn)

‎enterprise/aibridged/pool.go‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -51,11 +51,13 @@ type CachedBridgePool struct {
5151

5252
singleflight*singleflight.Group[string,*aibridge.RequestBridge]
5353

54+
metrics*aibridge.Metrics
55+
5456
shutDownOnce sync.Once
5557
shuttingDownChchanstruct{}
5658
}
5759

58-
funcNewCachedBridgePool(optionsPoolOptions,providers []aibridge.Provider,logger slog.Logger) (*CachedBridgePool,error) {
60+
funcNewCachedBridgePool(optionsPoolOptions,providers []aibridge.Provider,metrics*aibridge.Metrics,logger slog.Logger) (*CachedBridgePool,error) {
5961
cache,err:=ristretto.NewCache(&ristretto.Config[string,*aibridge.RequestBridge]{
6062
NumCounters:options.MaxItems*10,// Docs suggest setting this 10x number of keys.
6163
MaxCost:options.MaxItems*cacheCost,// Up to n instances.
@@ -88,6 +90,8 @@ func NewCachedBridgePool(options PoolOptions, providers []aibridge.Provider, log
8890

8991
singleflight:&singleflight.Group[string,*aibridge.RequestBridge]{},
9092

93+
metrics:metrics,
94+
9195
shuttingDownCh:make(chanstruct{}),
9296
},nil
9397
}
@@ -154,7 +158,7 @@ func (p *CachedBridgePool) Acquire(ctx context.Context, req Request, clientFn Cl
154158
}
155159
}
156160

157-
bridge,err:=aibridge.NewRequestBridge(ctx,p.providers,p.logger,recorder,mcpServers)
161+
bridge,err:=aibridge.NewRequestBridge(ctx,p.providers,recorder,mcpServers,p.metrics,p.logger)
158162
iferr!=nil {
159163
returnnil,xerrors.Errorf("create new request bridge: %w",err)
160164
}
@@ -167,7 +171,7 @@ func (p *CachedBridgePool) Acquire(ctx context.Context, req Request, clientFn Cl
167171
returninstance,err
168172
}
169173

170-
func (p*CachedBridgePool)Metrics()PoolMetrics {
174+
func (p*CachedBridgePool)CacheMetrics()PoolMetrics {
171175
ifp.cache==nil {
172176
returnnil
173177
}

‎enterprise/aibridged/pool_test.go‎

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ func TestPool(t *testing.T) {
3030
mcpProxy:=mcpmock.NewMockServerProxier(ctrl)
3131

3232
opts:= aibridged.PoolOptions{MaxItems:1,TTL:time.Second}
33-
pool,err:=aibridged.NewCachedBridgePool(opts,nil,logger)
33+
pool,err:=aibridged.NewCachedBridgePool(opts,nil,nil,logger)
3434
require.NoError(t,err)
3535
t.Cleanup(func() {pool.Shutdown(context.Background()) })
3636

@@ -63,11 +63,11 @@ func TestPool(t *testing.T) {
6363
require.NoError(t,err,"acquire pool instance")
6464
require.Same(t,inst,instB)
6565

66-
metrics:=pool.Metrics()
67-
require.EqualValues(t,1,metrics.KeysAdded())
68-
require.EqualValues(t,0,metrics.KeysEvicted())
69-
require.EqualValues(t,1,metrics.Hits())
70-
require.EqualValues(t,1,metrics.Misses())
66+
cacheMetrics:=pool.CacheMetrics()
67+
require.EqualValues(t,1,cacheMetrics.KeysAdded())
68+
require.EqualValues(t,0,cacheMetrics.KeysEvicted())
69+
require.EqualValues(t,1,cacheMetrics.Hits())
70+
require.EqualValues(t,1,cacheMetrics.Misses())
7171

7272
// This will get called again because a new instance will be created.
7373
mcpProxy.EXPECT().Init(gomock.Any()).Times(1).Return(nil)
@@ -81,11 +81,11 @@ func TestPool(t *testing.T) {
8181
require.NoError(t,err,"acquire pool instance")
8282
require.NotSame(t,inst,inst2)
8383

84-
metrics=pool.Metrics()
85-
require.EqualValues(t,2,metrics.KeysAdded())
86-
require.EqualValues(t,1,metrics.KeysEvicted())
87-
require.EqualValues(t,1,metrics.Hits())
88-
require.EqualValues(t,2,metrics.Misses())
84+
cacheMetrics=pool.CacheMetrics()
85+
require.EqualValues(t,2,cacheMetrics.KeysAdded())
86+
require.EqualValues(t,1,cacheMetrics.KeysEvicted())
87+
require.EqualValues(t,1,cacheMetrics.Hits())
88+
require.EqualValues(t,2,cacheMetrics.Misses())
8989

9090
// This will get called again because a new instance will be created.
9191
mcpProxy.EXPECT().Init(gomock.Any()).Times(1).Return(nil)
@@ -99,11 +99,11 @@ func TestPool(t *testing.T) {
9999
require.NoError(t,err,"acquire pool instance 2B")
100100
require.NotSame(t,inst2,inst2B)
101101

102-
metrics=pool.Metrics()
103-
require.EqualValues(t,3,metrics.KeysAdded())
104-
require.EqualValues(t,2,metrics.KeysEvicted())
105-
require.EqualValues(t,1,metrics.Hits())
106-
require.EqualValues(t,3,metrics.Misses())
102+
cacheMetrics=pool.CacheMetrics()
103+
require.EqualValues(t,3,cacheMetrics.KeysAdded())
104+
require.EqualValues(t,2,cacheMetrics.KeysEvicted())
105+
require.EqualValues(t,1,cacheMetrics.Hits())
106+
require.EqualValues(t,3,cacheMetrics.Misses())
107107

108108
// TODO: add test for expiry.
109109
// This requires Go 1.25's [synctest](https://pkg.go.dev/testing/synctest) since the

‎enterprise/aibridged/translator.go‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,16 @@ func (t *recorderTranslation) RecordPromptUsage(ctx context.Context, req *aibrid
5757
}
5858

5959
func (t*recorderTranslation)RecordTokenUsage(ctx context.Context,req*aibridge.TokenUsageRecord)error {
60+
merged:=req.Metadata
61+
ifmerged==nil {
62+
merged= aibridge.Metadata{}
63+
}
64+
65+
// Merge the token usage values into metadata; later we might want to store some of these in their own fields.
66+
fork,v:=rangereq.ExtraTokenTypes {
67+
merged[k]=v
68+
}
69+
6070
_,err:=t.client.RecordTokenUsage(ctx,&proto.RecordTokenUsageRequest{
6171
InterceptionId:req.InterceptionID,
6272
MsgId:req.MsgID,

‎enterprise/cli/aibridged.go‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@ import (
77

88
"golang.org/x/xerrors"
99

10+
"github.com/prometheus/client_golang/prometheus"
11+
1012
"github.com/coder/aibridge"
1113
"github.com/coder/coder/v2/codersdk"
1214
"github.com/coder/coder/v2/enterprise/aibridged"
@@ -31,8 +33,11 @@ func newAIBridgeDaemon(coderAPI *coderd.API) (*aibridged.Server, error) {
3133
},getBedrockConfig(coderAPI.DeploymentValues.AI.BridgeConfig.Bedrock)),
3234
}
3335

36+
reg:=prometheus.WrapRegistererWithPrefix("coder_aibridged_",coderAPI.PrometheusRegistry)
37+
metrics:=aibridge.NewMetrics(reg)
38+
3439
// Create pool for reusable stateful [aibridge.RequestBridge] instances (one per user).
35-
pool,err:=aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions,providers,logger.Named("pool"))// TODO: configurable.
40+
pool,err:=aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions,providers,metrics,logger.Named("pool"))// TODO: configurable size.
3641
iferr!=nil {
3742
returnnil,xerrors.Errorf("create request pool: %w",err)
3843
}

‎go.mod‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -165,9 +165,9 @@ require (
165165
github.com/pkg/diffv0.0.0-20210226163009-20ebb0f2a09e
166166
github.com/pkg/sftpv1.13.7
167167
github.com/prometheus-community/pro-bingv0.7.0
168-
github.com/prometheus/client_golangv1.23.0
168+
github.com/prometheus/client_golangv1.23.2
169169
github.com/prometheus/client_modelv0.6.2
170-
github.com/prometheus/commonv0.65.0
170+
github.com/prometheus/commonv0.66.1
171171
github.com/quasilyte/go-ruleguard/dslv0.3.22
172172
github.com/robfig/cron/v3v3.0.1
173173
github.com/shirou/gopsutil/v4v4.25.5
@@ -476,7 +476,7 @@ require (
476476
github.com/anthropics/anthropic-sdk-gov1.18.0
477477
github.com/brianvoe/gofakeit/v7v7.9.0
478478
github.com/coder/agentapi-sdk-gov0.0.0-20250505131810-560d1d88d225
479-
github.com/coder/aibridgev0.1.7
479+
github.com/coder/aibridgev0.2.0
480480
github.com/coder/aisdk-gov0.0.9
481481
github.com/coder/boundaryv1.0.1-0.20250925154134-55a44f2a7945
482482
github.com/coder/previewv1.0.4

‎go.sum‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -919,8 +919,8 @@ github.com/cncf/xds/go v0.0.0-20250501225837-2ac532fd4443 h1:aQ3y1lwWyqYPiWZThqv
919919
github.com/cncf/xds/gov0.0.0-20250501225837-2ac532fd4443/go.mod h1:W+zGtBO5Y1IgJhy4+A9GOqVhqLpfZi+vwmdNXUehLA8=
920920
github.com/coder/agentapi-sdk-gov0.0.0-20250505131810-560d1d88d225 h1:tRIViZ5JRmzdOEo5wUWngaGEFBG8OaE1o2GIHN5ujJ8=
921921
github.com/coder/agentapi-sdk-gov0.0.0-20250505131810-560d1d88d225/go.mod h1:rNLVpYgEVeu1Zk29K64z6Od8RBP9DwqCu9OfCzh8MR4=
922-
github.com/coder/aibridgev0.1.7 h1:GTAM8nHawXMeb/pxAIwvzr76dyVGu9hw9qV6Gvpc7nw=
923-
github.com/coder/aibridgev0.1.7/go.mod h1:7GhrLbzf6uM3sCA7OPaDzvq9QNrCjNuzMy+WgipYwfQ=
922+
github.com/coder/aibridgev0.2.0 h1:kAWhHD6fsmDLH1WxIwXPu9Ineijj+lVniko45C003Vo=
923+
github.com/coder/aibridgev0.2.0/go.mod h1:2T0RSnIX1WTqFajzXsaNsoNe6mmNsNeCTxiHBWEsFnE=
924924
github.com/coder/aisdk-gov0.0.9 h1:Vzo/k2qwVGLTR10ESDeP2Ecek1SdPfZlEjtTfMveiVo=
925925
github.com/coder/aisdk-gov0.0.9/go.mod h1:KF6/Vkono0FJJOtWtveh5j7yfNrSctVTpwgweYWSp5M=
926926
github.com/coder/boundaryv1.0.1-0.20250925154134-55a44f2a7945 h1:hDUf02kTX8EGR3+5B+v5KdYvORs4YNfDPci0zCs+pC0=
@@ -1718,15 +1718,15 @@ github.com/power-devops/perfstat v0.0.0-20240221224432-82ca36839d55 h1:o4JXh1EVt
17181718
github.com/power-devops/perfstatv0.0.0-20240221224432-82ca36839d55/go.mod h1:OmDBASR4679mdNQnz2pUhc2G8CO2JrUAVFDRBDP/hJE=
17191719
github.com/prometheus-community/pro-bingv0.7.0 h1:KFYFbxC2f2Fp6c+TyxbCOEarf7rbnzr9Gw8eIb0RfZA=
17201720
github.com/prometheus-community/pro-bingv0.7.0/go.mod h1:Moob9dvlY50Bfq6i88xIwfyw7xLFHH69LUgx9n5zqCE=
1721-
github.com/prometheus/client_golangv1.23.0 h1:ust4zpdl9r4trLY/gSjlm07PuiBq2ynaXXlptpfy8Uc=
1722-
github.com/prometheus/client_golangv1.23.0/go.mod h1:i/o0R9ByOnHX0McrTMTyhYvKE4haaf2mW08I+jGAjEE=
1721+
github.com/prometheus/client_golangv1.23.2 h1:Je96obch5RDVy3FDMndoUsjAhG5Edi49h0RJWRi/o0o=
1722+
github.com/prometheus/client_golangv1.23.2/go.mod h1:Tb1a6LWHB3/SPIzCoaDXI4I8UHKeFTEQ1YCr+0Gyqmg=
17231723
github.com/prometheus/client_modelv0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
17241724
github.com/prometheus/client_modelv0.2.0/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
17251725
github.com/prometheus/client_modelv0.3.0/go.mod h1:LDGWKZIo7rky3hgvBe+caln+Dr3dPggB5dvjtD7w9+w=
17261726
github.com/prometheus/client_modelv0.6.2 h1:oBsgwpGs7iVziMvrGhE53c/GrLUsZdHnqNwqPLxwZyk=
17271727
github.com/prometheus/client_modelv0.6.2/go.mod h1:y3m2F6Gdpfy6Ut/GBsUqTWZqCUvMVzSfMLjcu6wAwpE=
1728-
github.com/prometheus/commonv0.65.0 h1:QDwzd+G1twt//Kwj/Ww6E9FQq1iVMmODnILtW1t2VzE=
1729-
github.com/prometheus/commonv0.65.0/go.mod h1:0gZns+BLRQ3V6NdaerOhMbwwRbNh9hkGINtQAsP5GS8=
1728+
github.com/prometheus/commonv0.66.1 h1:h5E0h5/Y8niHc5DlaLlWLArTQI7tMrsfQjHV+d9ZoGs=
1729+
github.com/prometheus/commonv0.66.1/go.mod h1:gcaUsgf3KfRSwHY4dIMXLPV0K/Wg1oZ8+SbZk/HH/dA=
17301730
github.com/prometheus/procfsv0.16.1 h1:hZ15bTNuirocR6u0JZ6BAHHmwS1p8B4P6MRqxtzMyRg=
17311731
github.com/prometheus/procfsv0.16.1/go.mod h1:teAbpZRB1iIAJYREa1LsoWUXykVXA1KlTmWl8x/U+Is=
17321732
github.com/puzpuzpuz/xsync/v3v3.5.1 h1:GJYJZwO6IdxN/IKbneznS6yPkVC+c3zyY/j19c++5Fg=

0 commit comments

Comments
 (0)
[8]ページ先頭
©2009-2025 Movatter.jp