Skip to content

Commit 9e9491b

Browse files
dannykoppingdannykopping
authored
feat: expose prometheus metrics (#62)
* WIP Signed-off-by: Danny Kopping <danny@coder.com> * chore: fix tests Signed-off-by: Danny Kopping <danny@coder.com> * feat: token usage tracking Signed-off-by: Danny Kopping <danny@coder.com> * feat: track tool usage Signed-off-by: Danny Kopping <danny@coder.com> * chore: passthrough Signed-off-by: Danny Kopping <danny@coder.com> * chore: inflight Signed-off-by: Danny Kopping <danny@coder.com> * chore: cardinality estimates Signed-off-by: Danny Kopping <danny@coder.com> * chore: tests Signed-off-by: Danny Kopping <danny@coder.com> * chore: pr comments Signed-off-by: Danny Kopping <danny@coder.com> * chore: make fmt Signed-off-by: Danny Kopping <danny@coder.com> * chore: drop tool initiator partitioning Signed-off-by: Danny Kopping <danny@coder.com> --------- Signed-off-by: Danny Kopping <danny@coder.com>
1 parent ba198a7 commit 9e9491b

14 files changed

+566
-46
lines changed

‎api.go‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,8 +25,12 @@ type TokenUsageRecord struct {
2525
InterceptionID string
2626
MsgID string
2727
Input, Output int64
28-
Metadata Metadata
29-
CreatedAt time.Time
28+
// ExtraTokenTypes holds token types which *may* exist over and above input/output.
29+
// These should ultimately get merged into [Metadata], but it's useful to keep these
30+
// with their actual type (int64) since [Metadata] is a map[string]any.
31+
ExtraTokenTypes map[string]int64
32+
Metadata Metadata
33+
CreatedAt time.Time
3034
}
3135

3236
type PromptUsageRecord struct {

‎bridge.go‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -47,20 +47,20 @@ var _ http.Handler = &RequestBridge{}
4747
// A [Recorder] is also required to record prompt, tool, and token use.
4848
//
4949
// mcpProxy will be closed when the [RequestBridge] is closed.
50-
func NewRequestBridge(ctx context.Context, providers []Provider, logger slog.Logger, recorder Recorder, mcpProxy mcp.ServerProxier) (*RequestBridge, error) {
50+
func NewRequestBridge(ctx context.Context, providers []Provider, recorder Recorder, mcpProxy mcp.ServerProxier, metrics *Metrics, logger slog.Logger) (*RequestBridge, error) {
5151
mux := http.NewServeMux()
5252

5353
for _, provider := range providers {
5454
// Add the known provider-specific routes which are bridged (i.e. intercepted and augmented).
5555
for _, path := range provider.BridgedRoutes() {
56-
mux.HandleFunc(path, newInterceptionProcessor(provider, logger, recorder, mcpProxy))
56+
mux.HandleFunc(path, newInterceptionProcessor(provider, logger, recorder, mcpProxy, metrics))
5757
}
5858

5959
// Any requests which passthrough to this will be reverse-proxied to the upstream.
6060
//
61-
// We have to whitelist the known-safe routes because an API key with elevant privileges (i.e. admin) might be
61+
// We have to whitelist the known-safe routes because an API key with elevated privileges (i.e. admin) might be
6262
// configured, so we should just reverse-proxy known-safe routes.
63-
ftr := newPassthroughRouter(provider, logger.Named(fmt.Sprintf("passthrough.%s", provider.Name())))
63+
ftr := newPassthroughRouter(provider, logger.Named(fmt.Sprintf("passthrough.%s", provider.Name())), metrics)
6464
for _, path := range provider.PassthroughRoutes() {
6565
prefix := fmt.Sprintf("/%s", provider.Name())
6666
route := fmt.Sprintf("%s%s", prefix, path)

‎bridge_integration_test.go‎

Lines changed: 18 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -133,7 +133,7 @@ func TestAnthropicMessages(t *testing.T) {
133133
recorderClient := &mockRecorderClient{}
134134

135135
logger := slogtest.Make(t, &slogtest.Options{}).Leveled(slog.LevelDebug)
136-
b, err := aibridge.NewRequestBridge(ctx, []aibridge.Provider{aibridge.NewAnthropicProvider(anthropicCfg(srv.URL, apiKey), nil)}, logger, recorderClient, mcp.NewServerProxyManager(nil))
136+
b, err := aibridge.NewRequestBridge(ctx, []aibridge.Provider{aibridge.NewAnthropicProvider(anthropicCfg(srv.URL, apiKey), nil)}, recorderClient, mcp.NewServerProxyManager(nil), nil, logger)
137137
require.NoError(t, err)
138138

139139
mockSrv := httptest.NewUnstartedServer(b)
@@ -214,7 +214,7 @@ func TestAWSBedrockIntegration(t *testing.T) {
214214
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
215215
b, err := aibridge.NewRequestBridge(ctx, []aibridge.Provider{
216216
aibridge.NewAnthropicProvider(anthropicCfg("http://unused", apiKey), bedrockCfg),
217-
}, logger, recorderClient, mcp.NewServerProxyManager(nil))
217+
}, recorderClient, mcp.NewServerProxyManager(nil), nil, logger)
218218
require.NoError(t, err)
219219

220220
mockSrv := httptest.NewUnstartedServer(b)
@@ -312,7 +312,7 @@ func TestAWSBedrockIntegration(t *testing.T) {
312312
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
313313
b, err := aibridge.NewRequestBridge(
314314
ctx, []aibridge.Provider{aibridge.NewAnthropicProvider(anthropicCfg(srv.URL, apiKey), bedrockCfg)},
315-
logger, recorderClient, mcp.NewServerProxyManager(nil))
315+
recorderClient, mcp.NewServerProxyManager(nil), nil, logger)
316316
require.NoError(t, err)
317317

318318
mockBridgeSrv := httptest.NewUnstartedServer(b)
@@ -399,7 +399,7 @@ func TestOpenAIChatCompletions(t *testing.T) {
399399
recorderClient := &mockRecorderClient{}
400400

401401
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: false}).Leveled(slog.LevelDebug)
402-
b, err := aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{aibridge.NewOpenAIProvider(openaiCfg(srv.URL, apiKey))}, logger, recorderClient, mcp.NewServerProxyManager(nil))
402+
b, err := aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{aibridge.NewOpenAIProvider(openaiCfg(srv.URL, apiKey))}, recorderClient, mcp.NewServerProxyManager(nil), nil, logger)
403403
require.NoError(t, err)
404404

405405
mockSrv := httptest.NewUnstartedServer(b)
@@ -466,7 +466,7 @@ func TestSimple(t *testing.T) {
466466
fixture: antSimple,
467467
configureFunc: func(addr string, client aibridge.Recorder) (*aibridge.RequestBridge, error) {
468468
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: false}).Leveled(slog.LevelDebug)
469-
return aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{aibridge.NewAnthropicProvider(anthropicCfg(addr, apiKey), nil)}, logger, client, mcp.NewServerProxyManager(nil))
469+
return aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{aibridge.NewAnthropicProvider(anthropicCfg(addr, apiKey), nil)}, client, mcp.NewServerProxyManager(nil), nil, logger)
470470
},
471471
getResponseIDFunc: func(streaming bool, resp *http.Response) (string, error) {
472472
if streaming {
@@ -504,7 +504,7 @@ func TestSimple(t *testing.T) {
504504
fixture: oaiSimple,
505505
configureFunc: func(addr string, client aibridge.Recorder) (*aibridge.RequestBridge, error) {
506506
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: false}).Leveled(slog.LevelDebug)
507-
return aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{aibridge.NewOpenAIProvider(openaiCfg(addr, apiKey))}, logger, client, mcp.NewServerProxyManager(nil))
507+
return aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{aibridge.NewOpenAIProvider(openaiCfg(addr, apiKey))}, client, mcp.NewServerProxyManager(nil), nil, logger)
508508
},
509509
getResponseIDFunc: func(streaming bool, resp *http.Response) (string, error) {
510510
if streaming {
@@ -645,7 +645,7 @@ func TestFallthrough(t *testing.T) {
645645
configureFunc: func(addr string, client aibridge.Recorder) (aibridge.Provider, *aibridge.RequestBridge) {
646646
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: false}).Leveled(slog.LevelDebug)
647647
provider := aibridge.NewAnthropicProvider(anthropicCfg(addr, apiKey), nil)
648-
bridge, err := aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{provider}, logger, client, mcp.NewServerProxyManager(nil))
648+
bridge, err := aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{provider}, client, mcp.NewServerProxyManager(nil), nil, logger)
649649
require.NoError(t, err)
650650
return provider, bridge
651651
},
@@ -656,7 +656,7 @@ func TestFallthrough(t *testing.T) {
656656
configureFunc: func(addr string, client aibridge.Recorder) (aibridge.Provider, *aibridge.RequestBridge) {
657657
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: false}).Leveled(slog.LevelDebug)
658658
provider := aibridge.NewOpenAIProvider(openaiCfg(addr, apiKey))
659-
bridge, err := aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{provider}, logger, client, mcp.NewServerProxyManager(nil))
659+
bridge, err := aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{provider}, client, mcp.NewServerProxyManager(nil), nil, logger)
660660
require.NoError(t, err)
661661
return provider, bridge
662662
},
@@ -762,7 +762,7 @@ func TestAnthropicInjectedTools(t *testing.T) {
762762

763763
configureFn := func(addr string, client aibridge.Recorder, srvProxyMgr *mcp.ServerProxyManager) (*aibridge.RequestBridge, error) {
764764
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: false}).Leveled(slog.LevelDebug)
765-
return aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{aibridge.NewAnthropicProvider(anthropicCfg(addr, apiKey), nil)}, logger, client, srvProxyMgr)
765+
return aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{aibridge.NewAnthropicProvider(anthropicCfg(addr, apiKey), nil)}, client, srvProxyMgr, nil, logger)
766766
}
767767

