feat: aibridged mcp handling (#19911)

If you have used AI to produce some or all of this PR, please ensure you have read our [AI Contribution guidelines](https://coder.com/docs/about/contributing/AI_CONTRIBUTING) before submitting.

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

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,

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.

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

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)
+			}
+		})
+	}
+}

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

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

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
 }

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())