diff --git a/enterprise/x/aibridged/aibridged.go b/enterprise/x/aibridged/aibridged.go index ee1df6084f..04ae617c20 100644 --- a/enterprise/x/aibridged/aibridged.go +++ b/enterprise/x/aibridged/aibridged.go @@ -139,7 +139,7 @@ func (s *Server) GetRequestHandler(ctx context.Context, req Request) (http.Handl return nil, xerrors.New("nil requestBridgePool") } - reqBridge, err := s.requestBridgePool.Acquire(ctx, req, s.Client) + reqBridge, err := s.requestBridgePool.Acquire(ctx, req, s.Client, NewMCPProxyFactory(s.logger, s.Client)) if err != nil { return nil, xerrors.Errorf("acquire request bridge: %w", err) } diff --git a/enterprise/x/aibridged/aibridged_test.go b/enterprise/x/aibridged/aibridged_test.go index 0d6541260c..22210a00f3 100644 --- a/enterprise/x/aibridged/aibridged_test.go +++ b/enterprise/x/aibridged/aibridged_test.go @@ -122,7 +122,7 @@ func TestServeHTTP_FailureModes(t *testing.T) { // Should pass authorization. client.EXPECT().IsAuthorized(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.IsAuthorizedResponse{OwnerId: uuid.NewString()}, nil) // But fail when acquiring a pool instance. - pool.EXPECT().Acquire(gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes().Return(nil, xerrors.New("oops")) + pool.EXPECT().Acquire(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes().Return(nil, xerrors.New("oops")) }, expectedErr: aibridged.ErrAcquireRequestHandler, expectedStatus: http.StatusInternalServerError, diff --git a/enterprise/x/aibridged/aibridgedmock/poolmock.go b/enterprise/x/aibridged/aibridgedmock/poolmock.go index 80ce4a35b9..bf3b39ed2a 100644 --- a/enterprise/x/aibridged/aibridgedmock/poolmock.go +++ b/enterprise/x/aibridged/aibridgedmock/poolmock.go @@ -43,18 +43,18 @@ func (m *MockPooler) EXPECT() *MockPoolerMockRecorder { } // Acquire mocks base method. -func (m *MockPooler) Acquire(ctx context.Context, req aibridged.Request, clientFn aibridged.ClientFunc) (http.Handler, error) { +func (m *MockPooler) Acquire(ctx context.Context, req aibridged.Request, clientFn aibridged.ClientFunc, mcpBootstrapper aibridged.MCPProxyBuilder) (http.Handler, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "Acquire", ctx, req, clientFn) + ret := m.ctrl.Call(m, "Acquire", ctx, req, clientFn, mcpBootstrapper) ret0, _ := ret[0].(http.Handler) ret1, _ := ret[1].(error) return ret0, ret1 } // Acquire indicates an expected call of Acquire. -func (mr *MockPoolerMockRecorder) Acquire(ctx, req, clientFn any) *gomock.Call { +func (mr *MockPoolerMockRecorder) Acquire(ctx, req, clientFn, mcpBootstrapper any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Acquire", reflect.TypeOf((*MockPooler)(nil).Acquire), ctx, req, clientFn) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Acquire", reflect.TypeOf((*MockPooler)(nil).Acquire), ctx, req, clientFn, mcpBootstrapper) } // Shutdown mocks base method. diff --git a/enterprise/x/aibridged/mcp.go b/enterprise/x/aibridged/mcp.go new file mode 100644 index 0000000000..4b42287e02 --- /dev/null +++ b/enterprise/x/aibridged/mcp.go @@ -0,0 +1,191 @@ +package aibridged + +import ( + "context" + "fmt" + "regexp" + "time" + + "golang.org/x/xerrors" + + "cdr.dev/slog" + "github.com/coder/aibridge/mcp" + "github.com/coder/coder/v2/enterprise/x/aibridged/proto" +) + +var ( + ErrEmptyConfig = xerrors.New("empty config given") + ErrCompileRegex = xerrors.New("compile tool regex") +) + +const ( + InternalMCPServerID = "coder" +) + +type MCPProxyBuilder interface { + // Build creates a [mcp.ServerProxier] for the given request initiator. + // At minimum, the Coder MCP server will be proxied. + // The SessionKey from [Request] is used to authenticate against the Coder MCP server. + // + // NOTE: the [mcp.ServerProxier] instance may be proxying one or more MCP servers. + Build(ctx context.Context, req Request) (mcp.ServerProxier, error) +} + +var _ MCPProxyBuilder = &MCPProxyFactory{} + +type MCPProxyFactory struct { + logger slog.Logger + clientFn ClientFunc +} + +func NewMCPProxyFactory(logger slog.Logger, clientFn ClientFunc) *MCPProxyFactory { + return &MCPProxyFactory{ + logger: logger, + clientFn: clientFn, + } +} + +func (m *MCPProxyFactory) Build(ctx context.Context, req Request) (mcp.ServerProxier, error) { + proxiers, err := m.retrieveMCPServerConfigs(ctx, req) + if err != nil { + return nil, xerrors.Errorf("resolve configs: %w", err) + } + + return mcp.NewServerProxyManager(proxiers), nil +} + +func (m *MCPProxyFactory) retrieveMCPServerConfigs(ctx context.Context, req Request) (map[string]mcp.ServerProxier, error) { + client, err := m.clientFn() + if err != nil { + return nil, xerrors.Errorf("acquire client: %w", err) + } + + srvCfgCtx, srvCfgCancel := context.WithTimeout(ctx, time.Second*10) + defer srvCfgCancel() + + // Fetch MCP server configs. + mcpSrvCfgs, err := client.GetMCPServerConfigs(srvCfgCtx, &proto.GetMCPServerConfigsRequest{ + UserId: req.InitiatorID.String(), + }) + if err != nil { + return nil, xerrors.Errorf("get MCP server configs: %w", err) + } + + proxiers := make(map[string]mcp.ServerProxier, len(mcpSrvCfgs.GetExternalAuthMcpConfigs())+1) // Extra one for Coder MCP server. + + if mcpSrvCfgs.GetCoderMcpConfig() != nil { + // Setup the Coder MCP server proxy. + coderMCPProxy, err := m.newStreamableHTTPServerProxy(mcpSrvCfgs.GetCoderMcpConfig(), req.SessionKey) // The session key is used to auth against our internal MCP server. + if err != nil { + m.logger.Warn(ctx, "failed to create MCP server proxy", slog.F("mcp_server_id", mcpSrvCfgs.GetCoderMcpConfig().GetId()), slog.Error(err)) + } else { + proxiers[InternalMCPServerID] = coderMCPProxy + } + } + + if len(mcpSrvCfgs.GetExternalAuthMcpConfigs()) == 0 { + return proxiers, nil + } + + serverIDs := make([]string, 0, len(mcpSrvCfgs.GetExternalAuthMcpConfigs())) + for _, cfg := range mcpSrvCfgs.GetExternalAuthMcpConfigs() { + serverIDs = append(serverIDs, cfg.GetId()) + } + + accTokCtx, accTokCancel := context.WithTimeout(ctx, time.Second*10) + defer accTokCancel() + + // Request a batch of access tokens, one per given server ID. + resp, err := client.GetMCPServerAccessTokensBatch(accTokCtx, &proto.GetMCPServerAccessTokensBatchRequest{ + UserId: req.InitiatorID.String(), + McpServerConfigIds: serverIDs, + }) + if err != nil { + m.logger.Warn(ctx, "failed to retrieve access token(s)", slog.F("server_ids", serverIDs), slog.Error(err)) + } + + if resp == nil { + m.logger.Warn(ctx, "nil response given to mcp access tokens call") + return proxiers, nil + } + tokens := resp.GetAccessTokens() + if len(tokens) == 0 { + return proxiers, nil + } + + // Iterate over all External Auth configurations which are configured for MCP and attempt to setup + // a [mcp.ServerProxier] for it using the access token retrieved above. + for _, cfg := range mcpSrvCfgs.GetExternalAuthMcpConfigs() { + if err, ok := resp.GetErrors()[cfg.GetId()]; ok { + m.logger.Debug(ctx, "failed to get access token", slog.F("mcp_server_id", cfg.GetId()), slog.F("error", err)) + continue + } + + token, ok := tokens[cfg.GetId()] + if !ok { + m.logger.Warn(ctx, "no access token found", slog.F("mcp_server_id", cfg.GetId())) + continue + } + + proxy, err := m.newStreamableHTTPServerProxy(cfg, token) + if err != nil { + m.logger.Warn(ctx, "failed to create MCP server proxy", slog.F("mcp_server_id", cfg.GetId()), slog.Error(err)) + continue + } + + proxiers[cfg.Id] = proxy + } + return proxiers, nil +} + +// newStreamableHTTPServerProxy creates an MCP server capable of proxying requests using the Streamable HTTP transport. +// +// TODO: support SSE transport. +func (m *MCPProxyFactory) newStreamableHTTPServerProxy(cfg *proto.MCPServerConfig, accessToken string) (mcp.ServerProxier, error) { + if cfg == nil { + return nil, ErrEmptyConfig + } + + var ( + allowlist, denylist *regexp.Regexp + err error + ) + if cfg.GetToolAllowRegex() != "" { + allowlist, err = regexp.Compile(cfg.GetToolAllowRegex()) + if err != nil { + return nil, ErrCompileRegex + } + } + if cfg.GetToolDenyRegex() != "" { + denylist, err = regexp.Compile(cfg.GetToolDenyRegex()) + if err != nil { + return nil, ErrCompileRegex + } + } + + // TODO: future improvement: + // + // The access token provided here may expire at any time, or the connection to the MCP server could be severed. + // Instead of passing through an access token directly, rather provide an interface through which to retrieve + // an access token imperatively. In the event of a tool call failing, we could Ping() the MCP server to establish + // whether the connection is still active. If not, this indicates that the access token is probably expired/revoked. + // (It could also mean the server has a problem, which we should account for.) + // The proxy could then use its interface to retrieve a new access token and re-establish a connection. + // For now though, the short TTL of this cache should mostly mask this problem. + srv, err := mcp.NewStreamableHTTPServerProxy( + m.logger.Named(fmt.Sprintf("mcp-server-proxy-%s", cfg.GetId())), + cfg.GetId(), + cfg.GetUrl(), + // See https://modelcontextprotocol.io/specification/2025-06-18/basic/authorization#token-requirements. + map[string]string{ + "Authorization": fmt.Sprintf("Bearer %s", accessToken), + }, + allowlist, + denylist, + ) + if err != nil { + return nil, xerrors.Errorf("create streamable HTTP MCP server proxy: %w", err) + } + + return srv, nil +} diff --git a/enterprise/x/aibridged/mcp_internal_test.go b/enterprise/x/aibridged/mcp_internal_test.go new file mode 100644 index 0000000000..20edf79d06 --- /dev/null +++ b/enterprise/x/aibridged/mcp_internal_test.go @@ -0,0 +1,61 @@ +package aibridged + +import ( + "testing" + + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/enterprise/x/aibridged/proto" + "github.com/coder/coder/v2/testutil" +) + +func TestMCPRegex(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + allowRegex, denyRegex string + expectedErr error + }{ + { + name: "invalid allow regex", + allowRegex: `\`, + expectedErr: ErrCompileRegex, + }, + { + name: "invalid deny regex", + denyRegex: `+`, + expectedErr: ErrCompileRegex, + }, + { + name: "valid empty", + }, + { + name: "valid", + allowRegex: "(allowed|allowed2)", + denyRegex: ".*disallowed.*", + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + logger := testutil.Logger(t) + f := NewMCPProxyFactory(logger, nil) + + _, err := f.newStreamableHTTPServerProxy(&proto.MCPServerConfig{ + Id: "mock", + Url: "mock/mcp", + ToolAllowRegex: tc.allowRegex, + ToolDenyRegex: tc.denyRegex, + }, "") + + if tc.expectedErr == nil { + require.NoError(t, err) + } else { + require.ErrorIs(t, err, tc.expectedErr) + } + }) + } +} diff --git a/enterprise/x/aibridged/pool.go b/enterprise/x/aibridged/pool.go index 01de7059a3..97c08703c7 100644 --- a/enterprise/x/aibridged/pool.go +++ b/enterprise/x/aibridged/pool.go @@ -14,6 +14,7 @@ import ( "cdr.dev/slog" "github.com/coder/aibridge" + "github.com/coder/aibridge/mcp" ) const ( @@ -23,7 +24,7 @@ const ( // Pooler describes a pool of [*aibridge.RequestBridge] instances from which instances can be retrieved. // One [*aibridge.RequestBridge] instance is created per given key. type Pooler interface { - Acquire(ctx context.Context, req Request, clientFn ClientFunc) (http.Handler, error) + Acquire(ctx context.Context, req Request, clientFn ClientFunc, mcpBootstrapper MCPProxyBuilder) (http.Handler, error) Shutdown(ctx context.Context) error } @@ -102,7 +103,7 @@ func NewCachedBridgePool(options PoolOptions, providers []aibridge.Provider, log // // Each returned [*aibridge.RequestBridge] is safe for concurrent use. // Each [*aibridge.RequestBridge] is stateful because it has MCP clients which maintain sessions to the configured MCP server. -func (p *CachedBridgePool) Acquire(ctx context.Context, req Request, clientFn ClientFunc) (http.Handler, error) { +func (p *CachedBridgePool) Acquire(ctx context.Context, req Request, clientFn ClientFunc, mcpProxyFactory MCPProxyBuilder) (http.Handler, error) { if err := ctx.Err(); err != nil { return nil, xerrors.Errorf("acquire: %w", err) } @@ -141,7 +142,25 @@ func (p *CachedBridgePool) Acquire(ctx context.Context, req Request, clientFn Cl // Creating an *aibridge.RequestBridge may take some time, so gate all subsequent callers behind the initial request and return the resulting value. // TODO: track startup time since it adds latency to first request (histogram count will also help us see how often this occurs). instance, err, _ := p.singleflight.Do(req.InitiatorID.String(), func() (*aibridge.RequestBridge, error) { - bridge, err := aibridge.NewRequestBridge(ctx, p.providers, p.logger, recorder, nil) + var ( + mcpServers mcp.ServerProxier + err error + ) + + mcpServers, err = mcpProxyFactory.Build(ctx, req) + if err != nil { + p.logger.Warn(ctx, "failed to create MCP server proxiers", slog.Error(err)) + // Don't fail here; MCP server injection can gracefully degrade. + } + + if mcpServers != nil { + // This will block while connections are established with upstream MCP server(s), and tools are listed. + if err := mcpServers.Init(ctx); err != nil { + p.logger.Warn(ctx, "failed to initialize MCP server proxier(s)", slog.Error(err)) + } + } + + bridge, err := aibridge.NewRequestBridge(ctx, p.providers, p.logger, recorder, mcpServers) if err != nil { return nil, xerrors.Errorf("create new request bridge: %w", err) } diff --git a/enterprise/x/aibridged/pool_test.go b/enterprise/x/aibridged/pool_test.go index c3442a80de..38cae85da9 100644 --- a/enterprise/x/aibridged/pool_test.go +++ b/enterprise/x/aibridged/pool_test.go @@ -11,6 +11,8 @@ import ( "go.uber.org/mock/gomock" "cdr.dev/slog/sloggers/slogtest" + "github.com/coder/aibridge/mcp" + "github.com/coder/aibridge/mcpmock" "github.com/coder/coder/v2/enterprise/x/aibridged" mock "github.com/coder/coder/v2/enterprise/x/aibridged/aibridgedmock" ) @@ -25,6 +27,7 @@ func TestPool(t *testing.T) { ctrl := gomock.NewController(t) client := mock.NewMockDRPCClient(ctrl) + mcpProxy := mcpmock.NewMockServerProxier(ctrl) opts := aibridged.PoolOptions{MaxItems: 1, TTL: time.Second} pool, err := aibridged.NewCachedBridgePool(opts, nil, logger) @@ -36,19 +39,25 @@ func TestPool(t *testing.T) { return client, nil } + // Once a pool instance is initialized, it will try setup its MCP proxier(s). + // This is called exactly once since the instance below is only created once. + mcpProxy.EXPECT().Init(gomock.Any()).Times(1).Return(nil) + // This is part of the lifecycle. + mcpProxy.EXPECT().Shutdown(gomock.Any()).AnyTimes().Return(nil) + // Acquiring a pool instance will create one the first time it sees an // initiator ID... inst, err := pool.Acquire(t.Context(), aibridged.Request{ SessionKey: "key", InitiatorID: id, - }, clientFn) + }, clientFn, newMockMCPFactory(mcpProxy)) require.NoError(t, err, "acquire pool instance") // ...and it will return it when acquired again. instB, err := pool.Acquire(t.Context(), aibridged.Request{ SessionKey: "key", InitiatorID: id, - }, clientFn) + }, clientFn, newMockMCPFactory(mcpProxy)) require.NoError(t, err, "acquire pool instance") require.Same(t, inst, instB) @@ -58,11 +67,14 @@ func TestPool(t *testing.T) { require.EqualValues(t, 1, metrics.Hits()) require.EqualValues(t, 1, metrics.Misses()) + // This will get called again because a new instance will be created. + mcpProxy.EXPECT().Init(gomock.Any()).Times(1).Return(nil) + // But that key will be evicted when a new initiator is seen (maxItems=1): inst2, err := pool.Acquire(t.Context(), aibridged.Request{ SessionKey: "key", InitiatorID: id2, - }, clientFn) + }, clientFn, newMockMCPFactory(mcpProxy)) require.NoError(t, err, "acquire pool instance") require.NotSame(t, inst, inst2) @@ -76,3 +88,17 @@ func TestPool(t *testing.T) { // This requires Go 1.25's [synctest](https://pkg.go.dev/testing/synctest) since the // internal cache lib cannot be tested using coder/quartz. } + +var _ aibridged.MCPProxyBuilder = &mockMCPFactory{} + +type mockMCPFactory struct { + proxy *mcpmock.MockServerProxier +} + +func newMockMCPFactory(proxy *mcpmock.MockServerProxier) *mockMCPFactory { + return &mockMCPFactory{proxy: proxy} +} + +func (m *mockMCPFactory) Build(ctx context.Context, req aibridged.Request) (mcp.ServerProxier, error) { + return m.proxy, nil +} diff --git a/enterprise/x/aibridgedserver/aibridgedserver.go b/enterprise/x/aibridgedserver/aibridgedserver.go index c3b0149989..9f36d90517 100644 --- a/enterprise/x/aibridgedserver/aibridgedserver.go +++ b/enterprise/x/aibridgedserver/aibridgedserver.go @@ -399,7 +399,7 @@ func getCoderMCPServerConfig(experiments codersdk.Experiments, accessURL string) } return &proto.MCPServerConfig{ - Id: "coder", + Id: aibridged.InternalMCPServerID, Url: u, }, nil } diff --git a/enterprise/x/aibridgedserver/aibridgedserver_test.go b/enterprise/x/aibridgedserver/aibridgedserver_test.go index 2e8e4c9c87..c464c7c368 100644 --- a/enterprise/x/aibridgedserver/aibridgedserver_test.go +++ b/enterprise/x/aibridgedserver/aibridgedserver_test.go @@ -28,6 +28,7 @@ import ( codermcp "github.com/coder/coder/v2/coderd/mcp" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/cryptorand" + "github.com/coder/coder/v2/enterprise/x/aibridged" "github.com/coder/coder/v2/enterprise/x/aibridged/proto" "github.com/coder/coder/v2/enterprise/x/aibridgedserver" "github.com/coder/coder/v2/testutil" @@ -255,7 +256,7 @@ func TestGetMCPServerConfigs(t *testing.T) { if tc.expectCoderMCP { coderConfig := resp.CoderMcpConfig require.NotNil(t, coderConfig) - require.Equal(t, "coder", coderConfig.GetId()) + require.Equal(t, aibridged.InternalMCPServerID, coderConfig.GetId()) expectedURL, err := url.JoinPath(accessURL, codermcp.MCPEndpoint) require.NoError(t, err) require.Equal(t, expectedURL, coderConfig.GetUrl())