Skip to content

Commit 597f3c0

Browse files
pawbanapawbana
committed
ServerProxyManager init tracing
1 parent 5fe4137 commit 597f3c0

File tree

6 files changed

+97
-44
lines changed

6 files changed

+97
-44
lines changed

‎aibtrace/aibtrace.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ const (
1818
Model = "model"
1919
Streaming = "streaming"
2020
IsBedrock = "aws_bedrock"
21+
MCPProxyName = "mcp_proxy_name"
2122
MCPToolName = "mcp_tool_name"
2223
PassthroughURL = "passthrough_url"
2324
PassthroughMethod = "passthrough_method"

‎bridge_integration_test.go‎

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -136,7 +136,7 @@ func TestAnthropicMessages(t *testing.T) {
136136

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

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

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

321321
mockBridgeSrv := httptest.NewUnstartedServer(b)
@@ -403,7 +403,7 @@ func TestOpenAIChatCompletions(t *testing.T) {
403403

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

409409
mockSrv := httptest.NewUnstartedServer(b)
@@ -471,7 +471,7 @@ func TestSimple(t *testing.T) {
471471
configureFunc: func(addr string, client aibridge.Recorder) (*aibridge.RequestBridge, error) {
472472
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: false}).Leveled(slog.LevelDebug)
473473
provider := []aibridge.Provider{aibridge.NewAnthropicProvider(anthropicCfg(addr, apiKey), nil)}
474-
return aibridge.NewRequestBridge(t.Context(), provider, client, mcp.NewServerProxyManager(nil), nil, defaultTracer, logger)
474+
return aibridge.NewRequestBridge(t.Context(), provider, client, mcp.NewServerProxyManager(nil, defaultTracer), nil, defaultTracer, logger)
475475
},
476476
getResponseIDFunc: func(streaming bool, resp *http.Response) (string, error) {
477477
if streaming {
@@ -510,7 +510,7 @@ func TestSimple(t *testing.T) {
510510
configureFunc: func(addr string, client aibridge.Recorder) (*aibridge.RequestBridge, error) {
511511
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: false}).Leveled(slog.LevelDebug)
512512
providers := []aibridge.Provider{aibridge.NewOpenAIProvider(openaiCfg(addr, apiKey))}
513-
return aibridge.NewRequestBridge(t.Context(), providers, client, mcp.NewServerProxyManager(nil), nil, defaultTracer, logger)
513+
return aibridge.NewRequestBridge(t.Context(), providers, client, mcp.NewServerProxyManager(nil, defaultTracer), nil, defaultTracer, logger)
514514
},
515515
getResponseIDFunc: func(streaming bool, resp *http.Response) (string, error) {
516516
if streaming {
@@ -642,7 +642,7 @@ func TestFallthrough(t *testing.T) {
642642
configureFunc: func(addr string, client aibridge.Recorder) (aibridge.Provider, *aibridge.RequestBridge) {
643643
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: false}).Leveled(slog.LevelDebug)
644644
provider := aibridge.NewAnthropicProvider(anthropicCfg(addr, apiKey), nil)
645-
bridge, err := aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{provider}, client, mcp.NewServerProxyManager(nil), nil, defaultTracer, logger)
645+
bridge, err := aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{provider}, client, mcp.NewServerProxyManager(nil, defaultTracer), nil, defaultTracer, logger)
646646
require.NoError(t, err)
647647
return provider, bridge
648648
},
@@ -653,7 +653,7 @@ func TestFallthrough(t *testing.T) {
653653
configureFunc: func(addr string, client aibridge.Recorder) (aibridge.Provider, *aibridge.RequestBridge) {
654654
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: false}).Leveled(slog.LevelDebug)
655655
provider := aibridge.NewOpenAIProvider(openaiCfg(addr, apiKey))
656-
bridge, err := aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{provider}, client, mcp.NewServerProxyManager(nil), nil, defaultTracer, logger)
656+
bridge, err := aibridge.NewRequestBridge(t.Context(), []aibridge.Provider{provider}, client, mcp.NewServerProxyManager(nil, defaultTracer), nil, defaultTracer, logger)
657657
require.NoError(t, err)
658658
return provider, bridge
659659
},
@@ -980,7 +980,7 @@ func setupInjectedToolTest(t *testing.T, fixture []byte, streaming bool, configu
980980
tools := setupMCPServerProxiesForTest(t)
981981

982982
// Configure the bridge with injected tools.
983-
mcpMgr := mcp.NewServerProxyManager(tools)
983+
mcpMgr := mcp.NewServerProxyManager(tools, defaultTracer)
984984
require.NoError(t, mcpMgr.Init(ctx))
985985
b, err := configureFn(mockSrv.URL, recorderClient, mcpMgr)
986986
require.NoError(t, err)
@@ -1096,7 +1096,7 @@ func TestErrorHandling(t *testing.T) {
10961096

10971097
recorderClient := &mockRecorderClient{}
10981098

1099-
b, err := tc.configureFunc(mockSrv.URL, recorderClient, mcp.NewServerProxyManager(nil))
1099+
b, err := tc.configureFunc(mockSrv.URL, recorderClient, mcp.NewServerProxyManager(nil, defaultTracer))
11001100
require.NoError(t, err)
11011101

11021102
// Invoke request to mocked API via aibridge.
@@ -1198,7 +1198,7 @@ func TestErrorHandling(t *testing.T) {
11981198

11991199
recorderClient := &mockRecorderClient{}
12001200

1201-
b, err := tc.configureFunc(mockSrv.URL, recorderClient, mcp.NewServerProxyManager(nil))
1201+
b, err := tc.configureFunc(mockSrv.URL, recorderClient, mcp.NewServerProxyManager(nil, defaultTracer))
12021202
require.NoError(t, err)
12031203

12041204
// Invoke request to mocked API via aibridge.
@@ -1267,7 +1267,7 @@ func TestStableRequestEncoding(t *testing.T) {
12671267
tools := setupMCPServerProxiesForTest(t)
12681268

12691269
// Configure the bridge with injected tools.
1270-
mcpMgr := mcp.NewServerProxyManager(tools)
1270+
mcpMgr := mcp.NewServerProxyManager(tools, defaultTracer)
12711271
require.NoError(t, mcpMgr.Init(ctx))
12721272

12731273
arc := txtar.Parse(tc.fixture)
@@ -1358,7 +1358,7 @@ func TestEnvironmentDoNotLeak(t *testing.T) {
13581358
configureFunc: func(addr string, client aibridge.Recorder) (*aibridge.RequestBridge, error) {
13591359
logger := slogtest.Make(t, &slogtest.Options{}).Leveled(slog.LevelDebug)
13601360
providers := []aibridge.Provider{aibridge.NewAnthropicProvider(anthropicCfg(addr, apiKey), nil)}
1361-
return aibridge.NewRequestBridge(t.Context(), providers, client, mcp.NewServerProxyManager(nil), nil, defaultTracer, logger)
1361+
return aibridge.NewRequestBridge(t.Context(), providers, client, mcp.NewServerProxyManager(nil, defaultTracer), nil, defaultTracer, logger)
13621362
},
13631363
createRequest: createAnthropicMessagesReq,
13641364
envVars: map[string]string{
@@ -1372,7 +1372,7 @@ func TestEnvironmentDoNotLeak(t *testing.T) {
13721372
configureFunc: func(addr string, client aibridge.Recorder) (*aibridge.RequestBridge, error) {
13731373
logger := slogtest.Make(t, &slogtest.Options{}).Leveled(slog.LevelDebug)
13741374
providers := []aibridge.Provider{aibridge.NewOpenAIProvider(openaiCfg(addr, apiKey))}
1375-
return aibridge.NewRequestBridge(t.Context(), providers, client, mcp.NewServerProxyManager(nil), nil, defaultTracer, logger)
1375+
return aibridge.NewRequestBridge(t.Context(), providers, client, mcp.NewServerProxyManager(nil, defaultTracer), nil, defaultTracer, logger)
13761376
},
13771377
createRequest: createOpenAIChatCompletionsReq,
13781378
envVars: map[string]string{

‎mcp/mcp_test.go‎

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

1414
"cdr.dev/slog"
1515
"cdr.dev/slog/sloggers/slogtest"
16+
"go.opentelemetry.io/otel"
1617
"go.uber.org/goleak"
1718

1819
"github.com/coder/aibridge/mcp"
@@ -324,7 +325,7 @@ func TestToolInjectionOrder(t *testing.T) {
324325
mgr := mcp.NewServerProxyManager(map[string]mcp.ServerProxier{
325326
"coder": proxy,
326327
"shmoder": proxy2,
327-
})
328+
}, otel.GetTracerProvider().Tracer("test"))
328329
require.NoError(t, mgr.Init(ctx))
329330

330331
// Then: the tools from both servers should be collectively sorted stably.

‎mcp/server_proxy_manager.go‎

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,11 @@ import (
77
"strings"
88
"sync"
99

10+
"github.com/coder/aibridge/aibtrace"
1011
"github.com/coder/aibridge/utils"
1112
"github.com/mark3labs/mcp-go/mcp"
13+
"go.opentelemetry.io/otel/attribute"
14+
"go.opentelemetry.io/otel/trace"
1215
)
1316

1417
var _ ServerProxier = &ServerProxyManager{}
@@ -18,14 +21,18 @@ var _ ServerProxier = &ServerProxyManager{}
1821
// for the purpose of injection into bridged requests and invocation.
1922
type ServerProxyManager struct {
2023
proxiers map[string]ServerProxier
24+
tracer trace.Tracer
2125

2226
// Protects access to the tools map.
2327
toolsMu sync.RWMutex
2428
tools map[string]*Tool
2529
}
2630

27-
func NewServerProxyManager(proxiers map[string]ServerProxier) *ServerProxyManager {
28-
return &ServerProxyManager{proxiers: proxiers}
31+
func NewServerProxyManager(proxiers map[string]ServerProxier, tracer trace.Tracer) *ServerProxyManager {
32+
return &ServerProxyManager{
33+
proxiers: proxiers,
34+
tracer: tracer,
35+
}
2936
}
3037

3138
func (s *ServerProxyManager) addTools(tools []*Tool) {
@@ -42,10 +49,16 @@ func (s *ServerProxyManager) addTools(tools []*Tool) {
4249
}
4350

4451
// Init concurrently initializes all of its [ServerProxier]s.
45-
func (s *ServerProxyManager) Init(ctx context.Context) error {
52+
func (s *ServerProxyManager) Init(ctx context.Context) (outErr error) {
53+
ctx, span := s.tracer.Start(ctx, "ServerProxyManager.Init")
54+
defer aibtrace.EndSpanErr(span, &outErr)
55+
4656
cg := utils.NewConcurrentGroup()
47-
for _, proxy := range s.proxiers {
57+
for name, proxy := range s.proxiers {
4858
cg.Go(func() error {
59+
ctx, span := s.tracer.Start(ctx, "ServerProxyManager.Init.Proxy", trace.WithAttributes(attribute.String(aibtrace.MCPProxyName, name)))
60+
defer aibtrace.EndSpanErr(span, &outErr)
61+
4962
return proxy.Init(ctx)
5063
})
5164
}

‎metrics_integration_test.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -238,7 +238,7 @@ func TestMetrics_InjectedToolUseCount(t *testing.T) {
238238

239239
// Setup mocked MCP server & tools.
240240
tools := setupMCPServerProxiesForTest(t)
241-
mcpMgr := mcp.NewServerProxyManager(tools)
241+
mcpMgr := mcp.NewServerProxyManager(tools, defaultTracer)
242242
require.NoError(t, mcpMgr.Init(ctx))
243243

244244
bridge, err := aibridge.NewRequestBridge(ctx, []aibridge.Provider{provider}, recorder, mcpMgr, metrics, defaultTracer, logger)
@@ -283,7 +283,7 @@ func newTestSrv(t *testing.T, ctx context.Context, provider aibridge.Provider, m
283283
}
284284
wrappedRecorder := aibridge.NewRecorder(logger, tracer, clientFn)
285285

286-
bridge, err := aibridge.NewRequestBridge(ctx, []aibridge.Provider{provider}, wrappedRecorder, mcp.NewServerProxyManager(nil), metrics, tracer, logger)
286+
bridge, err := aibridge.NewRequestBridge(ctx, []aibridge.Provider{provider}, wrappedRecorder, mcp.NewServerProxyManager(nil, defaultTracer), metrics, tracer, logger)
287287
require.NoError(t, err)
288288

289289
srv := httptest.NewUnstartedServer(bridge)

0 commit comments

Comments
 (0)