768768
// Build the requirements & make the assertions which are common to all providers.
@@ -843,7 +843,7 @@ func TestOpenAIInjectedTools(t *testing.T) {
843843

844844
configureFn := func(addr string, client aibridge.Recorder, srvProxyMgr *mcp.ServerProxyManager) (*aibridge.RequestBridge, error) {
845845
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: false}).Leveled(slog.LevelDebug)
846-
return aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{aibridge.NewOpenAIProvider(openaiCfg(addr, apiKey))}, logger, client, srvProxyMgr)
846+
return aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{aibridge.NewOpenAIProvider(openaiCfg(addr, apiKey))}, client, srvProxyMgr, nil, logger)
847847
}
848848

849849
// Build the requirements & make the assertions which are common to all providers.
@@ -1029,7 +1029,7 @@ func TestErrorHandling(t *testing.T) {
10291029
createRequestFunc: createAnthropicMessagesReq,
10301030
configureFunc: func(addr string, client aibridge.Recorder, srvProxyMgr *mcp.ServerProxyManager) (*aibridge.RequestBridge, error) {
10311031
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: false}).Leveled(slog.LevelDebug)
1032-
return aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{aibridge.NewAnthropicProvider(anthropicCfg(addr, apiKey), nil)}, logger, client, srvProxyMgr)
1032+
return aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{aibridge.NewAnthropicProvider(anthropicCfg(addr, apiKey), nil)}, client, srvProxyMgr, nil, logger)
10331033
},
10341034
responseHandlerFn: func(resp *http.Response) {
10351035
require.Equal(t, http.StatusBadRequest, resp.StatusCode)
@@ -1046,7 +1046,7 @@ func TestErrorHandling(t *testing.T) {
10461046
createRequestFunc: createOpenAIChatCompletionsReq,
10471047
configureFunc: func(addr string, client aibridge.Recorder, srvProxyMgr *mcp.ServerProxyManager) (*aibridge.RequestBridge, error) {
10481048
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: false}).Leveled(slog.LevelDebug)
1049-
return aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{aibridge.NewOpenAIProvider(openaiCfg(addr, apiKey))}, logger, client, srvProxyMgr)
1049+
return aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{aibridge.NewOpenAIProvider(openaiCfg(addr, apiKey))}, client, srvProxyMgr, nil, logger)
10501050
},
10511051
responseHandlerFn: func(resp *http.Response) {
10521052
require.Equal(t, http.StatusBadRequest, resp.StatusCode)
@@ -1134,7 +1134,7 @@ func TestErrorHandling(t *testing.T) {
11341134
createRequestFunc: createAnthropicMessagesReq,
11351135
configureFunc: func(addr string, client aibridge.Recorder, srvProxyMgr *mcp.ServerProxyManager) (*aibridge.RequestBridge, error) {
11361136
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: false}).Leveled(slog.LevelDebug)
1137-
return aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{aibridge.NewAnthropicProvider(anthropicCfg(addr, apiKey), nil)}, logger, client, srvProxyMgr)
1137+
return aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{aibridge.NewAnthropicProvider(anthropicCfg(addr, apiKey), nil)}, client, srvProxyMgr, nil, logger)
11381138
},
11391139
responseHandlerFn: func(resp *http.Response) {
11401140
// Server responds first with 200 OK then starts streaming.
@@ -1152,7 +1152,7 @@ func TestErrorHandling(t *testing.T) {
11521152
createRequestFunc: createOpenAIChatCompletionsReq,
11531153
configureFunc: func(addr string, client aibridge.Recorder, srvProxyMgr *mcp.ServerProxyManager) (*aibridge.RequestBridge, error) {
11541154
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: false}).Leveled(slog.LevelDebug)
1155-
return aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{aibridge.NewOpenAIProvider(openaiCfg(addr, apiKey))}, logger, client, srvProxyMgr)
1155+
return aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{aibridge.NewOpenAIProvider(openaiCfg(addr, apiKey))}, client, srvProxyMgr, nil, logger)
11561156
},
11571157
responseHandlerFn: func(resp *http.Response) {
11581158
// Server responds first with 200 OK then starts streaming.
@@ -1238,15 +1238,15 @@ func TestStableRequestEncoding(t *testing.T) {
12381238
fixture: antSimple,
12391239
createRequestFunc: createAnthropicMessagesReq,
12401240
configureFunc: func(addr string, client aibridge.Recorder, srvProxyMgr *mcp.ServerProxyManager) (*aibridge.RequestBridge, error) {
1241-
return aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{aibridge.NewAnthropicProvider(anthropicCfg(addr, apiKey), nil)}, logger, client, srvProxyMgr)
1241+
return aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{aibridge.NewAnthropicProvider(anthropicCfg(addr, apiKey), nil)}, client, srvProxyMgr, nil, logger)
12421242
},
12431243
},
12441244
{
12451245
name: aibridge.ProviderOpenAI,
12461246
fixture: oaiSimple,
12471247
createRequestFunc: createOpenAIChatCompletionsReq,
12481248
configureFunc: func(addr string, client aibridge.Recorder, srvProxyMgr *mcp.ServerProxyManager) (*aibridge.RequestBridge, error) {
1249-
return aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{aibridge.NewOpenAIProvider(openaiCfg(addr, apiKey))}, logger, client, srvProxyMgr)
1249+
return aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{aibridge.NewOpenAIProvider(openaiCfg(addr, apiKey))}, client, srvProxyMgr, nil, logger)
12501250
},
12511251
},
12521252
}
@@ -1352,7 +1352,7 @@ func TestEnvironmentDoNotLeak(t *testing.T) {
13521352
fixture: antSimple,
13531353
configureFunc: func(addr string, client aibridge.Recorder) (*aibridge.RequestBridge, error) {
13541354
logger := slogtest.Make(t, &slogtest.Options{}).Leveled(slog.LevelDebug)
1355-
return aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{aibridge.NewAnthropicProvider(anthropicCfg(addr, apiKey), nil)}, logger, client, mcp.NewServerProxyManager(nil))
1355+
return aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{aibridge.NewAnthropicProvider(anthropicCfg(addr, apiKey), nil)}, client, mcp.NewServerProxyManager(nil), nil, logger)
13561356
},
13571357
createRequest: createAnthropicMessagesReq,
13581358
envVars: map[string]string{
@@ -1365,7 +1365,7 @@ func TestEnvironmentDoNotLeak(t *testing.T) {
13651365
fixture: oaiSimple,
13661366
configureFunc: func(addr string, client aibridge.Recorder) (*aibridge.RequestBridge, error) {
13671367
logger := slogtest.Make(t, &slogtest.Options{}).Leveled(slog.LevelDebug)
1368-
return aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{aibridge.NewOpenAIProvider(openaiCfg(addr, apiKey))}, logger, client, mcp.NewServerProxyManager(nil))
1368+
return aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{aibridge.NewOpenAIProvider(openaiCfg(addr, apiKey))}, client, mcp.NewServerProxyManager(nil), nil, logger)
13691369
},
13701370
createRequest: createOpenAIChatCompletionsReq,
13711371
envVars: map[string]string{

‎go.mod‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,8 @@ require (
88
github.com/google/uuid v1.6.0
99
github.com/hashicorp/go-multierror v1.1.1
1010
github.com/mark3labs/mcp-go v0.38.0
11-
github.com/stretchr/testify v1.10.0
11+
github.com/prometheus/client_golang v1.23.2
12+
github.com/stretchr/testify v1.11.1
1213
github.com/tidwall/gjson v1.18.0
1314
github.com/tidwall/sjson v1.2.5
1415
go.uber.org/goleak v1.3.0
@@ -40,18 +41,25 @@ require (
4041
github.com/aws/smithy-go v1.20.3 // indirect
4142
github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
4243
github.com/bahlo/generic-list-go v0.2.0 // indirect
44+
github.com/beorn7/perks v1.0.1 // indirect
4345
github.com/buger/jsonparser v1.1.1 // indirect
46+
github.com/cespare/xxhash/v2 v2.3.0 // indirect
4447
github.com/charmbracelet/lipgloss v0.7.1 // indirect
4548
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
4649
github.com/hashicorp/errwrap v1.0.0 // indirect
4750
github.com/invopop/jsonschema v0.13.0 // indirect
51+
github.com/kylelemons/godebug v1.1.0 // indirect
4852
github.com/lucasb-eyer/go-colorful v1.2.0 // indirect
4953
github.com/mailru/easyjson v0.7.7 // indirect
5054
github.com/mattn/go-isatty v0.0.20 // indirect
5155
github.com/mattn/go-runewidth v0.0.15 // indirect
5256
github.com/muesli/reflow v0.3.0 // indirect
5357
github.com/muesli/termenv v0.15.2 // indirect
58+
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
5459
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
60+
github.com/prometheus/client_model v0.6.2 // indirect
61+
github.com/prometheus/common v0.66.1 // indirect
62+
github.com/prometheus/procfs v0.16.1 // indirect
5563
github.com/rivo/uniseg v0.4.4 // indirect
5664
github.com/rogpeppe/go-internal v1.13.1 // indirect
5765
github.com/spf13/cast v1.7.1 // indirect
@@ -61,11 +69,11 @@ require (
6169
github.com/yosida95/uritemplate/v3 v3.0.2 // indirect
6270
go.opentelemetry.io/otel v1.33.0 // indirect
6371
go.opentelemetry.io/otel/trace v1.33.0 // indirect
72+
go.yaml.in/yaml/v2 v2.4.2 // indirect
6473
golang.org/x/sys v0.35.0 // indirect
6574
golang.org/x/term v0.34.0 // indirect
66-
golang.org/x/text v0.28.0 // indirect
6775
golang.org/x/xerrors v0.0.0-20240716161551-93cc26a95ae9 // indirect
6876
google.golang.org/genproto/googleapis/api v0.0.0-20240318140521-94a12d6c2237 // indirect
69-
google.golang.org/protobuf v1.36.3 // indirect
77+
google.golang.org/protobuf v1.36.8 // indirect
7078
gopkg.in/yaml.v3 v3.0.1 // indirect
7179
)

0 commit comments

Comments
 (0)