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

Commitb72efe4

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

File tree

7 files changed

+308
-12
lines changed

7 files changed

+308
-12
lines changed

‎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
}

‎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,

‎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.

‎aibridged/mcp.go‎

Lines changed: 190 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,190 @@
1+
package aibridged
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"regexp"
7+
"time"
8+
9+
"golang.org/x/xerrors"
10+
11+
"cdr.dev/slog"
12+
"github.com/coder/aibridge/mcp"
13+
"github.com/coder/coder/v2/aibridged/proto"
14+
)
15+
16+
var (
17+
ErrEmptyConfig=xerrors.New("empty config given")
18+
ErrCompileRegex=xerrors.New("compile tool regex")
19+
)
20+
21+
typeMCPProxyBuilderinterface {
22+
// Build creates a [mcp.ServerProxier] for the given request initiator.
23+
// At minimum, the Coder MCP server will be proxied.
24+
// The SessionKey from [Request] is used to authenticate against the Coder MCP server.
25+
//
26+
// NOTE: the [mcp.ServerProxier] instance may be proxying one or more MCP servers.
27+
Build(ctx context.Context,reqRequest) (mcp.ServerProxier,error)
28+
}
29+
30+
var_MCPProxyBuilder=&MCPProxyFactory{}
31+
32+
typeMCPProxyFactorystruct {
33+
logger slog.Logger
34+
clientFnClientFunc
35+
}
36+
37+
funcNewMCPProxyFactory(logger slog.Logger,clientFnClientFunc)*MCPProxyFactory {
38+
return&MCPProxyFactory{
39+
logger:logger,
40+
clientFn:clientFn,
41+
}
42+
}
43+
44+
func (m*MCPProxyFactory)Build(ctx context.Context,reqRequest) (mcp.ServerProxier,error) {
45+
proxiers,err:=m.retrieveMCPServerConfigs(ctx,req)
46+
iferr!=nil {
47+
returnnil,xerrors.Errorf("resolve configs: %w",err)
48+
}
49+
50+
returnmcp.NewServerProxyManager(proxiers),nil
51+
}
52+
53+
func (m*MCPProxyFactory)retrieveMCPServerConfigs(ctx context.Context,reqRequest) (map[string]mcp.ServerProxier,error) {
54+
client,err:=m.clientFn()
55+
iferr!=nil {
56+
returnnil,xerrors.Errorf("acquire client: %w",err)
57+
}
58+
59+
srvCfgCtx,srvCfgCancel:=context.WithTimeout(ctx,time.Second*10)
60+
defersrvCfgCancel()
61+
62+
// Fetch MCP server configs.
63+
mcpSrvCfgs,err:=client.GetMCPServerConfigs(srvCfgCtx,&proto.GetMCPServerConfigsRequest{
64+
UserId:req.InitiatorID.String(),
65+
})
66+
iferr!=nil {
67+
returnnil,xerrors.Errorf("get MCP server configs: %w",err)
68+
}
69+
70+
proxiers:=make(map[string]mcp.ServerProxier,len(mcpSrvCfgs.GetExternalAuthMcpConfigs())+1)// Extra one for Coder MCP server.
71+
72+
ifmcpSrvCfgs.GetCoderMcpConfig()!=nil {
73+
// Setup the Coder MCP server proxy.
74+
coderMCPProxy,err:=m.newStreamableHTTPServerProxy(mcpSrvCfgs.GetCoderMcpConfig(),req.SessionKey)// The session key is used to auth against our internal MCP server.
75+
iferr!=nil {
76+
m.logger.Warn(ctx,"failed to create MCP server proxy",slog.F("mcp_server_id",mcpSrvCfgs.GetCoderMcpConfig().GetId()),slog.Error(err))
77+
}else {
78+
proxiers["coder"]=coderMCPProxy
79+
}
80+
}
81+
82+
iflen(mcpSrvCfgs.GetExternalAuthMcpConfigs())==0 {
83+
returnproxiers,nil
84+
}
85+
86+
serverIDs:=make([]string,0,len(mcpSrvCfgs.GetExternalAuthMcpConfigs()))
87+
for_,cfg:=rangemcpSrvCfgs.GetExternalAuthMcpConfigs() {
88+
serverIDs=append(serverIDs,cfg.GetId())
89+
}
90+
91+
accTokCtx,accTokCancel:=context.WithTimeout(ctx,time.Second*10)
92+
deferaccTokCancel()
93+
94+
// Request a batch of access tokens, one per given server ID.
95+
resp,err:=client.GetMCPServerAccessTokensBatch(accTokCtx,&proto.GetMCPServerAccessTokensBatchRequest{
96+
UserId:req.InitiatorID.String(),
97+
McpServerConfigIds:serverIDs,
98+
})
99+
iferr!=nil {
100+
m.logger.Warn(ctx,"failed to retrieve access token(s)",slog.F("server_ids",serverIDs),slog.Error(err))
101+
}
102+
103+
ifresp==nil {
104+
returnproxiers,nil
105+
}
106+
tokens:=resp.GetAccessTokens()
107+
iflen(tokens)==0 {
108+
returnproxiers,nil
109+
}
110+
111+
forid,tokErr:=rangeresp.GetErrors() {
112+
m.logger.Warn(ctx,"failed to retrieve access token",slog.F("server_id",id),slog.F("error",tokErr))
113+
}
114+
115+
// Iterate over all External Auth configurations which are configured for MCP and attempt to setup
116+
// a [mcp.ServerProxier] for it using the access token retrieved above.
117+
for_,cfg:=rangemcpSrvCfgs.GetExternalAuthMcpConfigs() {
118+
iferr,ok:=resp.GetErrors()[cfg.GetId()];ok {
119+
m.logger.Warn(ctx,"failed to get access token",slog.F("mcp_server_id",cfg.GetId()),slog.F("error",err))
120+
continue
121+
}
122+
123+
token,ok:=tokens[cfg.GetId()]
124+
if!ok {
125+
m.logger.Warn(ctx,"no access token found",slog.F("mcp_server_id",cfg.GetId()))
126+
continue
127+
}
128+
129+
proxy,err:=m.newStreamableHTTPServerProxy(cfg,token)
130+
iferr!=nil {
131+
m.logger.Warn(ctx,"failed to create MCP server proxy",slog.F("mcp_server_id",cfg.GetId()),slog.Error(err))
132+
continue
133+
}
134+
135+
proxiers[cfg.Id]=proxy
136+
}
137+
returnproxiers,nil
138+
}
139+
140+
// newStreamableHTTPServerProxy creates an MCP server capable of proxying requests using the Streamable HTTP transport.
141+
//
142+
// TODO: support SSE transport.
143+
func (m*MCPProxyFactory)newStreamableHTTPServerProxy(cfg*proto.MCPServerConfig,accessTokenstring) (mcp.ServerProxier,error) {
144+
ifcfg==nil {
145+
returnnil,ErrEmptyConfig
146+
}
147+
148+
var (
149+
allowlist,denylist*regexp.Regexp
150+
errerror
151+
)
152+
ifcfg.GetToolAllowRegex()!="" {
153+
allowlist,err=regexp.Compile(cfg.GetToolAllowRegex())
154+
iferr!=nil {
155+
returnnil,ErrCompileRegex
156+
}
157+
}
158+
ifcfg.GetToolDenyRegex()!="" {
159+
denylist,err=regexp.Compile(cfg.GetToolDenyRegex())
160+
iferr!=nil {
161+
returnnil,ErrCompileRegex
162+
}
163+
}
164+
165+
// TODO: future improvement:
166+
//
167+
// The access token provided here may expire at any time, or the connection to the MCP server could be severed.
168+
// Instead of passing through an access token directly, rather provide an interface through which to retrieve
169+
// an access token imperatively. In the event of a tool call failing, we could Ping() the MCP server to establish
170+
// whether the connection is still active. If not, this indicates that the access token is probably expired/revoked.
171+
// (It could also mean the server has a problem, which we should account for.)
172+
// The proxy could then use its interface to retrieve a new access token and re-establish a connection.
173+
// For now though, the short TTL of this cache should mostly mask this problem.
174+
srv,err:=mcp.NewStreamableHTTPServerProxy(
175+
m.logger.Named(fmt.Sprintf("mcp-server-proxy-%s",cfg.GetId())),
176+
cfg.GetId(),
177+
cfg.GetUrl(),
178+
// See https://modelcontextprotocol.io/specification/2025-06-18/basic/authorization#token-requirements.
179+
map[string]string{
180+
"Authorization":fmt.Sprintf("Bearer %s",accessToken),
181+
},
182+
allowlist,
183+
denylist,
184+
)
185+
iferr!=nil {
186+
returnnil,xerrors.Errorf("create streamable HTTP MCP server proxy: %w",err)
187+
}
188+
189+
returnsrv,nil
190+
}

