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

Commitd09c5e8

Browse files
committed
feat: aibridged mcp handling
Signed-off-by: Danny Kopping <danny@coder.com>
1 parentb6bb2fe commitd09c5e8

File tree

5 files changed

+57
-12
lines changed

5 files changed

+57
-12
lines changed

‎x/aibridged/aibridged.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -139,7 +139,7 @@ func (s *Server) GetRequestHandler(ctx context.Context, req Request) (http.Handl
139139
returnnil,xerrors.New("nil requestBridgePool")
140140
}
141141

142-
reqBridge,err:=s.requestBridgePool.Acquire(ctx,req,s.Client)
142+
reqBridge,err:=s.requestBridgePool.Acquire(ctx,req,s.Client,NewMCPProxyFactory(s.logger,s.Client))
143143
iferr!=nil {
144144
returnnil,xerrors.Errorf("acquire request bridge: %w",err)
145145
}

‎x/aibridged/aibridged_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -122,7 +122,7 @@ func TestServeHTTP_FailureModes(t *testing.T) {
122122
// Should pass authorization.
123123
client.EXPECT().IsAuthorized(gomock.Any(),gomock.Any()).AnyTimes().Return(&proto.IsAuthorizedResponse{OwnerId:uuid.NewString()},nil)
124124
// But fail when acquiring a pool instance.
125-
pool.EXPECT().Acquire(gomock.Any(),gomock.Any(),gomock.Any()).AnyTimes().Return(nil,xerrors.New("oops"))
125+
pool.EXPECT().Acquire(gomock.Any(),gomock.Any(),gomock.Any(),gomock.Any()).AnyTimes().Return(nil,xerrors.New("oops"))
126126
},
127127
expectedErr:aibridged.ErrAcquireRequestHandler,
128128
expectedStatus:http.StatusInternalServerError,

‎x/aibridged/aibridgedmock/poolmock.go‎

Lines changed: 4 additions & 4 deletions
Some generated files are not rendered by default. Learn more aboutcustomizing how changed files appear on GitHub.

‎x/aibridged/pool.go‎

Lines changed: 22 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ import (
1414
"cdr.dev/slog"
1515

1616
"github.com/coder/aibridge"
17+
"github.com/coder/aibridge/mcp"
1718
)
1819

