feat: backend for per user headers mcp auth added

Author: Pratham-Mishra04

core/bifrost.go

@@ -233,7 +233,7 @@ func Init(ctx context.Context, config schemas.BifrostConfig) (*Bifrost, error) {
 		requestQueues: sync.Map{},
 		waitGroups:    sync.Map{},
 		keySelector:   config.KeySelector,
-		mcpCredStore:  credstore.NewCredStore(config.OAuth2Provider, config.Logger),
+		mcpCredStore:  credstore.NewCredStore(config.OAuth2Provider, config.MCPHeadersProvider, config.Logger),
 		logger:        config.Logger,
 		kvStore:       config.KVStore,
 	}
@@ -3824,6 +3824,34 @@ func (bifrost *Bifrost) VerifyPerUserOAuthConnection(ctx context.Context, config
 	return bifrost.MCPManager.VerifyPerUserOAuthConnection(ctx, config, accessToken)
 }
 
+// VerifyHeadersConnection delegates to the MCP manager to verify an MCP
+// server using caller-supplied header values (admin sample or user-submitted)
+// and discover available tools. Mirrors VerifyPerUserOAuthConnection's lazy
+// MCP-manager init.
+func (bifrost *Bifrost) VerifyHeadersConnection(ctx context.Context, config *schemas.MCPClientConfig, userHeaders map[string]string) (map[string]schemas.ChatTool, map[string]string, error) {
+	if bifrost.MCPManager == nil {
+		bifrost.mcpInitOnce.Do(func() {
+			mcpConfig := schemas.MCPConfig{
+				ClientConfigs: []*schemas.MCPClientConfig{},
+			}
+			mcpConfig.PluginPipelineProvider = func() interface{} {
+				return bifrost.getPluginPipeline()
+			}
+			mcpConfig.ReleasePluginPipeline = func(pipeline interface{}) {
+				if pp, ok := pipeline.(*PluginPipeline); ok {
+					bifrost.releasePluginPipeline(pp)
+				}
+			}
+			codeMode := starlark.NewStarlarkCodeMode(nil, bifrost.logger)
+			bifrost.MCPManager = mcp.NewMCPManager(bifrost.ctx, mcpConfig, bifrost.mcpCredStore, bifrost.logger, codeMode)
+		})
+	}
+	if bifrost.MCPManager == nil {
+		return nil, nil, fmt.Errorf("MCP manager is not initialized")
+	}
+	return bifrost.MCPManager.VerifyHeadersConnection(ctx, config, userHeaders)
+}
+
 // SetClientTools delegates to the MCP manager to update the tool map for an
 // existing MCP client.
 func (bifrost *Bifrost) SetClientTools(clientID string, tools map[string]schemas.ChatTool, toolNameMapping map[string]string) {

core/mcp/agent.go

@@ -286,7 +286,7 @@ func (a *AgentModeExecutor) executeAgent(
 			wg := sync.WaitGroup{}
 			wg.Add(len(autoExecutableTools))
 			channelToolResults := make(chan *schemas.ChatMessage, len(autoExecutableTools))
-			var authRequiredErr *schemas.MCPUserOAuthRequiredError
+			var authRequiredErr *schemas.MCPAuthRequiredError
 			var authRequiredOnce sync.Once
 			for _, toolCall := range autoExecutableTools {
 				go func(toolCall schemas.ChatAssistantMessageToolCall) {
@@ -304,11 +304,11 @@ func (a *AgentModeExecutor) executeAgent(
 
 					mcpResponse, toolErr := executeToolFunc(toolCtx, mcpRequest)
 					if toolErr != nil {
-						// Check if this is a per-user OAuth auth-required error
-						var oauthErr *schemas.MCPUserOAuthRequiredError
-						if errors.As(toolErr, &oauthErr) {
+						// Check if this is a per-user auth-required error
+						var authErr *schemas.MCPAuthRequiredError
+						if errors.As(toolErr, &authErr) {
 							authRequiredOnce.Do(func() {
-								authRequiredErr = oauthErr
+								authRequiredErr = authErr
 							})
 							channelToolResults <- createToolResultMessage(toolCall, "", toolErr)
 							return

core/mcp/clientmanager.go

@@ -66,18 +66,20 @@ func (m *MCPManager) AcquireClientConn(ctx *schemas.BifrostContext, state *schem
 	// Closure-captured outputs from the op so the caller can CallTool on the
 	// live client after the gate returns.
 	var tempClient *client.Client
-	// MCPUserOAuthRequiredError is wrapped into a generic BifrostError by the
+	// MCPAuthRequiredError is wrapped into a generic BifrostError by the
 	// pipeline before PostConnectionHook runs, so capture it out-of-band to
-	// preserve the typed-error info for the envelope path.
-	var oauthErr *schemas.MCPUserOAuthRequiredError
+	// preserve the typed-error info for the envelope path. Same capture
+	// covers both per-user-OAuth (Kind=oauth) and per-user-headers
+	// (Kind=headers) surfaces.
+	var authRequiredErr *schemas.MCPAuthRequiredError
 	start := time.Now()
 
 	_, gateErr := m.runConnectWithPluginPipeline(ctx, connectReq, func(preReq *schemas.BifrostMCPConnectRequest) (*schemas.BifrostMCPConnectResponse, error) {
 		// Resolve auth headers AFTER PreConnectionHook ran. Plugins never see
 		// the Authorization header — it lives only on the wire transport.
 		authHeaders, credErr := m.credStore.ConnectionHeaders(ctx, config)
 		if credErr != nil {
-			errors.As(credErr, &oauthErr)
+			errors.As(credErr, &authRequiredErr)
 			return nil, credErr
 		}
 
@@ -158,8 +160,8 @@ func (m *MCPManager) AcquireClientConn(ctx *schemas.BifrostContext, state *schem
 		if tempClient != nil {
 			_ = tempClient.Close()
 		}
-		if oauthErr != nil {
-			return nil, nil, oauthErr
+		if authRequiredErr != nil {
+			return nil, nil, authRequiredErr
 		}
 		if gateErr.Error != nil {
 			return nil, nil, fmt.Errorf("%s", gateErr.Error.Message)
@@ -326,17 +328,20 @@ func (m *MCPManager) AddClient(config *schemas.MCPClientConfig) error {
 				url := config.ConnectionString.GetValue()
 				client.ConnectionInfo.ConnectionURL = &url
 			}
-			// Restore discovered tools from config (persisted in DB across restarts)
+			// Restore discovered tools from config (persisted in DB across restarts).
+			// Applies to every per-call-connection auth type — currently per-user
+			// OAuth and per-user headers — since both populate DiscoveredTools at
+			// admin-test time and never hold a persistent client.Conn.
 			if len(config.DiscoveredTools) > 0 {
 				for toolName, tool := range config.DiscoveredTools {
 					client.ToolMap[toolName] = tool
 				}
 				client.ToolNameMapping = config.DiscoveredToolNameMapping
 				client.State = schemas.MCPConnectionStateConnected
-				m.logger.Debug("%s Per-user OAuth MCP client '%s' restored with %d tools", MCPLogPrefix, config.Name, len(config.DiscoveredTools))
+				m.logger.Debug("%s Per-user (%s) MCP client '%s' restored with %d tools", MCPLogPrefix, config.AuthType, config.Name, len(config.DiscoveredTools))
 			} else {
 				client.State = schemas.MCPConnectionStatePendingTools
-				m.logger.Debug("%s Per-user OAuth MCP client '%s' registered (connection deferred to runtime)", MCPLogPrefix, config.Name)
+				m.logger.Debug("%s Per-user (%s) MCP client '%s' registered (connection deferred to runtime)", MCPLogPrefix, config.AuthType, config.Name)
 			}
 		}
 		m.mu.Unlock()
@@ -492,6 +497,145 @@ func (m *MCPManager) VerifyPerUserOAuthConnection(ctx context.Context, config *s
 	return tools, toolNameMapping, nil
 }
 
+// VerifyHeadersConnection creates a temporary MCP connection using the
+// provided user-submitted header values to verify the server is reachable
+// and discover available tools. The connection is closed after verification.
+//
+// Used in two paths:
+//   - Admin test flow: admin enters sample values during MCP client creation,
+//     this runs an Initialize handshake against the upstream to validate the
+//     schema (PerUserHeaderKeys) + discover tools. The discovered tools then
+//     persist on the MCPClient row; the sample values are discarded.
+//   - User submission flow: an end user submits their own values via the
+//     workspace submit URL surfaced inline by MCPAuthRequiredError. The
+//     handler runs this before upserting the row so a bad submission returns
+//     422 immediately instead of failing on the next tool call.
+//
+// Parameters:
+//   - config: MCP client configuration (connection URL, name, PerUserHeaderKeys, etc.)
+//   - userHeaders: caller-supplied header_name → value map (must cover every
+//     PerUserHeaderKeys entry; the caller validates that before invoking).
+//
+// Returns:
+//   - map[string]schemas.ChatTool: discovered tools keyed by prefixed name
+//   - map[string]string: tool name mapping (sanitized → original MCP name)
+//   - error: any error during verification
+func (m *MCPManager) VerifyHeadersConnection(ctx context.Context, config *schemas.MCPClientConfig, userHeaders map[string]string) (map[string]schemas.ChatTool, map[string]string, error) {
+	if config.ConnectionString == nil || config.ConnectionString.GetValue() == "" {
+		return nil, nil, fmt.Errorf("connection URL is required for per-user headers verification")
+	}
+	if len(userHeaders) == 0 {
+		return nil, nil, fmt.Errorf("user headers are required for per-user headers verification")
+	}
+
+	// Build prepared inputs for the typed connect plugin gate. Static admin
+	// headers (minus Authorization and minus any PerUserHeaderKeys) are
+	// plugin-visible; user-supplied credentials are layered AFTER PreHooks
+	// run so plugins cannot read or rewrite them. Mirrors
+	// VerifyPerUserOAuthConnection's Authorization-injection pattern.
+	url := config.ConnectionString.GetValue()
+	preparedHeaders := utils.FlattenHeaders(utils.StaticConfigHeaders(config))
+	connectReq := &schemas.BifrostMCPConnectRequest{
+		ClientName:       config.Name,
+		ConnectionType:   schemas.MCPConnectionTypeHTTP,
+		AuthType:         config.AuthType,
+		ConnectionString: &url,
+		Headers:          preparedHeaders,
+	}
+
+	verifyCtx, cancel := context.WithTimeout(ctx, MCPClientConnectionEstablishTimeout)
+	defer cancel()
+	gateCtx := schemas.NewBifrostContext(verifyCtx, schemas.NoDeadline)
+
+	var tempClient *client.Client
+	defer func() {
+		if tempClient != nil {
+			tempClient.Close()
+		}
+	}()
+	start := time.Now()
+
+	_, gateErr := m.runConnectWithPluginPipeline(gateCtx, connectReq, func(preReq *schemas.BifrostMCPConnectRequest) (*schemas.BifrostMCPConnectResponse, error) {
+		finalURL := url
+		if preReq.ConnectionString != nil {
+			finalURL = *preReq.ConnectionString
+		}
+
+		// Copy mutated headers, then layer the user's credential values on
+		// top. Copying (rather than mutating preReq.Headers in place) avoids
+		// leaking the values back into the request that PreHook plugins may
+		// still reference.
+		finalHeaders := make(map[string]string, len(preReq.Headers)+len(userHeaders))
+		maps.Copy(finalHeaders, preReq.Headers)
+		for k, v := range userHeaders {
+			finalHeaders[k] = v
+		}
+
+		httpTransport, hErr := transport.NewStreamableHTTP(finalURL, transport.WithHTTPHeaders(finalHeaders))
+		if hErr != nil {
+			return nil, fmt.Errorf("failed to create HTTP transport for verification: %w", hErr)
+		}
+		tempClient = client.NewClient(httpTransport)
+		if startErr := tempClient.Start(verifyCtx); startErr != nil {
+			return nil, fmt.Errorf("failed to start MCP connection for verification: %w", startErr)
+		}
+
+		initRequest := mcp.InitializeRequest{
+			Params: mcp.InitializeParams{
+				ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION,
+				Capabilities:    mcp.ClientCapabilities{},
+				ClientInfo: mcp.Implementation{
+					Name:    fmt.Sprintf("Bifrost-%s-verify", config.Name),
+					Version: "1.0.0",
+				},
+			},
+		}
+		initResult, initErr := tempClient.Initialize(verifyCtx, initRequest)
+		if initErr != nil {
+			return nil, fmt.Errorf("failed to initialize MCP connection for verification: %w", initErr)
+		}
+
+		resp := &schemas.BifrostMCPConnectResponse{
+			ConnectionInfo: &schemas.MCPClientConnectionInfo{
+				Type:          schemas.MCPConnectionTypeHTTP,
+				ConnectionURL: &finalURL,
+			},
+			ExtraFields: schemas.BifrostMCPResponseExtraFields{
+				Latency: time.Since(start).Milliseconds(),
+			},
+		}
+		if initResult != nil {
+			resp.ProtocolVersion = initResult.ProtocolVersion
+			resp.ServerInfo = &schemas.MCPServerInfo{
+				Name:    initResult.ServerInfo.Name,
+				Version: initResult.ServerInfo.Version,
+			}
+			resp.ServerCapabilities = &schemas.MCPServerCapabilities{
+				Tools:     initResult.Capabilities.Tools != nil,
+				Resources: initResult.Capabilities.Resources != nil,
+				Prompts:   initResult.Capabilities.Prompts != nil,
+				Logging:   initResult.Capabilities.Logging != nil,
+			}
+		}
+		return resp, nil
+	})
+
+	if gateErr != nil {
+		return nil, nil, fmt.Errorf("failed to verify MCP connection: %s", gateErr.GetErrorString())
+	}
+	if tempClient == nil {
+		return nil, nil, fmt.Errorf("headers verification was short-circuited by plugin; cannot discover tools without a live connection")
+	}
+
+	tools, toolNameMapping, err := m.runListToolsWithHooks(verifyCtx, tempClient, config.Name)
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to discover tools during verification: %w", err)
+	}
+
+	m.logger.Info("%s Per-user headers verification succeeded for '%s': discovered %d tools", MCPLogPrefix, config.Name, len(tools))
+	return tools, toolNameMapping, nil
+}
+
 // SetClientTools updates the tool map and name mapping for an existing client.
 // This is used to populate tools discovered during per-user OAuth verification,
 // where tool discovery happens separately from client creation.
@@ -798,6 +942,7 @@ func (m *MCPManager) UpdateClient(id string, updatedConfig *schemas.MCPClientCon
 		ToolSyncInterval:      updatedConfig.ToolSyncInterval,
 		AllowOnAllVirtualKeys: updatedConfig.AllowOnAllVirtualKeys,
 		Disabled:              updatedConfig.Disabled,
+		PerUserHeaderKeys:     slices.Clone(updatedConfig.PerUserHeaderKeys),
 	}
 
 	// Atomically replace the config pointer

core/mcp/codemode/starlark/executecode.go

@@ -528,7 +528,7 @@ func (s *StarlarkCodeMode) callMCPTool(ctx *schemas.BifrostContext, clientName,
 	// Acquire a connection through the shared ClientManager abstraction:
 	// shared-mode clients return their persistent state.Conn (release is a
 	// no-op); per-user clients get a fresh ephemeral transport that the
-	// release function closes. Credential errors (e.g. MCPUserOAuthRequiredError)
+	// release function closes. Credential errors (e.g. MCPAuthRequiredError)
 	// surface here.
 	conn, release, err := s.clientManager.AcquireClientConn(nestedCtx, client)
 	if err != nil {

core/mcp/credstore/credstore.go

@@ -31,14 +31,18 @@ type CredStore struct {
 
 // NewCredStore constructs the canonical MCPCredentialStore with one resolver
 // per known MCPAuthType. The oauth2Provider is injected into the OAuth-
-// flavored resolvers only; the None and StaticHeaders resolvers are stateless.
-func NewCredStore(oauth2Provider schemas.OAuth2Provider, logger schemas.Logger) *CredStore {
+// flavored resolvers only; the None and StaticHeaders resolvers are
+// stateless. The headersProvider is injected into the per-user-headers
+// resolver — pass nil if the configstore-backed provider isn't wired up
+// (the resolver returns a clear error rather than nil-pointering at use).
+func NewCredStore(oauth2Provider schemas.OAuth2Provider, headersProvider schemas.MCPHeadersProvider, logger schemas.Logger) *CredStore {
 	return &CredStore{
 		resolvers: map[schemas.MCPAuthType]resolver{
-			schemas.MCPAuthTypeNone:         &noneResolver{},
-			schemas.MCPAuthTypeHeaders:      &staticHeadersResolver{},
-			schemas.MCPAuthTypeOauth:        &serverOAuthResolver{provider: oauth2Provider},
-			schemas.MCPAuthTypePerUserOauth: &perUserOAuthResolver{provider: oauth2Provider},
+			schemas.MCPAuthTypeNone:           &noneResolver{},
+			schemas.MCPAuthTypeHeaders:        &staticHeadersResolver{},
+			schemas.MCPAuthTypeOauth:          &serverOAuthResolver{provider: oauth2Provider},
+			schemas.MCPAuthTypePerUserOauth:   &perUserOAuthResolver{provider: oauth2Provider},
+			schemas.MCPAuthTypePerUserHeaders: &perUserHeadersResolver{provider: headersProvider},
 		},
 		logger: logger,
 	}