‎aibridged/mcp_internal_test.go‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
package aibridged
2+
3+
import (
4+
"testing"
5+
6+
"github.com/stretchr/testify/require"
7+
8+
"github.com/coder/coder/v2/aibridged/proto"
9+
"github.com/coder/coder/v2/testutil"
10+
)
11+
12+
funcTestMCPRegex(t*testing.T) {
13+
t.Parallel()
14+
15+
cases:= []struct {
16+
namestring
17+
allowRegex,denyRegexstring
18+
expectedErrerror
19+
}{
20+
{
21+
name:"invalid allow regex",
22+
allowRegex:`\`,
23+
expectedErr:ErrCompileRegex,
24+
},
25+
{
26+
name:"invalid deny regex",
27+
denyRegex:`+`,
28+
expectedErr:ErrCompileRegex,
29+
},
30+
{
31+
name:"valid empty",
32+
},
33+
{
34+
name:"valid",
35+
allowRegex:"(allowed|allowed2)",
36+
denyRegex:".*disallowed.*",
37+
},
38+
}
39+
40+
for_,tc:=rangecases {
41+
t.Run(tc.name,func(t*testing.T) {
42+
t.Parallel()
43+
44+
logger:=testutil.Logger(t)
45+
f:=NewMCPProxyFactory(logger,nil)
46+
47+
_,err:=f.newStreamableHTTPServerProxy(&proto.MCPServerConfig{
48+
Id:"mock",
49+
Url:"mock/mcp",
50+
ToolAllowRegex:tc.allowRegex,
51+
ToolDenyRegex:tc.denyRegex,
52+
},"")
53+
54+
iftc.expectedErr==nil {
55+
require.NoError(t,err)
56+
}else {
57+
require.ErrorIs(t,err,tc.expectedErr)
58+
}
59+
})
60+
}
61+
}

‎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
}

0 commit comments

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