1920
const (
@@ -23,7 +24,7 @@ const (
2324
// Pooler describes a pool of [*aibridge.RequestBridge] instances from which instances can be retrieved.
2425
// One [*aibridge.RequestBridge] instance is created per given key.
2526
typePoolerinterface {
26-
Acquire(ctx context.Context,reqRequest,clientFnClientFunc) (http.Handler,error)
27+
Acquire(ctx context.Context,reqRequest,clientFnClientFunc,mcpBootstrapperMCPProxyBuilder) (http.Handler,error)
2728
Shutdown(ctx context.Context)error
2829
}
2930

@@ -102,7 +103,7 @@ func NewCachedBridgePool(options PoolOptions, providers []aibridge.Provider, log
102103
//
103104
// Each returned [*aibridge.RequestBridge] is safe for concurrent use.
104105
// Each [*aibridge.RequestBridge] is stateful because it has MCP clients which maintain sessions to the configured MCP server.
105-
func (p*CachedBridgePool)Acquire(ctx context.Context,reqRequest,clientFnClientFunc) (http.Handler,error) {
106+
func (p*CachedBridgePool)Acquire(ctx context.Context,reqRequest,clientFnClientFunc,mcpProxyFactoryMCPProxyBuilder) (http.Handler,error) {
106107
iferr:=ctx.Err();err!=nil {
107108
returnnil,xerrors.Errorf("acquire: %w",err)
108109
}
@@ -141,7 +142,25 @@ func (p *CachedBridgePool) Acquire(ctx context.Context, req Request, clientFn Cl
141142
// Creating an *aibridge.RequestBridge may take some time, so gate all subsequent callers behind the initial request and return the resulting value.
142143
// TODO: track startup time since it adds latency to first request (histogram count will also help us see how often this occurs).
143144
instance,err,_:=p.singleflight.Do(req.InitiatorID.String(),func() (*aibridge.RequestBridge,error) {
144-
bridge,err:=aibridge.NewRequestBridge(ctx,p.providers,p.logger,recorder,nil)
145+
var (
146+
mcpServers mcp.ServerProxier
147+
errerror
148+
)
149+
150+
mcpServers,err=mcpProxyFactory.Build(ctx,req)
151+
iferr!=nil {
152+
p.logger.Warn(ctx,"failed to create MCP server proxiers",slog.Error(err))
153+
// Don't fail here; MCP server injection can gracefully degrade.
154+
}
155+
156+
ifmcpServers!=nil {
157+
// This will block while connections are established with upstream MCP server(s), and tools are listed.
158+
iferr:=mcpServers.Init(ctx);err!=nil {
159+
p.logger.Warn(ctx,"failed to initialize MCP server proxier(s)",slog.Error(err))
160+
}
161+
}
162+
163+
bridge,err:=aibridge.NewRequestBridge(ctx,p.providers,p.logger,recorder,mcpServers)
145164
iferr!=nil {
146165
returnnil,xerrors.Errorf("create new request bridge: %w",err)
147166
}

‎x/aibridged/pool_test.go‎

Lines changed: 29 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@ import (
1111
"go.uber.org/mock/gomock"
1212

1313
"cdr.dev/slog/sloggers/slogtest"
14+
"github.com/coder/aibridge/mcp"
15+
"github.com/coder/aibridge/mcpmock"
1416
"github.com/coder/coder/v2/x/aibridged"
1517
mock"github.com/coder/coder/v2/x/aibridged/aibridgedmock"
1618
)
@@ -25,6 +27,7 @@ func TestPool(t *testing.T) {
2527

2628
ctrl:=gomock.NewController(t)
2729
client:=mock.NewMockDRPCClient(ctrl)
30+
mcpProxy:=mcpmock.NewMockServerProxier(ctrl)
2831

2932
opts:= aibridged.PoolOptions{MaxItems:1,TTL:time.Second}
3033
pool,err:=aibridged.NewCachedBridgePool(opts,nil,logger)
@@ -36,19 +39,25 @@ func TestPool(t *testing.T) {
3639
returnclient,nil
3740
}
3841

42+
// Once a pool instance is initialized, it will try setup its MCP proxier(s).
43+
// This is called exactly once since the instance below is only created once.
44+
mcpProxy.EXPECT().Init(gomock.Any()).Times(1).Return(nil)
45+
// This is part of the lifecycle.
46+
mcpProxy.EXPECT().Shutdown(gomock.Any()).AnyTimes().Return(nil)
47+
3948
// Acquiring a pool instance will create one the first time it sees an
4049
// initiator ID...
4150
inst,err:=pool.Acquire(t.Context(), aibridged.Request{
4251
SessionKey:"key",
4352
InitiatorID:id,
44-
},clientFn)
53+
},clientFn,newMockMCPFactory(mcpProxy))
4554
require.NoError(t,err,"acquire pool instance")
4655

4756
// ...and it will return it when acquired again.
4857
instB,err:=pool.Acquire(t.Context(), aibridged.Request{
4958
SessionKey:"key",
5059
InitiatorID:id,
51-
},clientFn)
60+
},clientFn,newMockMCPFactory(mcpProxy))
5261
require.NoError(t,err,"acquire pool instance")
5362
require.Same(t,inst,instB)
5463

@@ -58,11 +67,14 @@ func TestPool(t *testing.T) {
5867
require.EqualValues(t,1,metrics.Hits())
5968
require.EqualValues(t,1,metrics.Misses())
6069

70+
// This will get called again because a new instance will be created.
71+
mcpProxy.EXPECT().Init(gomock.Any()).Times(1).Return(nil)
72+
6173
// But that key will be evicted when a new initiator is seen (maxItems=1):
6274
inst2,err:=pool.Acquire(t.Context(), aibridged.Request{
6375
SessionKey:"key",
6476
InitiatorID:id2,
65-
},clientFn)
77+
},clientFn,newMockMCPFactory(mcpProxy))
6678
require.NoError(t,err,"acquire pool instance")
6779
require.NotSame(t,inst,inst2)
6880

@@ -76,3 +88,17 @@ func TestPool(t *testing.T) {
7688
// This requires Go 1.25's [synctest](https://pkg.go.dev/testing/synctest) since the
7789
// internal cache lib cannot be tested using coder/quartz.
7890
}
91+
92+
var_ aibridged.MCPProxyBuilder=&mockMCPFactory{}
93+
94+
typemockMCPFactorystruct {
95+
proxy*mcpmock.MockServerProxier
96+
}
97+
98+
funcnewMockMCPFactory(proxy*mcpmock.MockServerProxier)*mockMCPFactory {
99+
return&mockMCPFactory{proxy:proxy}
100+
}
101+
102+
func (m*mockMCPFactory)Build(ctx context.Context,req aibridged.Request) (mcp.ServerProxier,error) {
103+
returnm.proxy,nil
104+
}

0 commit comments

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