diff --git a/core/bifrost.go b/core/bifrost.go index 47d5cd59f8..27b1d55169 100644 --- a/core/bifrost.go +++ b/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) { diff --git a/core/mcp/agent.go b/core/mcp/agent.go index 59930bf612..28a4f4bb59 100644 --- a/core/mcp/agent.go +++ b/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 diff --git a/core/mcp/clientmanager.go b/core/mcp/clientmanager.go index 535a1201b3..b08fd38a77 100644 --- a/core/mcp/clientmanager.go +++ b/core/mcp/clientmanager.go @@ -66,10 +66,12 @@ 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) { @@ -77,7 +79,7 @@ func (m *MCPManager) AcquireClientConn(ctx *schemas.BifrostContext, state *schem // 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 diff --git a/core/mcp/codemode/starlark/executecode.go b/core/mcp/codemode/starlark/executecode.go index 505e71c492..560c668918 100644 --- a/core/mcp/codemode/starlark/executecode.go +++ b/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 { diff --git a/core/mcp/credstore/credstore.go b/core/mcp/credstore/credstore.go index ea2319cff6..194d0b2d6a 100644 --- a/core/mcp/credstore/credstore.go +++ b/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, } diff --git a/core/mcp/credstore/per_user_headers.go b/core/mcp/credstore/per_user_headers.go new file mode 100644 index 0000000000..a40b3ba442 --- /dev/null +++ b/core/mcp/credstore/per_user_headers.go @@ -0,0 +1,153 @@ +package credstore + +import ( + "errors" + "fmt" + "net/http" + + "github.com/maximhq/bifrost/core/mcp/utils" + "github.com/maximhq/bifrost/core/schemas" +) + +// perUserHeadersResolver handles MCPAuthTypePerUserHeaders — each caller's +// upstream API-key / signed-token headers are keyed by (auth_mode, identity, +// mcp_client) in the mcp_per_user_header_credentials table. On miss or stale +// schema, an inline submission flow is initiated and a *MCPAuthRequiredError +// with Kind="headers" is raised so the caller can complete the submission UI. +// +// ConnectionHeaders returns the user-submitted header values. Static admin +// headers are layered separately by AcquireClientConn via +// utils.StaticConfigHeaders (which excludes anything in +// config.PerUserHeaderKeys) so the connect-plugin gate never observes the +// caller's secret values. +// +// RequiresPerCallConnection is true: per-user-headers clients never hold a +// persistent upstream connection; AcquireClientConn opens a fresh ephemeral +// HTTP transport per call using the resolved user headers + plugin-mutated +// static headers. +type perUserHeadersResolver struct { + provider schemas.MCPHeadersProvider +} + +func (r *perUserHeadersResolver) ConnectionHeaders(ctx *schemas.BifrostContext, config *schemas.MCPClientConfig) (http.Header, error) { + if r.provider == nil { + return nil, fmt.Errorf("per-user headers requires an MCPHeadersProvider but none is configured") + } + if len(config.PerUserHeaderKeys) == 0 { + return nil, fmt.Errorf("per-user headers client %q has no PerUserHeaderKeys declared (admin config error)", config.Name) + } + + mode := ctx.MCPAuthMode() + identity := identityForMCPAuthMode(ctx, mode) + if identity == "" { + return nil, fmt.Errorf( + "per-user headers for %s requires an identity: send a Virtual Key (x-bf-vk), authenticate as a user, or set x-bf-mcp-session-id to any opaque string you'll re-send on subsequent calls", + config.Name, + ) + } + + cred, err := r.provider.GetCredentialByMode(ctx, mode, identity, config.ID) + switch { + case err == nil: + // Row present: intersect stored values with the current schema. If any + // required key is missing on the stored row, the schema has drifted + // since the user last submitted — surface the same submit-URL flow as + // "not found" but the underlying row stays (so the UI can prefill + // known values when the user resubmits). + if missing := missingRequiredHeaderKeys(config.PerUserHeaderKeys, cred.Headers); len(missing) > 0 { + return nil, r.buildAuthRequiredError(ctx, config) + } + return buildPerUserHeaderValues(config.PerUserHeaderKeys, cred.Headers), nil + case errors.Is(err, schemas.ErrHeadersCredentialNotFound), + errors.Is(err, schemas.ErrHeadersCredentialNeedsUpdate): + return nil, r.buildAuthRequiredError(ctx, config) + default: + return nil, fmt.Errorf("failed to load per-user header credential for %s: %w", config.Name, err) + } +} + +func (r *perUserHeadersResolver) RequiresPerCallConnection() bool { return true } + +// buildAuthRequiredError creates a pending mcp_per_user_header_flows row +// via the provider, then constructs the inline-401 payload pointing at +// that flow's auth-page URL. Mirrors per_user_oauth.go's call to +// InitiateUserOAuthFlow: the provider mints a temp-token bound to the +// flow ID and embeds it as a `#t=` URL fragment so anonymous +// browser visitors can complete the submission without a dashboard +// session. +func (r *perUserHeadersResolver) buildAuthRequiredError(ctx *schemas.BifrostContext, config *schemas.MCPClientConfig) error { + baseURL := utils.BuildMCPCallbackBaseURL(ctx) + if baseURL == "" { + return fmt.Errorf("per-user headers requires a callback base URL but none is available in context") + } + mode := ctx.MCPAuthMode() + identity := identityForMCPAuthMode(ctx, mode) + if identity == "" { + // Defensive — the caller already validated identity before invoking + // the auth-required path, but keep the guard so a future refactor + // can't accidentally start flow rows with empty identity columns. + return fmt.Errorf("per-user headers auth-required flow requires an identity") + } + initiation, err := r.provider.InitiateUserSubmissionFlow(ctx, mode, identity, config.ID, baseURL) + if err != nil { + return fmt.Errorf("failed to initiate per-user headers submission flow for %s: %w", config.Name, err) + } + return &schemas.MCPAuthRequiredError{ + Kind: schemas.MCPAuthRequiredKindHeaders, + MCPClientID: config.ID, + MCPClientName: config.Name, + SubmitURL: initiation.FrontendURL, + SessionID: initiation.FlowID, + RequiredHeaderKeys: append([]string(nil), config.PerUserHeaderKeys...), + AdminHeaderKeys: adminHeaderKeyNames(config), + // Include the URL in the message so plain-text clients (curl, basic + // SDK wrappers) that don't parse extra_fields still get an actionable + // hint. Matches per_user_oauth.go's behavior. + Message: fmt.Sprintf("Authentication required for %s. Visit %s to submit the required headers.", config.Name, initiation.FrontendURL), + } +} + +// missingRequiredHeaderKeys returns the names of any required header key +// that's absent or whose stored value is empty in storedHeaders. Comparison +// is case-insensitive at the wire level but the schema is the source of +// truth — we look up by the exact key the admin declared. +func missingRequiredHeaderKeys(required []string, storedHeaders map[string]string) []string { + if len(storedHeaders) == 0 { + return append([]string(nil), required...) + } + var missing []string + for _, key := range required { + if v, ok := storedHeaders[key]; !ok || v == "" { + missing = append(missing, key) + } + } + return missing +} + +// buildPerUserHeaderValues constructs the http.Header carrying just the +// user-submitted credential values for the required keys. Keys not declared +// by the current schema are dropped on purpose so a stale row that still +// stores a deprecated key cannot leak it onto the wire. +func buildPerUserHeaderValues(required []string, storedHeaders map[string]string) http.Header { + out := http.Header{} + for _, key := range required { + if v, ok := storedHeaders[key]; ok && v != "" { + out.Set(key, v) + } + } + return out +} + +// adminHeaderKeyNames returns the names (no values) of static admin headers +// declared on the MCP client. Surfaced to the submission UI so the user can +// see what context will accompany their request without exposing the values. +func adminHeaderKeyNames(config *schemas.MCPClientConfig) []string { + if config == nil || len(config.Headers) == 0 { + return nil + } + names := make([]string, 0, len(config.Headers)) + for name := range config.Headers { + names = append(names, name) + } + return names +} diff --git a/core/mcp/credstore/per_user_oauth.go b/core/mcp/credstore/per_user_oauth.go index f454ede228..41ddb6bfd8 100644 --- a/core/mcp/credstore/per_user_oauth.go +++ b/core/mcp/credstore/per_user_oauth.go @@ -53,7 +53,7 @@ func (r *perUserOAuthResolver) ConnectionHeaders(ctx *schemas.BifrostContext, co if config.OauthConfigID == nil || *config.OauthConfigID == "" { return nil, fmt.Errorf("per-user OAuth requires an OAuth config but MCP client %s has none", config.Name) } - redirectURI := utils.BuildRedirectURIFromContext(ctx) + redirectURI := utils.BuildOAuthRedirectURIFromContext(ctx) if redirectURI == "" { return nil, fmt.Errorf("per-user OAuth requires a redirect URI but none is available in context") } @@ -61,7 +61,8 @@ func (r *perUserOAuthResolver) ConnectionHeaders(ctx *schemas.BifrostContext, co if flowErr != nil { return nil, fmt.Errorf("failed to initiate per-user OAuth flow for %s: %w", config.Name, flowErr) } - return nil, &schemas.MCPUserOAuthRequiredError{ + return nil, &schemas.MCPAuthRequiredError{ + Kind: schemas.MCPAuthRequiredKindOAuth, MCPClientID: config.ID, MCPClientName: config.Name, AuthorizeURL: flowInitiation.AuthorizeURL, @@ -82,23 +83,3 @@ func (r *perUserOAuthResolver) ConnectionHeaders(ctx *schemas.BifrostContext, co } func (r *perUserOAuthResolver) RequiresPerCallConnection() bool { return true } - -// identityForMCPAuthMode returns the identity string to look up by, given the -// derived mode. Mirrors the priority used by ctx.MCPAuthMode(). -func identityForMCPAuthMode(ctx *schemas.BifrostContext, mode schemas.MCPAuthMode) string { - switch mode { - case schemas.MCPAuthModeUser: - if v, _ := ctx.Value(schemas.BifrostContextKeyUserID).(string); v != "" { - return v - } - case schemas.MCPAuthModeVK: - if v, _ := ctx.Value(schemas.BifrostContextKeyGovernanceVirtualKeyID).(string); v != "" { - return v - } - case schemas.MCPAuthModeSession: - if v, _ := ctx.Value(schemas.BifrostContextKeyMCPSessionID).(string); v != "" { - return v - } - } - return "" -} diff --git a/core/mcp/credstore/utils.go b/core/mcp/credstore/utils.go new file mode 100644 index 0000000000..9501d10c79 --- /dev/null +++ b/core/mcp/credstore/utils.go @@ -0,0 +1,28 @@ +package credstore + +import "github.com/maximhq/bifrost/core/schemas" + +// identityForMCPAuthMode returns the identity string to look up by, given the +// derived mode. Mirrors the priority used by ctx.MCPAuthMode(). +// +// Used by every resolver that keys persisted state by (mode, identity, +// mcp_client) — currently per-user OAuth and per-user headers. Lives in its +// own file so both resolvers can call it without duplication or accidental +// drift. +func identityForMCPAuthMode(ctx *schemas.BifrostContext, mode schemas.MCPAuthMode) string { + switch mode { + case schemas.MCPAuthModeUser: + if v, _ := ctx.Value(schemas.BifrostContextKeyUserID).(string); v != "" { + return v + } + case schemas.MCPAuthModeVK: + if v, _ := ctx.Value(schemas.BifrostContextKeyGovernanceVirtualKeyID).(string); v != "" { + return v + } + case schemas.MCPAuthModeSession: + if v, _ := ctx.Value(schemas.BifrostContextKeyMCPSessionID).(string); v != "" { + return v + } + } + return "" +} diff --git a/core/mcp/exec.go b/core/mcp/exec.go index 4e14de9204..fe5a863cb1 100644 --- a/core/mcp/exec.go +++ b/core/mcp/exec.go @@ -85,8 +85,9 @@ func (m *MCPManager) executeToolWithHooks( // Resolve the upstream client and acquire its connection BEFORE the plugin // gate runs. Connection lifecycle is the orchestrator's concern, not the // plugin op's — the plugin pipeline only wraps the actual CallTool. When - // AcquireClientConn fails (e.g. *MCPUserOAuthRequiredError for per-user - // clients that need re-auth), the plugin gate is never invoked. + // AcquireClientConn fails (e.g. *MCPAuthRequiredError for per-user + // clients that need re-auth or headers submission), the plugin gate is + // never invoked. state, conn, release, prepErr := m.prepareToolExecution(ctx, request) if prepErr != nil { bErr := &schemas.BifrostError{ @@ -94,9 +95,9 @@ func (m *MCPManager) executeToolWithHooks( Error: &schemas.ErrorField{Message: prepErr.Error()}, ExtraFields: schemas.BifrostErrorExtraFields{RequestType: requestType, MCPRequestType: request.RequestType}, } - var oauthErr *schemas.MCPUserOAuthRequiredError - if errors.As(prepErr, &oauthErr) { - bErr.ExtraFields.MCPAuthRequired = oauthErr + var authRequiredErr *schemas.MCPAuthRequiredError + if errors.As(prepErr, &authRequiredErr) { + bErr.ExtraFields.MCPAuthRequired = authRequiredErr } return nil, bErr } diff --git a/core/mcp/interface.go b/core/mcp/interface.go index 309bbc2a3b..a180094fa0 100644 --- a/core/mcp/interface.go +++ b/core/mcp/interface.go @@ -76,6 +76,11 @@ type MCPManagerInterface interface { // EnableClient reconnects a disabled client and restarts its workers EnableClient(id string) error + // VerifyHeadersConnection creates a temporary MCP connection using a set of + // caller-supplied header values to verify connectivity and discover tools. + // The connection is closed after verification. + VerifyHeadersConnection(ctx context.Context, config *schemas.MCPClientConfig, userHeaders map[string]string) (map[string]schemas.ChatTool, map[string]string, error) + // VerifyPerUserOAuthConnection creates a temporary MCP connection using a // test access token to verify connectivity and discover tools. The connection // is closed after verification. diff --git a/core/mcp/mcp.go b/core/mcp/mcp.go index c727e82b2e..cbb463d205 100644 --- a/core/mcp/mcp.go +++ b/core/mcp/mcp.go @@ -79,11 +79,12 @@ func NewMCPManager(ctx context.Context, config schemas.MCPConfig, credStore sche if logger == nil { logger = defaultLogger } - // Default to an OAuth-less CredentialStore so tests (and callers that - // don't wire OAuth) get a working store: static/headers/none resolvers - // stay functional, OAuth-flavored resolvers cleanly error on use. + // Default to a provider-less CredentialStore so tests (and callers that + // don't wire OAuth / per-user-headers) get a working store: static / + // headers / none resolvers stay functional, per-user resolvers cleanly + // error on use. if credStore == nil { - credStore = credstore.NewCredStore(nil, logger) + credStore = credstore.NewCredStore(nil, nil, logger) } // Set default values if config.ToolManagerConfig == nil { diff --git a/core/mcp/toolmanager.go b/core/mcp/toolmanager.go index ccc3895bc7..2831231040 100644 --- a/core/mcp/toolmanager.go +++ b/core/mcp/toolmanager.go @@ -152,7 +152,7 @@ func NewToolsManagerWithCodeMode( // transparently for None / StaticHeaders auth and surfaces a clear // "OAuth2 provider not available" error for OAuth-flavored clients. if credStore == nil { - credStore = credstore.NewCredStore(nil, logger) + credStore = credstore.NewCredStore(nil, nil, logger) } agentModeExecutor := &AgentModeExecutor{ diff --git a/core/mcp/utils.go b/core/mcp/utils.go index a6249e49bb..b281b079ae 100644 --- a/core/mcp/utils.go +++ b/core/mcp/utils.go @@ -636,6 +636,19 @@ func validateMCPClientConfig(config *schemas.MCPClientConfig) error { default: return fmt.Errorf("unknown connection type '%s' in client '%s'", config.ConnectionType, config.Name) } + if config.AuthType == schemas.MCPAuthTypePerUserHeaders { + if len(config.PerUserHeaderKeys) == 0 { + return fmt.Errorf("per_user_header_keys is required (non-empty) for per_user_headers auth type in client '%s'", config.Name) + } + for i, key := range config.PerUserHeaderKeys { + if strings.TrimSpace(key) == "" { + return fmt.Errorf("per_user_header_keys[%d] is empty in client '%s'", i, config.Name) + } + } + if config.OauthConfigID != nil && *config.OauthConfigID != "" { + return fmt.Errorf("oauth_config_id must not be set for per_user_headers auth type in client '%s'", config.Name) + } + } return nil } diff --git a/core/mcp/utils/utils.go b/core/mcp/utils/utils.go index ec53c65e4c..5605c556d8 100644 --- a/core/mcp/utils/utils.go +++ b/core/mcp/utils/utils.go @@ -22,24 +22,46 @@ func FlattenHeaders(h http.Header) map[string]string { return out } -// BuildRedirectURIFromContext extracts the OAuth redirect URI from context. -func BuildRedirectURIFromContext(ctx *schemas.BifrostContext) string { - if uri, ok := ctx.Value(schemas.BifrostContextKeyOAuthRedirectURI).(string); ok && uri != "" { - return uri +// BuildMCPCallbackBaseURL extracts the base URL set on the BifrostContext by +// the HTTP middleware (e.g. "https://host"). Per-user OAuth and per-user +// headers resolvers append their respective paths on top. +func BuildMCPCallbackBaseURL(ctx *schemas.BifrostContext) string { + if base, ok := ctx.Value(schemas.BifrostContextKeyMCPCallbackBaseURL).(string); ok && base != "" { + return base } return "" } +// BuildOAuthRedirectURIFromContext returns the full OAuth callback URL +// ("/api/oauth/callback") needed by the per-user OAuth flow, or empty +// if the base URL is unavailable. +func BuildOAuthRedirectURIFromContext(ctx *schemas.BifrostContext) string { + base := BuildMCPCallbackBaseURL(ctx) + if base == "" { + return "" + } + return base + "/api/oauth/callback" +} + // StaticConfigHeaders returns the admin-configured static headers from -// config.Headers MINUS any Authorization header. These are the headers that -// are safe to expose to MCP connect-plugins via the PreConnectionHook gate — -// plugins may add, remove, or rewrite them. +// config.Headers MINUS any header whose name is a credential — Authorization +// always, plus any name declared in config.PerUserHeaderKeys. These are the +// headers that are safe to expose to MCP connect-plugins via the +// PreConnectionHook gate — plugins may add, remove, or rewrite them. // -// Authorization is excluded by design even when an admin sets it manually -// in config.Headers (e.g. for MCPAuthTypeHeaders with a hard-coded bearer): -// it is a credential, and credentials are layered AFTER the plugin gate -// runs. The CredentialStore resolver for the relevant auth type emits the -// final Authorization value (either from config or from a dynamic token). +// Why exclude: +// - Authorization: credential by definition. The CredentialStore resolver +// for the active auth type emits the final value (config bearer for +// MCPAuthTypeHeaders; dynamic token for OAuth-flavored types). +// - PerUserHeaderKeys: credential schema for MCPAuthTypePerUserHeaders. If +// an admin accidentally (or deliberately) baked one of these names into +// config.Headers with a static value, exposing it to plugins would leak +// the static fallback. The per-user-headers resolver emits the caller's +// value; the static fallback should never reach the wire (and never +// reach plugins) for per-user-headers clients. +// +// Comparison is case-insensitive because HTTP headers are case-insensitive +// on the wire but case-sensitive in Go maps. func StaticConfigHeaders(config *schemas.MCPClientConfig) http.Header { headers := make(http.Header) if config == nil { @@ -49,11 +71,26 @@ func StaticConfigHeaders(config *schemas.MCPClientConfig) http.Header { if strings.EqualFold(key, "Authorization") { continue } + if matchesPerUserHeaderKey(key, config.PerUserHeaderKeys) { + continue + } headers.Add(key, value.GetValue()) } return headers } +// matchesPerUserHeaderKey reports whether name matches any entry in +// perUserKeys (case-insensitively). Used by StaticConfigHeaders to strip +// per-user credential keys from the plugin-visible static header set. +func matchesPerUserHeaderKey(name string, perUserKeys []string) bool { + for _, key := range perUserKeys { + if strings.EqualFold(name, key) { + return true + } + } + return false +} + // ExtractFilteredExtras returns just the per-request "extra" headers carried // in the BifrostContext (BifrostContextKeyMCPExtraHeaders), scoped by the // client's AllowedExtraHeaders. Static config headers are NOT included here — diff --git a/core/schemas/bifrost.go b/core/schemas/bifrost.go index 7be0acd810..c5e54b33d5 100644 --- a/core/schemas/bifrost.go +++ b/core/schemas/bifrost.go @@ -22,6 +22,7 @@ type BifrostConfig struct { LLMPlugins []LLMPlugin MCPPlugins []MCPPlugin OAuth2Provider OAuth2Provider + MCPHeadersProvider MCPHeadersProvider // Backend for MCPAuthTypePerUserHeaders credential storage; nil disables per-user-headers auth (resolver errors at use) Logger Logger Tracer Tracer // Tracer for distributed tracing (nil = NoOpTracer) InitialPoolSize int // Initial pool size for sync pools in Bifrost. Higher values will reduce memory allocations but will increase memory usage. @@ -244,7 +245,7 @@ const ( BifrostContextKeyTraceCompleter BifrostContextKey = "bifrost-trace-completer" // func([]PluginLogEntry) (callback to complete trace after streaming, receives transport plugin logs - set by tracing middleware) BifrostContextKeyAccumulatorID BifrostContextKey = "bifrost-accumulator-id" // string (ID for streaming accumulator lookup - set by tracer for accumulator operations) BifrostContextKeyMCPSessionID BifrostContextKey = "bifrost-mcp-session-id" // string (session-mode identity: any opaque value asserted by the caller via x-bf-mcp-session-id; binds the OAuth token row to subsequent /mcp calls when no VK or user is present) - BifrostContextKeyOAuthRedirectURI BifrostContextKey = "bifrost-oauth-redirect-uri" // string (OAuth callback URL, e.g. https://host/api/oauth/callback - set by HTTP middleware) + BifrostContextKeyMCPCallbackBaseURL BifrostContextKey = "bifrost-mcp-callback-base-url" // string (base URL like "https://host" — set by HTTP middleware. OAuth resolver appends /api/oauth/callback; headers resolver appends the workspace submit path. Used for both per-user OAuth and per-user headers auth flows) BifrostContextKeyIsMCPGateway BifrostContextKey = "bifrost-is-mcp-gateway" // bool (true when request is being handled via the MCP gateway path) BifrostContextKeyHasEmittedMessageDelta BifrostContextKey = "bifrost-has-emitted-message-delta" // bool (tracks whether message_delta was already emitted during streaming - avoids duplicates) BifrostContextKeySkipDBUpdate BifrostContextKey = "bifrost-skip-db-update" // bool (set by bifrost - DO NOT SET THIS MANUALLY)) @@ -1663,5 +1664,5 @@ type BifrostErrorExtraFields struct { ConvertedRequestType RequestType `json:"converted_request_type,omitempty"` DroppedCompatPluginParams []string `json:"dropped_compat_plugin_params,omitempty"` KeyStatuses []KeyStatus `json:"key_statuses,omitempty"` - MCPAuthRequired *MCPUserOAuthRequiredError `json:"mcp_auth_required,omitempty"` // Set when a per-user OAuth MCP tool requires authentication + MCPAuthRequired *MCPAuthRequiredError `json:"mcp_auth_required,omitempty"` // Set when a per-user MCP tool requires the caller to complete an inline auth flow (OAuth or headers) } diff --git a/core/schemas/mcp.go b/core/schemas/mcp.go index 00ba9b9361..ce601883b6 100644 --- a/core/schemas/mcp.go +++ b/core/schemas/mcp.go @@ -38,20 +38,59 @@ var ( ErrMCPReconnectNotApplicable = errors.New("reconnect is not applicable for this client type") ) -// MCPUserOAuthRequiredError is returned when a per-user OAuth MCP server requires -// the user to authenticate before tool execution can proceed. -type MCPUserOAuthRequiredError struct { +// MCPAuthRequiredKind discriminates the kind of inline-401 auth flow surfaced +// to the caller. The value lands in MCPAuthRequiredError.Kind and on the wire +// under extra_fields.mcp_auth_required.kind. +const ( + MCPAuthRequiredKindOAuth = "oauth" + MCPAuthRequiredKindHeaders = "headers" +) + +// MCPAuthRequiredError is returned when a per-user MCP credential is missing +// and the caller must complete an inline auth flow (OAuth dance or headers +// submission) before tool execution can proceed. +// +// Kind discriminates which set of fields is populated: +// - "oauth": AuthorizeURL, SessionID +// - "headers": SubmitURL, SessionID, RequiredHeaderKeys, AdminHeaderKeys +// +// SessionID is shared by both Kinds: for "oauth" it is the +// mcp_per_user_oauth_flows row ID, for "headers" the +// mcp_per_user_header_flows row ID. Either way it lets the caller +// reference the pending flow row without parsing the URL fragment. +// +// Common fields (MCPClientID, MCPClientName, Message) are always set. +type MCPAuthRequiredError struct { + Kind string `json:"kind"` MCPClientID string `json:"mcp_client_id"` MCPClientName string `json:"mcp_client_name"` - AuthorizeURL string `json:"authorize_url"` - SessionID string `json:"session_id"` Message string `json:"message"` + + // OAuth-specific fields (populated when Kind == "oauth"). SessionID is + // also populated for Kind == "headers" — see the type-level comment. + AuthorizeURL string `json:"authorize_url,omitempty"` + SessionID string `json:"session_id,omitempty"` + + // Headers-specific fields (populated when Kind == "headers"). SubmitURL is + // the workspace landing page where the user provides values for + // RequiredHeaderKeys; AdminHeaderKeys lists the admin-set static headers + // (names only, no values) for context display. + SubmitURL string `json:"submit_url,omitempty"` + RequiredHeaderKeys []string `json:"required_header_keys,omitempty"` + AdminHeaderKeys []string `json:"admin_header_keys,omitempty"` } -func (e *MCPUserOAuthRequiredError) Error() string { +func (e *MCPAuthRequiredError) Error() string { return e.Message } +// MCPUserOAuthRequiredError is an alias retained for backward compatibility +// with callers that referenced the OAuth-only error type before headers auth +// was added. New code should use MCPAuthRequiredError directly. +// +// Deprecated: use MCPAuthRequiredError. +type MCPUserOAuthRequiredError = MCPAuthRequiredError + // MCPCredentialStore is the single source of truth for MCP credential resolution. // It exposes three predicates that MCPManager consumes uniformly: // @@ -77,8 +116,9 @@ type MCPCredentialStore interface { // resolver returns the caller's full set (static + filtered // context-extras + per-user auth). // - // May return *MCPUserOAuthRequiredError when a per-user credential is - // missing and the caller must complete an auth flow before retrying. + // May return *MCPAuthRequiredError when a per-user credential is missing + // and the caller must complete an inline auth flow (OAuth dance or + // headers submission) before retrying. ConnectionHeaders(ctx *BifrostContext, config *MCPClientConfig) (http.Header, error) // RequestHeaders returns the per-message headers attached to each @@ -228,26 +268,35 @@ const ( type MCPAuthType string const ( - MCPAuthTypeNone MCPAuthType = "none" // No authentication - MCPAuthTypeHeaders MCPAuthType = "headers" // Header-based authentication (API keys, etc.) - MCPAuthTypeOauth MCPAuthType = "oauth" // OAuth 2.0 authentication (server-level, admin authenticates once) - MCPAuthTypePerUserOauth MCPAuthType = "per_user_oauth" // Per-user OAuth 2.0 authentication (each user authenticates individually) + MCPAuthTypeNone MCPAuthType = "none" // No authentication + MCPAuthTypeHeaders MCPAuthType = "headers" // Header-based authentication (API keys, etc.) + MCPAuthTypeOauth MCPAuthType = "oauth" // OAuth 2.0 authentication (server-level, admin authenticates once) + MCPAuthTypePerUserOauth MCPAuthType = "per_user_oauth" // Per-user OAuth 2.0 authentication (each user authenticates individually) + MCPAuthTypePerUserHeaders MCPAuthType = "per_user_headers" // Per-user header authentication (each user submits API keys / signed tokens; admin declares the required key names via PerUserHeaderKeys) ) // MCPClientConfig defines tool filtering for an MCP client. type MCPClientConfig struct { - ID string `json:"client_id"` // Client ID - Name string `json:"name"` // Client name - IsCodeModeClient bool `json:"is_code_mode_client"` // Whether the client is a code mode client - ConnectionType MCPConnectionType `json:"connection_type"` // How to connect (HTTP, STDIO, SSE, or InProcess) - ConnectionString *EnvVar `json:"connection_string,omitempty"` // HTTP or SSE URL (required for HTTP or SSE connections) - StdioConfig *MCPStdioConfig `json:"stdio_config,omitempty"` // STDIO configuration (required for STDIO connections) - AuthType MCPAuthType `json:"auth_type"` // Authentication type (none, headers, or oauth) - OauthConfigID *string `json:"oauth_config_id,omitempty"` // OAuth config ID (references oauth_configs table) - OauthClientID *EnvVar `json:"oauth_client_id,omitempty"` // Redacted OAuth client ID (populated on GET, not stored here) - OauthClientSecret *EnvVar `json:"oauth_client_secret,omitempty"` // Redacted OAuth client secret (populated on GET, not stored here) - State string `json:"state,omitempty"` // Connection state (connected, disconnected, error) - Headers map[string]EnvVar `json:"headers,omitempty"` // Headers to send with the request (for headers auth type) + ID string `json:"client_id"` // Client ID + Name string `json:"name"` // Client name + IsCodeModeClient bool `json:"is_code_mode_client"` // Whether the client is a code mode client + ConnectionType MCPConnectionType `json:"connection_type"` // How to connect (HTTP, STDIO, SSE, or InProcess) + ConnectionString *EnvVar `json:"connection_string,omitempty"` // HTTP or SSE URL (required for HTTP or SSE connections) + StdioConfig *MCPStdioConfig `json:"stdio_config,omitempty"` // STDIO configuration (required for STDIO connections) + AuthType MCPAuthType `json:"auth_type"` // Authentication type (none, headers, or oauth) + OauthConfigID *string `json:"oauth_config_id,omitempty"` // OAuth config ID (references oauth_configs table) + OauthClientID *EnvVar `json:"oauth_client_id,omitempty"` // Redacted OAuth client ID (populated on GET, not stored here) + OauthClientSecret *EnvVar `json:"oauth_client_secret,omitempty"` // Redacted OAuth client secret (populated on GET, not stored here) + State string `json:"state,omitempty"` // Connection state (connected, disconnected, error) + Headers map[string]EnvVar `json:"headers,omitempty"` // Headers to send with the request (for headers auth type) + // PerUserHeaderKeys lists the header *names* each caller must supply for + // MCPAuthTypePerUserHeaders clients. Admin-declared schema only — the + // values live per-user in the mcp_per_user_header_credentials table and + // are resolved at call time. Names in this list are stripped from + // utils.StaticConfigHeaders so admin-set values in `Headers` with the + // same name cannot leak through the plugin gate. Required (non-empty) + // when AuthType == per_user_headers; ignored otherwise. + PerUserHeaderKeys []string `json:"per_user_header_keys,omitempty"` AllowedExtraHeaders WhiteList `json:"allowed_extra_headers,omitempty"` // Allowlist of request-level headers that callers may forward to this MCP server at execution time InProcessServer *server.MCPServer `json:"-"` // MCP server instance for in-process connections (Go package only) ToolsToExecute WhiteList `json:"tools_to_execute,omitempty"` // Include-only list. diff --git a/core/schemas/mcp_headers.go b/core/schemas/mcp_headers.go new file mode 100644 index 0000000000..5785292da7 --- /dev/null +++ b/core/schemas/mcp_headers.go @@ -0,0 +1,108 @@ +//go:build !tinygo && !wasm + +package schemas + +import ( + "context" + "errors" + "time" +) + +// Per-user-headers errors. Mirrors the OAuth sentinels at the top of mcp.go; +// kept in this file so the headers feature surface is self-contained. +var ( + // ErrHeadersCredentialProviderNotAvailable signals that the headers + // provider isn't wired up — typically a misconfiguration (per_user_headers + // auth type used while running without a configstore-backed provider). + ErrHeadersCredentialProviderNotAvailable = errors.New("per-user headers credential provider not available") + + // ErrHeadersCredentialNotFound is the sentinel returned by + // MCPHeadersProvider.GetCredentialByMode when no row exists for the + // (mode, identity, mcp_client) triple. The resolver fans this out into an + // inline MCPAuthRequiredError so the caller can complete the submission + // flow. + ErrHeadersCredentialNotFound = errors.New("per-user headers credential not found for this identity and mcp client") + + // ErrHeadersCredentialNeedsUpdate signals that the stored credential is + // stale relative to the current MCPClientConfig.PerUserHeaderKeys schema + // (e.g. admin added a new required key). The resolver treats this like + // "not found" for inline-401 purposes but the row is preserved so the UI + // can prefill known values. + ErrHeadersCredentialNeedsUpdate = errors.New("per-user headers credential is missing keys required by the current schema") +) + +// MCPHeadersUserCredentialStatus mirrors the lifecycle states tracked on the +// mcp_per_user_header_credentials table. Storage-layer concerns; the resolver +// only cares about "is this row usable right now". +type MCPHeadersUserCredentialStatus string + +const ( + MCPHeadersUserCredentialStatusActive MCPHeadersUserCredentialStatus = "active" // Row matches the current schema and may be used + MCPHeadersUserCredentialStatusNeedsUpdate MCPHeadersUserCredentialStatus = "needs_update" // Schema (PerUserHeaderKeys) changed; user must resubmit + MCPHeadersUserCredentialStatusOrphaned MCPHeadersUserCredentialStatus = "orphaned" // Owner (VK / user) was deleted or detached; awaiting cleanup +) + +// MCPHeadersUserCredential is the in-memory view of a single per-user header +// credential row. The transport between core and framework treats Headers as +// plaintext — encryption at rest is the configstore's responsibility. +type MCPHeadersUserCredential struct { + ID string + MCPClientID string + AuthMode MCPAuthMode + UserID *string + VirtualKeyID *string + SessionID *string + Headers map[string]string // Decrypted header values + Status MCPHeadersUserCredentialStatus + CreatedAt time.Time + UpdatedAt time.Time +} + +// MCPHeadersFlowInitiation is the response returned by InitiateUserSubmissionFlow. +// Mirrors OAuth2FlowInitiation structurally so the resolver-side handling on +// the two per-user-auth surfaces stays uniform: a UUID, an auth-page URL, +// and an expiry. The "state" field is unused (no PKCE for headers); kept +// off the struct. +type MCPHeadersFlowInitiation struct { + FlowID string // Flow row primary key + FrontendURL string // {base}/workspace/mcp-sessions/auth?flow={id}#t={temp_token} + ExpiresAt time.Time // Flow expiration (15 min default; matches OAuth) +} + +// MCPHeadersProvider is the contract between the per-user-headers +// CredentialStore resolver and the configstore-backed implementation. Mirrors +// OAuth2Provider's per-user methods structurally so future provider +// implementations stay consistent. +type MCPHeadersProvider interface { + // GetCredentialByMode returns the persisted credential for a single + // identity dimension determined by mode. No fallback chain — exactly one + // identity column is queried. Returns ErrHeadersCredentialNotFound when + // the row is absent. Both 'active' and 'needs_update' rows are returned; + // orphaned rows are filtered out at the store layer. The runtime + // resolver's missing-keys check distinguishes usable from re-submission- + // required rows, so the caller doesn't need to inspect Status itself. + GetCredentialByMode(ctx context.Context, mode MCPAuthMode, identity, mcpClientID string) (*MCPHeadersUserCredential, error) + + // UpsertCredential persists a user-submitted set of header values for the + // (mode, identity, mcp_client_id) triple after a successful verify. The + // caller is expected to have run VerifyHeadersConnection before invoking + // this — the provider does not re-test the upstream connection. + UpsertCredential(ctx context.Context, cred *MCPHeadersUserCredential) error + + // DeleteCredential removes a credential row by its primary-key ID. + DeleteCredential(ctx context.Context, id string) error + + // InitiateUserSubmissionFlow creates a pending mcp_per_user_header_flows + // row keyed by (mode, identity, mcp_client_id), mints a mcp_headers_auth + // temp-token bound to the new row's ID, and returns the auth-page URL + // with the token embedded as a `#t=` fragment. Mirrors + // OAuth2Provider.InitiateUserOAuthFlow's role: the resolver calls this + // when an inline-401 fires, then puts the returned FrontendURL on the + // MCPAuthRequiredError so the caller can drive the submission flow. + // + // baseURL is the bifrost dashboard origin (e.g. "https://host") — the + // resolver pulls it from BifrostContextKeyMCPCallbackBaseURL and passes + // it in so the provider can construct the frontend URL without + // reaching into the BifrostContext itself. + InitiateUserSubmissionFlow(ctx context.Context, mode MCPAuthMode, identity, mcpClientID, baseURL string) (*MCPHeadersFlowInitiation, error) +} diff --git a/framework/configstore/migrations.go b/framework/configstore/migrations.go index aa1b9884b1..e8bc6e0c11 100644 --- a/framework/configstore/migrations.go +++ b/framework/configstore/migrations.go @@ -807,6 +807,12 @@ func triggerMigrations(ctx context.Context, db *gorm.DB) error { if err := migrationDropAzureAPIVersionColumn(ctx, db); err != nil { return err } + if err := migrationAddPerUserHeadersTables(ctx, db); err != nil { + return err + } + if err := migrationAddPerUserHeadersFlowsTable(ctx, db); err != nil { + return err + } return nil } @@ -8694,6 +8700,143 @@ func migrationDropVKAccessProfileIDColumn(ctx context.Context, db *gorm.DB) erro return nil } +// migrationAddPerUserHeadersTables creates the mcp_per_user_header_credentials +// table and adds the per_user_header_keys_json column to config_mcp_clients. +// Mirrors the partial-unique-index pattern from +// migrationAddOAuthAuthModeColumns so the per-user-headers credentials are +// keyed by (auth_mode, identity, mcp_client_id) the same way per-user OAuth +// tokens are. Forward-only on data — no rows exist yet. +func migrationAddPerUserHeadersTables(ctx context.Context, db *gorm.DB) error { + m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{ + ID: "add_mcp_per_user_header_credentials_table", + Migrate: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + mg := tx.Migrator() + + // 1) config_mcp_clients.per_user_header_keys_json (admin-defined + // schema of required header names; nullable / empty for all + // other auth types). + if !mg.HasColumn(&tables.TableMCPClient{}, "per_user_header_keys_json") { + if err := mg.AddColumn(&tables.TableMCPClient{}, "PerUserHeaderKeysJSON"); err != nil { + return fmt.Errorf("add per_user_header_keys_json column to config_mcp_clients: %w", err) + } + } + + // 2) mcp_per_user_header_credentials table. + if !mg.HasTable(&tables.TableMCPPerUserHeaderCredential{}) { + if err := mg.CreateTable(&tables.TableMCPPerUserHeaderCredential{}); err != nil { + return fmt.Errorf("create mcp_per_user_header_credentials table: %w", err) + } + } + + // 3) Partial unique indexes per auth_mode — matches the + // oauth_user_tokens layout so the cascade / orphan logic stays + // parallel. + partialUniques := []string{ + `CREATE UNIQUE INDEX IF NOT EXISTS idx_mcp_per_user_header_credentials_user_mcp + ON mcp_per_user_header_credentials (user_id, mcp_client_id) + WHERE auth_mode = 'user' AND user_id IS NOT NULL AND user_id != ''`, + `CREATE UNIQUE INDEX IF NOT EXISTS idx_mcp_per_user_header_credentials_vk_mcp + ON mcp_per_user_header_credentials (virtual_key_id, mcp_client_id) + WHERE auth_mode = 'vk' AND virtual_key_id IS NOT NULL AND virtual_key_id != ''`, + `CREATE UNIQUE INDEX IF NOT EXISTS idx_mcp_per_user_header_credentials_session_mcp + ON mcp_per_user_header_credentials (session_id, mcp_client_id) + WHERE auth_mode = 'session' AND session_id IS NOT NULL AND session_id != ''`, + } + for _, stmt := range partialUniques { + if err := tx.Exec(stmt).Error; err != nil { + return fmt.Errorf("create partial unique index on mcp_per_user_header_credentials: %w", err) + } + } + + // 4) Status-scoped partial indexes for cheap UI / cleanup queries. + statusIndexes := []string{ + `CREATE INDEX IF NOT EXISTS idx_mcp_per_user_header_credentials_orphaned + ON mcp_per_user_header_credentials (status) + WHERE status = 'orphaned'`, + `CREATE INDEX IF NOT EXISTS idx_mcp_per_user_header_credentials_needs_update + ON mcp_per_user_header_credentials (mcp_client_id) + WHERE status = 'needs_update'`, + } + for _, stmt := range statusIndexes { + if err := tx.Exec(stmt).Error; err != nil { + return fmt.Errorf("create status partial index on mcp_per_user_header_credentials: %w", err) + } + } + + return nil + }, + Rollback: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + mg := tx.Migrator() + for _, name := range []string{ + "idx_mcp_per_user_header_credentials_user_mcp", + "idx_mcp_per_user_header_credentials_vk_mcp", + "idx_mcp_per_user_header_credentials_session_mcp", + "idx_mcp_per_user_header_credentials_orphaned", + "idx_mcp_per_user_header_credentials_needs_update", + } { + if err := tx.Exec("DROP INDEX IF EXISTS " + name).Error; err != nil { + return fmt.Errorf("drop %s: %w", name, err) + } + } + if mg.HasTable(&tables.TableMCPPerUserHeaderCredential{}) { + if err := mg.DropTable(&tables.TableMCPPerUserHeaderCredential{}); err != nil { + return fmt.Errorf("drop mcp_per_user_header_credentials: %w", err) + } + } + if mg.HasColumn(&tables.TableMCPClient{}, "per_user_header_keys_json") { + if err := mg.DropColumn(&tables.TableMCPClient{}, "PerUserHeaderKeysJSON"); err != nil { + return fmt.Errorf("drop per_user_header_keys_json column from config_mcp_clients: %w", err) + } + } + return nil + }, + }}) + if err := m.Migrate(); err != nil { + return fmt.Errorf("error running add_mcp_per_user_header_credentials_table migration: %s", err.Error()) + } + return nil +} + +// migrationAddPerUserHeadersFlowsTable creates the +// mcp_per_user_header_flows table. Pending submission flow rows that +// mirror oauth_user_sessions for the per-user-headers surface — the +// resolver creates one when the inline-401 fires; the submit endpoint +// deletes the row on success; the sweep worker reaps expired pending +// rows. Lives in its own migration so it can land on DBs that already +// applied migrationAddPerUserHeadersTables (which only created the +// credentials table). +func migrationAddPerUserHeadersFlowsTable(ctx context.Context, db *gorm.DB) error { + m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{ + ID: "add_mcp_per_user_header_flows_table", + Migrate: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + mg := tx.Migrator() + if !mg.HasTable(&tables.TableMCPPerUserHeaderFlow{}) { + if err := mg.CreateTable(&tables.TableMCPPerUserHeaderFlow{}); err != nil { + return fmt.Errorf("create mcp_per_user_header_flows table: %w", err) + } + } + return nil + }, + Rollback: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + mg := tx.Migrator() + if mg.HasTable(&tables.TableMCPPerUserHeaderFlow{}) { + if err := mg.DropTable(&tables.TableMCPPerUserHeaderFlow{}); err != nil { + return fmt.Errorf("drop mcp_per_user_header_flows: %w", err) + } + } + return nil + }, + }}) + if err := m.Migrate(); err != nil { + return fmt.Errorf("error running add_mcp_per_user_header_flows_table migration: %s", err.Error()) + } + return nil +} + // migrationAddCreatedByUserIDColumnForVirtualKeys adds the created_by_user_id column to the governance_virtual_keys table. func migrationAddCreatedByUserIDColumnForVirtualKeys(ctx context.Context, db *gorm.DB) error { m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{ diff --git a/framework/configstore/rdb.go b/framework/configstore/rdb.go index 582b073e95..c73a966572 100644 --- a/framework/configstore/rdb.go +++ b/framework/configstore/rdb.go @@ -11,6 +11,7 @@ import ( "time" "github.com/bytedance/sonic" + "github.com/google/uuid" bifrost "github.com/maximhq/bifrost/core" providerUtils "github.com/maximhq/bifrost/core/providers/utils" "github.com/maximhq/bifrost/core/schemas" @@ -1503,6 +1504,7 @@ func (s *RDBConfigStore) GetMCPConfig(ctx context.Context) (*schemas.MCPConfig, Disabled: dbClient.Disabled, DiscoveredTools: dbClient.DiscoveredTools, DiscoveredToolNameMapping: dbClient.DiscoveredToolNameMapping, + PerUserHeaderKeys: dbClient.PerUserHeaderKeys, } } return &schemas.MCPConfig{ @@ -1543,6 +1545,7 @@ func (s *RDBConfigStore) GetMCPConfig(ctx context.Context) (*schemas.MCPConfig, ToolPricing: dbClient.ToolPricing, DiscoveredTools: dbClient.DiscoveredTools, DiscoveredToolNameMapping: dbClient.DiscoveredToolNameMapping, + PerUserHeaderKeys: dbClient.PerUserHeaderKeys, } } return &schemas.MCPConfig{ @@ -1628,6 +1631,7 @@ func (s *RDBConfigStore) GetMCPClientConfigByID(ctx context.Context, id string) ToolPricing: dbClient.ToolPricing, DiscoveredTools: dbClient.DiscoveredTools, DiscoveredToolNameMapping: dbClient.DiscoveredToolNameMapping, + PerUserHeaderKeys: dbClient.PerUserHeaderKeys, }, nil } @@ -1679,7 +1683,13 @@ func (s *RDBConfigStore) CreateMCPClientConfig(ctx context.Context, clientConfig // DiscoveredTools has json:"-" so deepCopy loses it; use original clientConfig DiscoveredTools: clientConfig.DiscoveredTools, DiscoveredToolNameMapping: clientConfig.DiscoveredToolNameMapping, - Disabled: clientConfigCopy.Disabled, + // PerUserHeaderKeys is the admin-declared schema for + // MCPAuthTypePerUserHeaders. Without this copy the BeforeSave + // hook persists an empty column, and on restart AddClient's + // validation rejects the row (empty PerUserHeaderKeys is + // invalid for per_user_headers), leaving the client orphaned. + PerUserHeaderKeys: clientConfigCopy.PerUserHeaderKeys, + Disabled: clientConfigCopy.Disabled, } if err := tx.WithContext(ctx).Create(&dbClient).Error; err != nil { return s.parseGormError(err) @@ -1778,6 +1788,16 @@ func (s *RDBConfigStore) UpdateMCPClientConfig(ctx context.Context, id string, c } toolNameMappingJSON = string(data) } + // Mirror BeforeSave for PerUserHeaderKeys — same map-based update + // path that bypasses GORM hooks for the other virtual fields. + perUserHeaderKeysJSON := "" + if clientConfig.PerUserHeaderKeys != nil { + data, marshalErr := json.Marshal(clientConfig.PerUserHeaderKeys) + if marshalErr != nil { + return fmt.Errorf("failed to marshal per_user_header_keys: %w", marshalErr) + } + perUserHeaderKeysJSON = string(data) + } headersJSONStr := string(headersJSON) if encrypt.IsEnabled() && headersJSONStr != "" && headersJSONStr != "{}" { @@ -1815,6 +1835,13 @@ func (s *RDBConfigStore) UpdateMCPClientConfig(ctx context.Context, id string, c if toolNameMappingJSON != "" { updates["tool_name_mapping_json"] = toolNameMappingJSON } + // Always persist PerUserHeaderKeys (empty string clears when the + // caller is dropping all keys). Treat absent (nil) as "preserve" by + // only writing when PerUserHeaderKeys was explicitly set on the + // update payload. + if clientConfig.PerUserHeaderKeys != nil { + updates["per_user_header_keys_json"] = perUserHeaderKeysJSON + } // Config-file driven reconciliation passes ConfigHash. In this mode we should // also sync connection/auth metadata from config.json and persist the hash. if clientConfigCopy.ConfigHash != "" { @@ -1879,8 +1906,9 @@ func (s *RDBConfigStore) DeleteMCPClientConfig(ctx context.Context, id string) e } } - // Delete per-user OAuth token + flow rows for this MCP client. Token - // rows reference mcp_client_id by the string client_id; nothing + // Delete per-user OAuth token + flow rows AND per-user header + // credentials + per-user header flow rows for this MCP client. All + // four reference mcp_client_id by the string client_id; nothing // auto-cascades, so we do it explicitly inside the same transaction // to keep cleanup atomic. if err := tx.WithContext(ctx).Where("mcp_client_id = ?", existingClient.ClientID).Delete(&tables.TableOauthUserToken{}).Error; err != nil { @@ -1889,6 +1917,12 @@ func (s *RDBConfigStore) DeleteMCPClientConfig(ctx context.Context, id string) e if err := tx.WithContext(ctx).Where("mcp_client_id = ?", existingClient.ClientID).Delete(&tables.TableOauthUserSession{}).Error; err != nil { return err } + if err := tx.WithContext(ctx).Where("mcp_client_id = ?", existingClient.ClientID).Delete(&tables.TableMCPPerUserHeaderCredential{}).Error; err != nil { + return err + } + if err := tx.WithContext(ctx).Where("mcp_client_id = ?", existingClient.ClientID).Delete(&tables.TableMCPPerUserHeaderFlow{}).Error; err != nil { + return err + } // Delete the client (this will also handle foreign key cascades) return tx.WithContext(ctx).Delete(&existingClient).Error @@ -2724,6 +2758,18 @@ func (s *RDBConfigStore) DeleteVirtualKey(ctx context.Context, id string, tx ... if err := txDB.WithContext(ctx).Where("virtual_key_id = ?", id).Delete(&tables.TableOauthUserToken{}).Error; err != nil { return err } + // Delete per-user MCP header credentials tied to this VK + if err := txDB.WithContext(ctx).Where("virtual_key_id = ?", id).Delete(&tables.TableMCPPerUserHeaderCredential{}).Error; err != nil { + return err + } + // Delete per-user MCP header flows tied to this VK — mirrors the + // OAuth session purge above. A pending flow's temp-token is valid + // for ~15 min; without this delete, a submission in flight could + // upsert a credential row pointing at the just-deleted VK and + // re-grant access after explicit revocation. + if err := txDB.WithContext(ctx).Where("virtual_key_id = ?", id).Delete(&tables.TableMCPPerUserHeaderFlow{}).Error; err != nil { + return err + } // Delete budgets owned by this virtual key if err := txDB.WithContext(ctx).Where("virtual_key_id = ?", id).Delete(&tables.TableBudget{}).Error; err != nil { return err @@ -5157,3 +5203,313 @@ func (s *RDBConfigStore) DeleteOrphanedOauthUserTokens(ctx context.Context, olde } return result.RowsAffected, nil } + +// ---------- Per-User MCP Header Credentials ---------- + +// GetMCPPerUserHeaderCredentialByMode looks up a usable per-user header +// credential by a single identity dimension. Returns both 'active' and +// 'needs_update' rows; the runtime resolver's missing-keys check +// distinguishes them — needs_update rows that genuinely lack keys for the +// current schema trigger the auth-required flow, while rows where the +// schema only narrowed still satisfy and get used. Orphaned rows are +// filtered at SQL because they mean the user lost grant: neither runtime +// resolution nor the flow-detail prefill UX should surface them. Mirrors +// GetOauthUserTokenByMode (which is stricter — OAuth has no needs_update +// equivalent because tokens are opaque and resubmission is the full IdP +// dance). +func (s *RDBConfigStore) GetMCPPerUserHeaderCredentialByMode(ctx context.Context, mode schemas.MCPAuthMode, identity, mcpClientID string) (*tables.TableMCPPerUserHeaderCredential, error) { + if identity == "" || mcpClientID == "" { + return nil, nil + } + var cred tables.TableMCPPerUserHeaderCredential + var result *gorm.DB + statuses := []string{"active", "needs_update"} + switch mode { + case schemas.MCPAuthModeUser: + result = s.DB().WithContext(ctx). + Where("auth_mode = ? AND user_id = ? AND mcp_client_id = ? AND status IN ?", string(schemas.MCPAuthModeUser), identity, mcpClientID, statuses). + First(&cred) + case schemas.MCPAuthModeVK: + result = s.DB().WithContext(ctx). + Where("auth_mode = ? AND virtual_key_id = ? AND mcp_client_id = ? AND status IN ?", string(schemas.MCPAuthModeVK), identity, mcpClientID, statuses). + First(&cred) + case schemas.MCPAuthModeSession: + result = s.DB().WithContext(ctx). + Where("auth_mode = ? AND session_id = ? AND mcp_client_id = ? AND status IN ?", string(schemas.MCPAuthModeSession), identity, mcpClientID, statuses). + First(&cred) + default: + return nil, fmt.Errorf("unknown auth mode: %s", mode) + } + if result.Error != nil { + if errors.Is(result.Error, gorm.ErrRecordNotFound) { + return nil, nil + } + return nil, fmt.Errorf("failed to get mcp per-user header credential by mode %s: %w", mode, result.Error) + } + return &cred, nil +} + +// GetMCPPerUserHeaderCredentialByID looks up a single row by primary key. +// Returns nil, nil when not found. +func (s *RDBConfigStore) GetMCPPerUserHeaderCredentialByID(ctx context.Context, id string) (*tables.TableMCPPerUserHeaderCredential, error) { + if id == "" { + return nil, nil + } + var cred tables.TableMCPPerUserHeaderCredential + if err := s.ScopedDB(ctx).Where("id = ?", id).First(&cred).Error; err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, nil + } + return nil, fmt.Errorf("failed to get mcp per-user header credential by id %s: %w", id, err) + } + return &cred, nil +} + +// UpsertMCPPerUserHeaderCredential atomically inserts or updates a credential +// row keyed by (auth_mode, identity, mcp_client_id). Mirrors +// CreateOauthUserToken — the row represents the (identity, mcp_client) +// binding, so a re-submit preserves CreatedAt. +func (s *RDBConfigStore) UpsertMCPPerUserHeaderCredential(ctx context.Context, cred *tables.TableMCPPerUserHeaderCredential) error { + return s.DB().WithContext(ctx).Transaction(func(tx *gorm.DB) error { + var existing tables.TableMCPPerUserHeaderCredential + var lookupErr error + switch { + case cred.UserID != nil && *cred.UserID != "": + lookupErr = dbForUpdate(tx). + Where("auth_mode = ? AND user_id = ? AND mcp_client_id = ?", string(schemas.MCPAuthModeUser), *cred.UserID, cred.MCPClientID). + First(&existing).Error + case cred.VirtualKeyID != nil && *cred.VirtualKeyID != "": + lookupErr = dbForUpdate(tx). + Where("auth_mode = ? AND virtual_key_id = ? AND mcp_client_id = ?", string(schemas.MCPAuthModeVK), *cred.VirtualKeyID, cred.MCPClientID). + First(&existing).Error + case cred.SessionID != "": + lookupErr = dbForUpdate(tx). + Where("auth_mode = ? AND session_id = ? AND mcp_client_id = ?", string(schemas.MCPAuthModeSession), cred.SessionID, cred.MCPClientID). + First(&existing).Error + default: + lookupErr = gorm.ErrRecordNotFound + } + + if lookupErr == nil { + cred.ID = existing.ID + cred.CreatedAt = existing.CreatedAt + return tx.Save(cred).Error + } + if !errors.Is(lookupErr, gorm.ErrRecordNotFound) { + return fmt.Errorf("failed to query mcp per-user header credential: %w", lookupErr) + } + if cred.ID == "" { + cred.ID = uuid.New().String() + } + if err := tx.Create(cred).Error; err != nil { + return fmt.Errorf("failed to create mcp per-user header credential: %w", err) + } + return nil + }) +} + +// DeleteMCPPerUserHeaderCredential removes a credential row by its primary key. +func (s *RDBConfigStore) DeleteMCPPerUserHeaderCredential(ctx context.Context, id string) error { + if id == "" { + return nil + } + result := s.DB().WithContext(ctx).Where("id = ?", id).Delete(&tables.TableMCPPerUserHeaderCredential{}) + if result.Error != nil { + return fmt.Errorf("failed to delete mcp per-user header credential: %w", result.Error) + } + return nil +} + +// ListAllMCPPerUserHeaderCredentials returns every row regardless of status. +// The sessions UI surfaces non-active states (needs_update / orphaned) with +// distinct affordances; filtering here would only hide rows the user needs to +// act on. Runtime lookups apply their own status='active' filter and don't go +// through this method. +func (s *RDBConfigStore) ListAllMCPPerUserHeaderCredentials(ctx context.Context) ([]tables.TableMCPPerUserHeaderCredential, error) { + var creds []tables.TableMCPPerUserHeaderCredential + if err := s.ScopedDB(ctx). + Preload("MCPClient", func(db *gorm.DB) *gorm.DB { return db.Select("client_id, name") }). + Preload("VirtualKey", func(db *gorm.DB) *gorm.DB { return db.Select("id, name") }). + Order("created_at DESC"). + Find(&creds).Error; err != nil { + return nil, fmt.Errorf("failed to list all mcp per-user header credentials: %w", err) + } + return creds, nil +} + +// MarkMCPPerUserHeaderCredentialsNeedsUpdate flips status to 'needs_update' +// for every active row tied to mcpClientID. Called when the admin changes +// PerUserHeaderKeys on the MCP client config. +func (s *RDBConfigStore) MarkMCPPerUserHeaderCredentialsNeedsUpdate(ctx context.Context, mcpClientID string) error { + if mcpClientID == "" { + return nil + } + result := s.DB().WithContext(ctx). + Model(&tables.TableMCPPerUserHeaderCredential{}). + Where("mcp_client_id = ? AND status = ?", mcpClientID, "active"). + Update("status", "needs_update") + if result.Error != nil { + return fmt.Errorf("failed to mark mcp per-user header credentials needs_update for client %s: %w", mcpClientID, result.Error) + } + return nil +} + +// DeleteOrphanedMCPPerUserHeaderCredentials hard-deletes rows in 'orphaned' +// state longer than olderThan. Skipped silently when olderThan is zero or +// negative. +func (s *RDBConfigStore) DeleteOrphanedMCPPerUserHeaderCredentials(ctx context.Context, olderThan time.Duration) (int64, error) { + if olderThan <= 0 { + return 0, nil + } + cutoff := time.Now().Add(-olderThan) + result := s.DB().WithContext(ctx). + Where("status = ? AND updated_at < ?", "orphaned", cutoff). + Delete(&tables.TableMCPPerUserHeaderCredential{}) + if result.Error != nil { + return 0, fmt.Errorf("failed to delete orphaned mcp per-user header credentials: %w", result.Error) + } + return result.RowsAffected, nil +} + +// CreateMCPPerUserHeaderFlow persists a pending per-user-headers submission +// flow row. ID is set by the caller (typically a fresh UUID). +func (s *RDBConfigStore) CreateMCPPerUserHeaderFlow(ctx context.Context, flow *tables.TableMCPPerUserHeaderFlow) error { + if flow == nil { + return fmt.Errorf("flow is nil") + } + if err := s.DB().WithContext(ctx).Create(flow).Error; err != nil { + return fmt.Errorf("failed to create mcp per-user header flow: %w", err) + } + return nil +} + +// GetMCPPerUserHeaderFlowByID looks up a flow row by primary key. +// Returns nil, nil when not found. +func (s *RDBConfigStore) GetMCPPerUserHeaderFlowByID(ctx context.Context, id string) (*tables.TableMCPPerUserHeaderFlow, error) { + if id == "" { + return nil, nil + } + var flow tables.TableMCPPerUserHeaderFlow + if err := s.ScopedDB(ctx). + Preload("MCPClient", func(db *gorm.DB) *gorm.DB { + return db.Select("client_id, name, headers_json, allowed_extra_headers_json, per_user_header_keys_json") + }). + Preload("VirtualKey", func(db *gorm.DB) *gorm.DB { return db.Select("id, name") }). + Where("id = ?", id).First(&flow).Error; err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, nil + } + return nil, fmt.Errorf("failed to get mcp per-user header flow by id %s: %w", id, err) + } + return &flow, nil +} + +// DeleteMCPPerUserHeaderFlow hard-deletes a single flow row by primary key. +// Called on submit-success and on revoke; no-op when the row is absent so +// terminal-state transitions are idempotent. +func (s *RDBConfigStore) DeleteMCPPerUserHeaderFlow(ctx context.Context, id string) error { + if id == "" { + return nil + } + if err := s.DB().WithContext(ctx).Where("id = ?", id).Delete(&tables.TableMCPPerUserHeaderFlow{}).Error; err != nil { + return fmt.Errorf("failed to delete mcp per-user header flow %s: %w", id, err) + } + return nil +} + +// GetMCPPerUserHeaderFlowByModeIdentityAndMCPClient returns the canonical +// pending flow row for the (mode, identity, mcp_client) triple, or nil +// when none exists. Mirrors GetOauthUserSessionByModeIdentityAndMCPClient. +// Used by InitiateUserSubmissionFlow to keep at most one pending row per +// binding (re-init updates in place instead of inserting a duplicate). +func (s *RDBConfigStore) GetMCPPerUserHeaderFlowByModeIdentityAndMCPClient(ctx context.Context, mode schemas.MCPAuthMode, identity, mcpClientID string) (*tables.TableMCPPerUserHeaderFlow, error) { + if identity == "" || mcpClientID == "" { + return nil, nil + } + q := s.DB().WithContext(ctx). + Where("flow_mode = ? AND mcp_client_id = ?", string(mode), mcpClientID) + switch mode { + case schemas.MCPAuthModeUser: + q = q.Where("user_id = ?", identity) + case schemas.MCPAuthModeVK: + q = q.Where("virtual_key_id = ?", identity) + case schemas.MCPAuthModeSession: + q = q.Where("session_id = ?", identity) + default: + return nil, fmt.Errorf("unknown auth mode: %s", mode) + } + var flow tables.TableMCPPerUserHeaderFlow + if err := q.First(&flow).Error; err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, nil + } + return nil, fmt.Errorf("failed to get mcp per-user header flow by mode/identity/mcp_client: %w", err) + } + return &flow, nil +} + +// UpdateMCPPerUserHeaderFlow updates a flow row in place. +func (s *RDBConfigStore) UpdateMCPPerUserHeaderFlow(ctx context.Context, flow *tables.TableMCPPerUserHeaderFlow) error { + if flow == nil || flow.ID == "" { + return fmt.Errorf("flow id is required") + } + if err := s.DB().WithContext(ctx).Save(flow).Error; err != nil { + return fmt.Errorf("failed to update mcp per-user header flow: %w", err) + } + return nil +} + +// DeleteMCPPerUserHeaderFlowsByModeIdentityAndMCPClient hard-deletes any +// flow rows matching the binding. Mirrors +// DeleteOauthUserSessionsByModeIdentityAndMCPClient. +func (s *RDBConfigStore) DeleteMCPPerUserHeaderFlowsByModeIdentityAndMCPClient(ctx context.Context, mode schemas.MCPAuthMode, identity, mcpClientID string) error { + if identity == "" || mcpClientID == "" { + return nil + } + q := s.DB().WithContext(ctx). + Where("flow_mode = ? AND mcp_client_id = ?", string(mode), mcpClientID) + switch mode { + case schemas.MCPAuthModeUser: + q = q.Where("user_id = ?", identity) + case schemas.MCPAuthModeVK: + q = q.Where("virtual_key_id = ?", identity) + case schemas.MCPAuthModeSession: + q = q.Where("session_id = ?", identity) + default: + return fmt.Errorf("unknown auth mode: %s", mode) + } + if err := q.Delete(&tables.TableMCPPerUserHeaderFlow{}).Error; err != nil { + return fmt.Errorf("failed to delete mcp per-user header flows by mode/identity/mcp_client: %w", err) + } + return nil +} + +// ListAllPendingMCPPerUserHeaderFlows returns all pending header-submission +// flow rows whose expiry is in the future. Mirrors +// ListAllPendingOauthUserSessions. Visibility scoping is handled by the +// enterprise configstore layer via DAC; OSS sees everything. +func (s *RDBConfigStore) ListAllPendingMCPPerUserHeaderFlows(ctx context.Context) ([]tables.TableMCPPerUserHeaderFlow, error) { + var flows []tables.TableMCPPerUserHeaderFlow + if err := s.ScopedDB(ctx). + Preload("MCPClient", func(db *gorm.DB) *gorm.DB { return db.Select("client_id, name") }). + Preload("VirtualKey", func(db *gorm.DB) *gorm.DB { return db.Select("id, name") }). + Where("status = ? AND expires_at > ?", "pending", time.Now()). + Order("created_at DESC"). + Find(&flows).Error; err != nil { + return nil, fmt.Errorf("failed to list all pending mcp per-user header flows: %w", err) + } + return flows, nil +} + +// DeleteExpiredMCPPerUserHeaderFlows hard-deletes pending flow rows whose +// ExpiresAt has passed. Status filter excludes already-completed rows +// (which the submit path deletes immediately anyway). +func (s *RDBConfigStore) DeleteExpiredMCPPerUserHeaderFlows(ctx context.Context) (int64, error) { + result := s.DB().WithContext(ctx). + Where("expires_at < ? AND status = ?", time.Now(), "pending"). + Delete(&tables.TableMCPPerUserHeaderFlow{}) + if result.Error != nil { + return 0, fmt.Errorf("failed to delete expired mcp per-user header flows: %w", result.Error) + } + return result.RowsAffected, nil +} diff --git a/framework/configstore/store.go b/framework/configstore/store.go index 196a6ae3cd..57f14581da 100644 --- a/framework/configstore/store.go +++ b/framework/configstore/store.go @@ -386,6 +386,61 @@ type ConfigStore interface { // and updated_at is older than olderThan. Returns the number of rows removed. DeleteOrphanedOauthUserTokens(ctx context.Context, olderThan time.Duration) (int64, error) + // Per-user MCP header credential CRUD. Storage analog of per-user OAuth + // tokens for MCPAuthTypePerUserHeaders clients. The row holds an encrypted + // JSON blob of header_name → value pairs keyed by (auth_mode, identity, + // mcp_client_id). + GetMCPPerUserHeaderCredentialByMode(ctx context.Context, mode schemas.MCPAuthMode, identity, mcpClientID string) (*tables.TableMCPPerUserHeaderCredential, error) + GetMCPPerUserHeaderCredentialByID(ctx context.Context, id string) (*tables.TableMCPPerUserHeaderCredential, error) + UpsertMCPPerUserHeaderCredential(ctx context.Context, cred *tables.TableMCPPerUserHeaderCredential) error + DeleteMCPPerUserHeaderCredential(ctx context.Context, id string) error + // ListAllMCPPerUserHeaderCredentials returns every row regardless of + // status. Mirrors ListAllOauthUserTokens — the sessions UI surfaces + // non-active states (needs_update / orphaned) with distinct affordances. + ListAllMCPPerUserHeaderCredentials(ctx context.Context) ([]tables.TableMCPPerUserHeaderCredential, error) + // MarkMCPPerUserHeaderCredentialsNeedsUpdate flips status to 'needs_update' + // for every row tied to mcpClientID. Called when the admin changes + // PerUserHeaderKeys on the MCP client config: existing user submissions + // stay (so the UI can prefill known values) but are excluded from runtime + // lookups until the user re-submits. + MarkMCPPerUserHeaderCredentialsNeedsUpdate(ctx context.Context, mcpClientID string) error + // DeleteOrphanedMCPPerUserHeaderCredentials hard-deletes rows where + // status='orphaned' and updated_at is older than olderThan. + DeleteOrphanedMCPPerUserHeaderCredentials(ctx context.Context, olderThan time.Duration) (int64, error) + + // Per-user-headers submission flow CRUD. Mirrors the OAuth user-session + // surface — the resolver creates a pending flow row when the inline-401 + // fires, the submit endpoint deletes the row on success, and the sweep + // worker reaps expired pending rows. + CreateMCPPerUserHeaderFlow(ctx context.Context, flow *tables.TableMCPPerUserHeaderFlow) error + GetMCPPerUserHeaderFlowByID(ctx context.Context, id string) (*tables.TableMCPPerUserHeaderFlow, error) + // GetMCPPerUserHeaderFlowByModeIdentityAndMCPClient returns the canonical + // pending flow row for the (mode, identity, mcp_client) triple, if any. + // Companion to GetOauthUserSessionByModeIdentityAndMCPClient — used by + // InitiateUserSubmissionFlow to keep at most one pending row per binding + // (mirrors OAuth's single-row-per-binding invariant). + GetMCPPerUserHeaderFlowByModeIdentityAndMCPClient(ctx context.Context, mode schemas.MCPAuthMode, identity, mcpClientID string) (*tables.TableMCPPerUserHeaderFlow, error) + // UpdateMCPPerUserHeaderFlow updates a flow row in place. Used on the + // reauth/re-init path to rotate ExpiresAt without spawning a new row. + UpdateMCPPerUserHeaderFlow(ctx context.Context, flow *tables.TableMCPPerUserHeaderFlow) error + // DeleteMCPPerUserHeaderFlowsByModeIdentityAndMCPClient hard-deletes any + // pending flow rows for a binding. Called from revoke so a credential + // delete also clears any in-flight resubmission flow for the same + // (mode, identity, mcp_client). Mirrors + // DeleteOauthUserSessionsByModeIdentityAndMCPClient. + DeleteMCPPerUserHeaderFlowsByModeIdentityAndMCPClient(ctx context.Context, mode schemas.MCPAuthMode, identity, mcpClientID string) error + DeleteMCPPerUserHeaderFlow(ctx context.Context, id string) error + // ListAllPendingMCPPerUserHeaderFlows returns every non-expired flow row + // with status='pending', regardless of caller identity. Visibility scoping + // happens at the enterprise configstore layer via DAC scope; OSS sees + // everything. Used by the sessions list endpoint to surface pending + // submission flows alongside completed credentials. Mirrors + // ListAllPendingOauthUserSessions on the OAuth side. + ListAllPendingMCPPerUserHeaderFlows(ctx context.Context) ([]tables.TableMCPPerUserHeaderFlow, error) + // DeleteExpiredMCPPerUserHeaderFlows hard-deletes pending flow rows whose + // ExpiresAt has passed. Returns the number of rows removed. + DeleteExpiredMCPPerUserHeaderFlows(ctx context.Context) (int64, error) + // Not found retry wrapper RetryOnNotFound(ctx context.Context, fn func(ctx context.Context) (any, error), maxRetries int, retryDelay time.Duration) (any, error) diff --git a/framework/configstore/tables/mcp.go b/framework/configstore/tables/mcp.go index 87cf7e586b..0f71b493d5 100644 --- a/framework/configstore/tables/mcp.go +++ b/framework/configstore/tables/mcp.go @@ -33,10 +33,16 @@ type TableMCPClient struct { ToolNameMappingJSON string `gorm:"type:text" json:"-"` // JSON serialized map[string]string // OAuth authentication fields - AuthType string `gorm:"type:varchar(20);default:'headers'" json:"auth_type"` // "none", "headers", "oauth" + AuthType string `gorm:"type:varchar(20);default:'headers'" json:"auth_type"` // "none", "headers", "oauth", "per_user_oauth", "per_user_headers" OauthConfigID *string `gorm:"type:varchar(255);index;constraint:OnDelete:CASCADE" json:"oauth_config_id"` // Foreign key to oauth_configs.ID with CASCADE delete OauthConfig *TableOauthConfig `gorm:"foreignKey:OauthConfigID;references:ID;constraint:OnDelete:CASCADE" json:"-"` // Gorm relationship + // Per-user-headers schema: admin-declared list of header *names* that each + // caller must supply. Empty/null for all other auth types. Used by both + // the resolver (intersect with persisted user values) and by + // utils.StaticConfigHeaders (strip from plugin-visible static headers). + PerUserHeaderKeysJSON string `gorm:"type:text" json:"-"` // JSON serialized []string + AllowOnAllVirtualKeys bool `gorm:"default:false" json:"allow_on_all_virtual_keys"` // Whether to allow the MCP client to run on all virtual keys Disabled bool `gorm:"default:false" json:"disabled"` // Whether the client is intentionally disabled @@ -58,6 +64,7 @@ type TableMCPClient struct { ToolPricing map[string]float64 `gorm:"-" json:"tool_pricing"` DiscoveredTools map[string]schemas.ChatTool `gorm:"-" json:"-"` DiscoveredToolNameMapping map[string]string `gorm:"-" json:"-"` + PerUserHeaderKeys []string `gorm:"-" json:"per_user_header_keys"` } // TableName sets the table name for each model @@ -161,6 +168,16 @@ func (c *TableMCPClient) BeforeSave(tx *gorm.DB) error { c.ToolNameMappingJSON = string(data) } + if c.PerUserHeaderKeys != nil { + data, err := json.Marshal(c.PerUserHeaderKeys) + if err != nil { + return err + } + c.PerUserHeaderKeysJSON = string(data) + } else { + c.PerUserHeaderKeysJSON = "" + } + // Encrypt sensitive fields after serialization. // Always set EncryptionStatus when encryption is enabled so the startup // batch pass does not re-process this row indefinitely. @@ -249,5 +266,10 @@ func (c *TableMCPClient) AfterFind(tx *gorm.DB) error { return err } } + if c.PerUserHeaderKeysJSON != "" { + if err := sonic.Unmarshal([]byte(c.PerUserHeaderKeysJSON), &c.PerUserHeaderKeys); err != nil { + return err + } + } return nil } diff --git a/framework/configstore/tables/mcp_per_user_headers.go b/framework/configstore/tables/mcp_per_user_headers.go new file mode 100644 index 0000000000..ba863a1a6d --- /dev/null +++ b/framework/configstore/tables/mcp_per_user_headers.go @@ -0,0 +1,151 @@ +package tables + +import ( + "encoding/json" + "fmt" + "time" + + "github.com/maximhq/bifrost/framework/encrypt" + "gorm.io/gorm" +) + +// TableMCPPerUserHeaderFlow tracks pending per-user-headers submission +// flows. Mirrors TableOauthUserSession structurally so the per-user-auth +// surfaces (OAuth + headers) have identical lifecycles: an inline-401 +// from the resolver creates a flow row, the auth-page URL carries the +// flow's ID (with a temp-token in the URL fragment for unauthenticated +// callers), and the submission endpoint completes / deletes the row. +// +// Unlike OAuth, there is no PKCE state to round-trip — the only durable +// state this row carries is (mcp_client_id, identity) so the submission +// endpoint can scope the upsert. No state column either: the row exists +// only while the submission is pending; submit completes by deleting it. +type TableMCPPerUserHeaderFlow struct { + ID string `gorm:"type:varchar(255);primaryKey" json:"id"` // Flow UUID + MCPClientID string `gorm:"type:varchar(255);not null;index" json:"mcp_client_id"` // Which MCP server this submission is for + SessionID string `gorm:"type:varchar(255);index" json:"session_id,omitempty"` // Session-mode identity: client-asserted x-bf-mcp-session-id. Empty for vk/user mode rows. + VirtualKeyID *string `gorm:"type:varchar(255);index" json:"virtual_key_id"` // VK identity (vk-mode rows) + UserID *string `gorm:"type:varchar(255);index" json:"user_id"` // User identity (user-mode rows) + FlowMode string `gorm:"type:varchar(20);not null;default:'vk'" json:"flow_mode"` // 'user' | 'vk' | 'session' — mirrors the credential row's AuthMode; immutable after creation + Status string `gorm:"type:varchar(50);not null;index" json:"status"` // "pending", "completed", "expired" + ExpiresAt time.Time `gorm:"index;not null" json:"expires_at"` // Flow expiration (15 min default) + CreatedAt time.Time `gorm:"index;not null" json:"created_at"` + UpdatedAt time.Time `gorm:"index;not null" json:"updated_at"` + + // Display-only relations (no DB-level FK constraint; preloaded for sessions UI). + MCPClient *TableMCPClient `gorm:"foreignKey:MCPClientID;references:ClientID" json:"-"` + VirtualKey *TableVirtualKey `gorm:"foreignKey:VirtualKeyID;references:ID" json:"-"` + + // User mirrors TableOauthUserSession.User — populated post-fetch by the + // enterprise configstore wrapper for the sessions UI. OSS leaves it nil. + User *OauthUserSummary `gorm:"-" json:"-"` +} + +// TableName sets the table name. +func (TableMCPPerUserHeaderFlow) TableName() string { + return "mcp_per_user_header_flows" +} + +// BeforeSave defaults Status to 'pending' when unset. +func (f *TableMCPPerUserHeaderFlow) BeforeSave(tx *gorm.DB) error { + if f.Status == "" { + f.Status = "pending" + } + return nil +} + +// TableMCPPerUserHeaderCredential stores per-user header credentials for +// MCPAuthTypePerUserHeaders MCP clients. Each row holds the encrypted header +// values for a specific identity × MCP client pair. Exactly one identity +// column (UserID, VirtualKeyID, or SessionID) is populated per row; AuthMode +// records which one. Mirrors TableOauthUserToken structurally so cascade / +// orphan-sweep logic stays parallel between the two per-user auth surfaces. +// +// HeadersJSON holds a JSON-encoded map[string]string of header_name → value, +// encrypted at rest via the shared encrypt package (same key as +// oauth_user_tokens). Schema (i.e. the set of allowed header names) lives on +// TableMCPClient.PerUserHeaderKeysJSON; this table holds the values only. +type TableMCPPerUserHeaderCredential struct { + ID string `gorm:"type:varchar(255);primaryKey" json:"id"` // UUID + SessionID string `gorm:"type:varchar(255);index" json:"session_id,omitempty"` // Session-mode identity: client-asserted x-bf-mcp-session-id. Empty for vk/user mode rows. + VirtualKeyID *string `gorm:"type:varchar(255);index" json:"virtual_key_id"` // VK identity (vk-mode rows) + UserID *string `gorm:"type:varchar(255);index" json:"user_id"` // User identity (user-mode rows) + MCPClientID string `gorm:"type:varchar(255);not null;index" json:"mcp_client_id"` // Which MCP server + AuthMode string `gorm:"type:varchar(20);not null" json:"auth_mode"` // 'user' | 'vk' | 'session' — which identity column keys this row + Status string `gorm:"type:varchar(20);not null;default:'active'" json:"status"` // 'active' | 'orphaned' | 'needs_update' + HeadersJSON string `gorm:"type:text;not null" json:"-"` // Encrypted JSON map[string]string of user-supplied header values + EncryptionStatus string `gorm:"type:varchar(20);default:'plain_text'" json:"-"` + CreatedAt time.Time `gorm:"index;not null" json:"created_at"` + UpdatedAt time.Time `gorm:"index;not null" json:"updated_at"` + + // Display-only relations (no DB-level FK constraint; preloaded for sessions UI). + MCPClient *TableMCPClient `gorm:"foreignKey:MCPClientID;references:ClientID" json:"-"` + VirtualKey *TableVirtualKey `gorm:"foreignKey:VirtualKeyID;references:ID" json:"-"` + + // User mirrors TableOauthUserToken.User — populated post-fetch by enterprise + // configstore wrapper for the sessions UI. OSS leaves it nil. + User *OauthUserSummary `gorm:"-" json:"-"` +} + +func (TableMCPPerUserHeaderCredential) TableName() string { + return "mcp_per_user_header_credentials" +} + +// BeforeSave encrypts HeadersJSON when encryption is enabled. The JSON +// serialization is the caller's responsibility (see SetHeaders). When +// encryption is not configured (no BIFROST_ENCRYPTION_KEY), the field +// is stored as plaintext and EncryptionStatus stays "plain_text" — same +// convention as TableOauthUserToken. +func (c *TableMCPPerUserHeaderCredential) BeforeSave(tx *gorm.DB) error { + if c.Status == "" { + c.Status = "active" + } + if c.HeadersJSON == "" { + c.HeadersJSON = "{}" + } + if encrypt.IsEnabled() { + if err := encryptString(&c.HeadersJSON); err != nil { + return fmt.Errorf("failed to encrypt mcp per-user header credential headers: %w", err) + } + c.EncryptionStatus = EncryptionStatusEncrypted + } + return nil +} + +// AfterFind decrypts HeadersJSON when the row is marked encrypted. +func (c *TableMCPPerUserHeaderCredential) AfterFind(tx *gorm.DB) error { + if c.EncryptionStatus == EncryptionStatusEncrypted { + if err := decryptString(&c.HeadersJSON); err != nil { + return fmt.Errorf("failed to decrypt mcp per-user header credential headers: %w", err) + } + } + return nil +} + +// SetHeaders serializes the caller-supplied header map into HeadersJSON. +// Callers must use this rather than assigning HeadersJSON directly so the +// JSON shape stays consistent. +func (c *TableMCPPerUserHeaderCredential) SetHeaders(headers map[string]string) error { + if headers == nil { + headers = map[string]string{} + } + data, err := json.Marshal(headers) + if err != nil { + return fmt.Errorf("failed to serialize mcp per-user header credential headers: %w", err) + } + c.HeadersJSON = string(data) + return nil +} + +// GetHeaders deserializes HeadersJSON into a header map. Returns an empty map +// for the zero JSON (`{}` or empty string) so callers do not need to nil-check. +func (c *TableMCPPerUserHeaderCredential) GetHeaders() (map[string]string, error) { + headers := map[string]string{} + if c.HeadersJSON == "" || c.HeadersJSON == "{}" { + return headers, nil + } + if err := json.Unmarshal([]byte(c.HeadersJSON), &headers); err != nil { + return nil, fmt.Errorf("failed to deserialize mcp per-user header credential headers: %w", err) + } + return headers, nil +} diff --git a/framework/mcp_headers/main.go b/framework/mcp_headers/main.go new file mode 100644 index 0000000000..8f24681b7d --- /dev/null +++ b/framework/mcp_headers/main.go @@ -0,0 +1,301 @@ +// Package mcp_headers implements schemas.MCPHeadersProvider against the +// configstore. It is the storage backend for MCPAuthTypePerUserHeaders MCP +// clients — the parallel of framework/oauth2 for the OAuth-flavored per-user +// auth type. +// +// The provider is intentionally storage-only: it does not run upstream +// verification (that's clientmanager.VerifyHeadersConnection) and does not +// build inline-401 errors (that's the credstore resolver). It just maps +// between the table type and the in-memory schema view. +package mcp_headers + +import ( + "context" + "errors" + "fmt" + "strings" + "sync" + "time" + + "github.com/google/uuid" + bifrost "github.com/maximhq/bifrost/core" + "github.com/maximhq/bifrost/core/schemas" + "github.com/maximhq/bifrost/framework/configstore" + "github.com/maximhq/bifrost/framework/configstore/tables" + "github.com/maximhq/bifrost/framework/temptoken" +) + +// SubmissionFlowTTL caps how long a pending headers submission flow row +// (and the temp token bound to it) remains valid. Mirrors the OAuth flow +// expiry so the two per-user-auth surfaces feel uniform to the user. +const SubmissionFlowTTL = 15 * time.Minute + +// Provider implements schemas.MCPHeadersProvider. +type Provider struct { + configStore configstore.ConfigStore + logger schemas.Logger + + // tempTokens, when non-nil, is used by InitiateUserSubmissionFlow to + // mint a short-lived mcp_headers_auth token and embed it in the + // returned auth-page URL as a fragment. Optional — when nil, the URL + // is returned without a fragment and the page works only for callers + // already authenticated to the dashboard. Mirrors + // oauth2.OAuth2Provider.tempTokens exactly. + mu sync.RWMutex + tempTokens *temptoken.Service +} + +// NewProvider constructs a configstore-backed MCPHeadersProvider. Mirrors +// oauth2.NewOAuth2Provider so the wiring in transports/bifrost-http stays +// symmetric between the two per-user auth surfaces. +func NewProvider(configStore configstore.ConfigStore, logger schemas.Logger) *Provider { + if logger == nil { + logger = bifrost.NewDefaultLogger(schemas.LogLevelInfo) + } + return &Provider{configStore: configStore, logger: logger} +} + +// SetTempTokenService installs the temp-token service used by +// InitiateUserSubmissionFlow to mint the mcp_headers_auth token embedded +// in the auth-page URL fragment. Called by server startup once both +// services have been constructed (the provider is built first by +// lib/config.go, the service later by the HTTP transport). Mirrors +// oauth2.OAuth2Provider.SetTempTokenService. +func (p *Provider) SetTempTokenService(svc *temptoken.Service) { + p.mu.Lock() + defer p.mu.Unlock() + p.tempTokens = svc +} + +// GetCredentialByMode looks up the active credential row for the given +// identity dimension. Returns ErrHeadersCredentialNotFound when the row is +// absent so callers can switch on the sentinel. +func (p *Provider) GetCredentialByMode(ctx context.Context, mode schemas.MCPAuthMode, identity, mcpClientID string) (*schemas.MCPHeadersUserCredential, error) { + if p.configStore == nil { + return nil, schemas.ErrHeadersCredentialProviderNotAvailable + } + if strings.TrimSpace(identity) == "" || strings.TrimSpace(mcpClientID) == "" { + return nil, schemas.ErrHeadersCredentialNotFound + } + row, err := p.configStore.GetMCPPerUserHeaderCredentialByMode(ctx, mode, identity, mcpClientID) + if err != nil { + return nil, fmt.Errorf("load mcp per-user header credential: %w", err) + } + if row == nil { + return nil, schemas.ErrHeadersCredentialNotFound + } + cred, err := rowToCredential(row) + if err != nil { + return nil, err + } + return cred, nil +} + +// UpsertCredential persists the caller-supplied credential. The caller is +// expected to have run clientmanager.VerifyHeadersConnection before invoking +// this — the provider trusts the values and only handles serialization + +// storage. +func (p *Provider) UpsertCredential(ctx context.Context, cred *schemas.MCPHeadersUserCredential) error { + if p.configStore == nil { + return schemas.ErrHeadersCredentialProviderNotAvailable + } + if cred == nil { + return errors.New("nil credential") + } + if strings.TrimSpace(cred.MCPClientID) == "" { + return errors.New("mcp_client_id is required") + } + row, err := credentialToRow(cred) + if err != nil { + return err + } + if err := p.configStore.UpsertMCPPerUserHeaderCredential(ctx, row); err != nil { + return fmt.Errorf("upsert mcp per-user header credential: %w", err) + } + // Propagate the row ID back so the caller can reference it (e.g. revoke later). + cred.ID = row.ID + cred.CreatedAt = row.CreatedAt + cred.UpdatedAt = row.UpdatedAt + return nil +} + +// DeleteCredential removes a credential by primary key. +func (p *Provider) DeleteCredential(ctx context.Context, id string) error { + if p.configStore == nil { + return schemas.ErrHeadersCredentialProviderNotAvailable + } + if strings.TrimSpace(id) == "" { + return nil + } + if err := p.configStore.DeleteMCPPerUserHeaderCredential(ctx, id); err != nil { + return fmt.Errorf("delete mcp per-user header credential: %w", err) + } + return nil +} + +// InitiateUserSubmissionFlow creates a pending mcp_per_user_header_flows +// row keyed by (mode, identity, mcp_client_id), mints a +// mcp_headers_auth temp-token bound to the new row's ID, and returns the +// auth-page URL with the token embedded as a `#t=` fragment. +// Mirrors oauth2.OAuth2Provider.InitiateUserOAuthFlow. +// +// The temp-token mint is best-effort: if it fails, the URL is returned +// without a fragment and remains usable for callers already authenticated +// to the dashboard. The behavior is the same as the OAuth equivalent so +// the two surfaces feel uniform. +func (p *Provider) InitiateUserSubmissionFlow(ctx context.Context, mode schemas.MCPAuthMode, identity, mcpClientID, baseURL string) (*schemas.MCPHeadersFlowInitiation, error) { + if p.configStore == nil { + return nil, schemas.ErrHeadersCredentialProviderNotAvailable + } + if strings.TrimSpace(mcpClientID) == "" { + return nil, errors.New("mcp_client_id is required") + } + if strings.TrimSpace(identity) == "" { + return nil, errors.New("identity is required to initiate per-user-headers submission flow") + } + if strings.TrimSpace(baseURL) == "" { + return nil, errors.New("base URL is required to build the submission auth page URL") + } + + // Single canonical lookup: at most one pending row per (mode, identity, + // mcp_client). If a pending row already exists, refresh its ExpiresAt + // in place rather than spawning a duplicate (which a user clicking + // "Edit values" repeatedly would otherwise produce). Mirrors + // oauth2.InitiateUserOAuthFlow's find-or-update pattern. + existing, lookupErr := p.configStore.GetMCPPerUserHeaderFlowByModeIdentityAndMCPClient(ctx, mode, identity, mcpClientID) + if lookupErr != nil { + return nil, fmt.Errorf("look up existing header flow: %w", lookupErr) + } + + now := time.Now() + var flow *tables.TableMCPPerUserHeaderFlow + if existing != nil && existing.Status == "pending" { + // Re-init path: keep the same row, rotate the expiry. + existing.ExpiresAt = now.Add(SubmissionFlowTTL) + existing.UpdatedAt = now + if err := p.configStore.UpdateMCPPerUserHeaderFlow(ctx, existing); err != nil { + return nil, fmt.Errorf("update mcp per-user header flow: %w", err) + } + flow = existing + } else { + // Fresh row. Exactly one of (UserID, VirtualKeyID, SessionID) is + // populated based on mode — same convention as TableOauthUserSession + // so the sessions UI rendering stays uniform. + flow = &tables.TableMCPPerUserHeaderFlow{ + ID: uuid.NewString(), + MCPClientID: mcpClientID, + FlowMode: string(mode), + Status: "pending", + ExpiresAt: now.Add(SubmissionFlowTTL), + CreatedAt: now, + UpdatedAt: now, + } + switch mode { + case schemas.MCPAuthModeUser: + v := identity + flow.UserID = &v + case schemas.MCPAuthModeVK: + v := identity + flow.VirtualKeyID = &v + case schemas.MCPAuthModeSession: + flow.SessionID = identity + default: + return nil, fmt.Errorf("unknown auth mode for headers flow: %s", mode) + } + if err := p.configStore.CreateMCPPerUserHeaderFlow(ctx, flow); err != nil { + return nil, fmt.Errorf("create mcp per-user header flow: %w", err) + } + } + + // Build the frontend URL: {base}/workspace/mcp-sessions/auth?flow={id}&kind=headers. + // The auth page hosts both per-user-OAuth and per-user-headers flows on + // the same URL pattern; `kind=headers` tells the page to call the + // per-user-headers flow APIs instead of the OAuth ones. The OAuth + // counterpart omits the `kind` param (default branch). + frontendURL := strings.TrimRight(baseURL, "/") + "/workspace/mcp-sessions/auth?flow=" + flow.ID + "&kind=headers" + + // Mint a mcp_headers_auth temp token bound to this flow's row ID and + // embed it in the URL as a fragment so a browser hitting the auth page + // without a dashboard session can still call the per-user submission + // endpoints. The fragment never leaves the browser (not in server + // logs, not in upstream Referer), unlike a query param. Best-effort: + // mint failure does not fail the flow init. + p.mu.RLock() + tempTokens := p.tempTokens + p.mu.RUnlock() + if tempTokens != nil { + ttl := time.Until(flow.ExpiresAt) + if ttl > 0 { + plaintext, mintErr := tempTokens.Mint(ctx, temptoken.MCPHeadersAuthScopeName, flow.ID, ttl) + if mintErr != nil { + p.logger.Warn("Failed to mint mcp_headers_auth temp token for flow %s: %v (link still usable for dashboard-authenticated callers)", flow.ID, mintErr) + } else { + frontendURL = frontendURL + "#t=" + plaintext + } + } + } + + return &schemas.MCPHeadersFlowInitiation{ + FlowID: flow.ID, + FrontendURL: frontendURL, + ExpiresAt: flow.ExpiresAt, + }, nil +} + +// rowToCredential converts the gorm row into the in-memory schema view, +// decrypting and deserializing HeadersJSON in the process. +func rowToCredential(row *tables.TableMCPPerUserHeaderCredential) (*schemas.MCPHeadersUserCredential, error) { + headers, err := row.GetHeaders() + if err != nil { + return nil, err + } + return &schemas.MCPHeadersUserCredential{ + ID: row.ID, + MCPClientID: row.MCPClientID, + AuthMode: schemas.MCPAuthMode(row.AuthMode), + UserID: row.UserID, + VirtualKeyID: row.VirtualKeyID, + SessionID: nilIfEmpty(row.SessionID), + Headers: headers, + Status: schemas.MCPHeadersUserCredentialStatus(row.Status), + CreatedAt: row.CreatedAt, + UpdatedAt: row.UpdatedAt, + }, nil +} + +// credentialToRow builds a fresh table row from the in-memory credential. +// Sets timestamps when zero so a re-upsert behaves like an update. +func credentialToRow(cred *schemas.MCPHeadersUserCredential) (*tables.TableMCPPerUserHeaderCredential, error) { + status := string(cred.Status) + if status == "" { + status = string(schemas.MCPHeadersUserCredentialStatusActive) + } + row := &tables.TableMCPPerUserHeaderCredential{ + ID: cred.ID, + MCPClientID: cred.MCPClientID, + AuthMode: string(cred.AuthMode), + UserID: cred.UserID, + VirtualKeyID: cred.VirtualKeyID, + Status: status, + CreatedAt: cred.CreatedAt, + UpdatedAt: cred.UpdatedAt, + } + if cred.SessionID != nil { + row.SessionID = *cred.SessionID + } + if row.CreatedAt.IsZero() { + row.CreatedAt = time.Now() + } + row.UpdatedAt = time.Now() + if err := row.SetHeaders(cred.Headers); err != nil { + return nil, err + } + return row, nil +} + +func nilIfEmpty(s string) *string { + if s == "" { + return nil + } + return &s +} diff --git a/framework/mcp_headers/sweep.go b/framework/mcp_headers/sweep.go new file mode 100644 index 0000000000..4d50c68d20 --- /dev/null +++ b/framework/mcp_headers/sweep.go @@ -0,0 +1,139 @@ +package mcp_headers + +import ( + "context" + "sync" + "time" + + "github.com/maximhq/bifrost/core/schemas" +) + +// CredentialSweepWorker periodically purges stale per-user header credentials +// AND expired pending submission-flow rows. Two independent sweeps run on the +// same goroutine: +// +// - Orphan credential sweep: rows that have been in 'orphaned' state +// (VK lost access) longer than OrphanRetention are hard-deleted. +// - Expired flow sweep: pending flow rows whose ExpiresAt has passed +// (caller never completed the submission, link aged out) are +// hard-deleted. Tighter cadence than the orphan sweep because flow rows +// are short-lived (15 min TTL). +// +// Mirrors oauth2.PerUserOAuthSweepWorker, which combines the same two +// concerns on the OAuth side. +// +// Defaults: 24h orphan cadence, 15 min flow-expiry cadence. Non-positive +// orphanRetention disables the orphan sweep entirely; the expired-flow sweep +// always runs because flow rows have no semantic value past their expiry. +type CredentialSweepWorker struct { + provider *Provider + orphanSweepEvery time.Duration + orphanRetention time.Duration + expiredFlowEvery time.Duration + stopCh chan struct{} + stopOnce sync.Once + logger schemas.Logger +} + +// NewCredentialSweepWorker creates a sweep worker with sensible defaults. +// orphanRetention <= 0 disables the orphan sweep (the worker still starts but +// the tick is a no-op — keeps wiring uniform). +func NewCredentialSweepWorker(provider *Provider, orphanRetention time.Duration, logger schemas.Logger) *CredentialSweepWorker { + if provider == nil || provider.configStore == nil { + if logger != nil { + logger.Warn("per-user headers credential sweep worker not started: provider or config store is nil") + } + return nil + } + return &CredentialSweepWorker{ + provider: provider, + orphanSweepEvery: 24 * time.Hour, + orphanRetention: orphanRetention, + expiredFlowEvery: 15 * time.Minute, + stopCh: make(chan struct{}), + logger: logger, + } +} + +// Start begins the sweep worker in a background goroutine. +func (w *CredentialSweepWorker) Start(ctx context.Context) { + go w.run(ctx) + if w.logger != nil { + w.logger.Info("Per-user headers sweep worker started (orphan=%s, retention=%s, expired_flow=%s)", + w.orphanSweepEvery, w.orphanRetention, w.expiredFlowEvery) + } +} + +// Stop gracefully stops the sweep worker. sync.Once guards against double-close +// panics from redundant shutdown paths. +func (w *CredentialSweepWorker) Stop() { + w.stopOnce.Do(func() { + close(w.stopCh) + if w.logger != nil { + w.logger.Info("Per-user headers credential sweep worker stopped") + } + }) +} + +func (w *CredentialSweepWorker) run(ctx context.Context) { + orphanTicker := time.NewTicker(w.orphanSweepEvery) + defer orphanTicker.Stop() + expiredFlowTicker := time.NewTicker(w.expiredFlowEvery) + defer expiredFlowTicker.Stop() + + // Run once on start so a deploy doesn't have to wait a full interval. + w.sweepOrphanedCredentials(ctx) + w.sweepExpiredFlows(ctx) + + for { + select { + case <-orphanTicker.C: + w.sweepOrphanedCredentials(ctx) + case <-expiredFlowTicker.C: + w.sweepExpiredFlows(ctx) + case <-w.stopCh: + return + case <-ctx.Done(): + return + } + } +} + +func (w *CredentialSweepWorker) sweepOrphanedCredentials(ctx context.Context) { + if w.orphanRetention <= 0 { + return + } + n, err := w.provider.configStore.DeleteOrphanedMCPPerUserHeaderCredentials(ctx, w.orphanRetention) + if err != nil { + if w.logger != nil { + w.logger.Error("per-user headers orphan sweep failed: %v", err) + } + return + } + if n > 0 && w.logger != nil { + w.logger.Info("per-user headers orphan sweep removed %d rows older than %s", n, w.orphanRetention) + } +} + +func (w *CredentialSweepWorker) sweepExpiredFlows(ctx context.Context) { + n, err := w.provider.configStore.DeleteExpiredMCPPerUserHeaderFlows(ctx) + if err != nil { + if w.logger != nil { + w.logger.Error("per-user headers expired-flow sweep failed: %v", err) + } + return + } + if n > 0 && w.logger != nil { + w.logger.Info("per-user headers expired-flow sweep removed %d rows", n) + } +} + +// SetOrphanSweepInterval updates the orphan-sweep cadence (for testing). +// Non-positive durations are ignored — run() feeds the field straight into +// time.NewTicker, which panics on d <= 0. +func (w *CredentialSweepWorker) SetOrphanSweepInterval(d time.Duration) { + if d <= 0 { + return + } + w.orphanSweepEvery = d +} diff --git a/framework/temptoken/scope.go b/framework/temptoken/scope.go index 490fedbc78..1e3a7c3ed8 100644 --- a/framework/temptoken/scope.go +++ b/framework/temptoken/scope.go @@ -17,6 +17,12 @@ const ( // auth page to call the per-user flow endpoints. Bound resource_id is the // OAuth flow ID. See transports/bifrost-http/handlers/temp_token_scopes.go. MCPAuthScopeName = "mcp_auth" + + // MCPHeadersAuthScopeName names the scope that authorizes the MCP + // per-user-headers auth page to call the per-user-headers submission + // flow endpoints. Bound resource_id is the headers flow ID. Parallel + // of MCPAuthScopeName for the per-user-headers surface. + MCPHeadersAuthScopeName = "mcp_headers_auth" ) // RoutePattern is one (method, path) pair a Scope grants access to. The path diff --git a/plugins/governance/main.go b/plugins/governance/main.go index 152993e2f2..5658a19f81 100644 --- a/plugins/governance/main.go +++ b/plugins/governance/main.go @@ -1700,6 +1700,60 @@ func (p *GovernancePlugin) PostMCPHook(ctx *schemas.BifrostContext, resp *schema return resp, bifrostErr, nil } +// PreMCPConnectionHook resolves the caller's identity onto the BifrostContext +// before the connect-plugin gate releases control to the credential-store +// resolver. This is the only point in the MCP connect lifecycle where we can +// turn the raw x-bf-vk header into the resolved VK row ID — anything later +// (PreMCPHook / PostMCPHook) runs after the resolver has already needed that +// row ID, and per-user auth types (per_user_oauth, per_user_headers) key +// their stored credentials by it. +// +// The hook is intentionally narrow: it ONLY populates the identity context +// keys (VK row ID, name, team / customer fan-out). Policy checks (budget, +// rate limit, tool allow-list) stay on PreMCPHook for the actual CallTool — +// Connect is transport setup, not the gated operation. +// +// No short-circuit returned even when the VK isn't recognized: bad-VK +// rejection belongs on the tool-call path so the caller gets a stable +// error format. An unknown VK here simply leaves the row ID empty, and the +// resolver will surface the "requires an identity" error itself. +func (p *GovernancePlugin) PreMCPConnectionHook(ctx *schemas.BifrostContext, req *schemas.BifrostMCPConnectRequest) (*schemas.BifrostMCPConnectRequest, *schemas.MCPConnectionShortCircuit, error) { + virtualKeyValue := bifrost.GetStringFromContext(ctx, schemas.BifrostContextKeyVirtualKey) + if virtualKeyValue == "" { + return req, nil, nil + } + vk, ok := p.store.GetVirtualKey(ctx, virtualKeyValue) + if !ok || vk == nil { + // Unknown VK — leave identity unset; the resolver will surface the + // appropriate error on the per-user auth path. For shared-connection + // auth types this is a no-op (they don't read these keys). + return req, nil, nil + } + ctx.SetValue(schemas.BifrostContextKeyGovernanceVirtualKeyID, vk.ID) + ctx.SetValue(schemas.BifrostContextKeyGovernanceVirtualKeyName, vk.Name) + if vk.Team != nil { + ctx.SetValue(schemas.BifrostContextKeyGovernanceTeamID, vk.Team.ID) + ctx.SetValue(schemas.BifrostContextKeyGovernanceTeamName, vk.Team.Name) + if vk.Team.Customer != nil { + ctx.SetValue(schemas.BifrostContextKeyGovernanceCustomerID, vk.Team.Customer.ID) + ctx.SetValue(schemas.BifrostContextKeyGovernanceCustomerName, vk.Team.Customer.Name) + } + } + if vk.Customer != nil { + ctx.SetValue(schemas.BifrostContextKeyGovernanceCustomerID, vk.Customer.ID) + ctx.SetValue(schemas.BifrostContextKeyGovernanceCustomerName, vk.Customer.Name) + } + return req, nil, nil +} + +// PostMCPConnectionHook is a pass-through; the identity resolution that +// PreMCPConnectionHook performs is observation-only and has no post-connect +// cleanup. Implementing this satisfies MCPConnectionPlugin so the typed +// PreMCPConnectionHook is dispatched by the plugin pipeline. +func (p *GovernancePlugin) PostMCPConnectionHook(ctx *schemas.BifrostContext, resp *schemas.BifrostMCPConnectResponse, bifrostErr *schemas.BifrostError) (*schemas.BifrostMCPConnectResponse, *schemas.BifrostError, error) { + return resp, bifrostErr, nil +} + // Cleanup shuts down all components gracefully func (p *GovernancePlugin) Cleanup() error { var cleanupErr error diff --git a/transports/bifrost-http/handlers/mcp.go b/transports/bifrost-http/handlers/mcp.go index bec5d87276..aa095e9c37 100644 --- a/transports/bifrost-http/handlers/mcp.go +++ b/transports/bifrost-http/handlers/mcp.go @@ -36,6 +36,11 @@ type MCPManager interface { // VerifyPerUserOAuthConnection verifies an MCP server using a temporary access // token and discovers available tools. The connection is closed after verification. VerifyPerUserOAuthConnection(ctx context.Context, config *schemas.MCPClientConfig, accessToken string) (map[string]schemas.ChatTool, map[string]string, error) + // VerifyHeadersConnection verifies an MCP server using a caller-supplied set + // of header values (admin sample or user-submitted) and discovers available + // tools. The connection is closed after verification. Mirrors + // VerifyPerUserOAuthConnection's role for MCPAuthTypePerUserHeaders. + VerifyHeadersConnection(ctx context.Context, config *schemas.MCPClientConfig, userHeaders map[string]string) (map[string]schemas.ChatTool, map[string]string, error) // SetClientTools updates the tool map for an existing client. SetClientTools(clientID string, tools map[string]schemas.ChatTool, toolNameMapping map[string]string) } @@ -223,6 +228,7 @@ func (h *MCPHandler) getMCPClientsPaginated(ctx *fasthttp.RequestCtx, limitStr, ToolPricing: dbClient.ToolPricing, AllowOnAllVirtualKeys: dbClient.AllowOnAllVirtualKeys, Disabled: dbClient.Disabled, + PerUserHeaderKeys: dbClient.PerUserHeaderKeys, } // Populate oauth client credentials from pre-fetched batch if dbClient.OauthConfigID != nil { @@ -322,10 +328,17 @@ type OAuthConfigRequest struct { Scopes []string `json:"scopes"` } -// MCPClientRequest represents the full MCP client creation request with OAuth support +// MCPClientRequest represents the full MCP client creation request with OAuth support. +// +// UserHeaders carries a sample set of per-user-headers values used only for +// upstream verification + tool discovery during create. Mirrors the per-user +// OAuth flow where the admin's temp access token is used the same way: the +// server runs discovery, attaches DiscoveredTools to the persisted config, +// and discards the credentials. Ignored for non-per_user_headers auth types. type MCPClientRequest struct { configstoreTables.TableMCPClient OauthConfig *OAuthConfigRequest `json:"oauth_config,omitempty"` + UserHeaders map[string]string `json:"user_headers,omitempty"` } // MCPVKConfigRequest represents a per-VK tool access config for an MCP client @@ -382,10 +395,107 @@ func (h *MCPHandler) addMCPClient(ctx *fasthttp.RequestCtx) { return } + // Handle per-user headers: admin declares the required key names (schema) + // AND supplies a sample set of values inline so the server can verify + // upstream + discover tools in a single round-trip. Mirrors the per-user + // OAuth flow exactly — the sample values are used once for verification + // and discarded (never persisted); each end-user submits their own values + // later via the inline-401 flow. + if req.AuthType == string(schemas.MCPAuthTypePerUserHeaders) { + if len(req.PerUserHeaderKeys) == 0 { + SendError(ctx, fasthttp.StatusBadRequest, "per_user_header_keys must be a non-empty list when auth_type is 'per_user_headers'") + return + } + normalisedHeaders := make([]string, 0, len(req.PerUserHeaderKeys)) + for i, key := range req.PerUserHeaderKeys { + if strings.TrimSpace(key) == "" { + SendError(ctx, fasthttp.StatusBadRequest, fmt.Sprintf("per_user_header_keys[%d] is empty", i)) + return + } + normalisedHeaders = append(normalisedHeaders, strings.ToLower(strings.TrimSpace(key))) + } + // HTTP header names are case-insensitive on the wire — reject duplicates + // like ["X-Api-Key", "x-api-key"] so downstream change-detection and + // credential storage stay correct. + if lib.HasDuplicates(normalisedHeaders) { + SendError(ctx, fasthttp.StatusBadRequest, "per_user_header_keys contains duplicate entries") + return + } + if missing := missingPerUserHeaderValues(req.PerUserHeaderKeys, req.UserHeaders); len(missing) > 0 { + SendError(ctx, fasthttp.StatusBadRequest, fmt.Sprintf("sample user_headers missing values for required keys: %s", strings.Join(missing, ", "))) + return + } + + toolSyncInterval := mcp.DefaultToolSyncInterval + if req.ToolSyncInterval != 0 { + toolSyncInterval = time.Duration(req.ToolSyncInterval) * time.Minute + } else { + config, cfgErr := h.store.ConfigStore.GetClientConfig(ctx) + if cfgErr == nil && config != nil { + toolSyncInterval = time.Duration(config.MCPToolSyncInterval) * time.Minute + } + } + + isPingAvailable := true + if req.IsPingAvailable != nil { + isPingAvailable = *req.IsPingAvailable + } + + schemasConfig := &schemas.MCPClientConfig{ + ID: req.ClientID, + Name: req.Name, + IsCodeModeClient: req.IsCodeModeClient, + IsPingAvailable: &isPingAvailable, + ToolSyncInterval: toolSyncInterval, + ConnectionType: schemas.MCPConnectionType(req.ConnectionType), + ConnectionString: req.ConnectionString, + StdioConfig: req.StdioConfig, + AuthType: schemas.MCPAuthTypePerUserHeaders, + PerUserHeaderKeys: req.PerUserHeaderKeys, + ToolsToExecute: req.ToolsToExecute, + ToolsToAutoExecute: req.ToolsToAutoExecute, + ToolPricing: req.ToolPricing, + Headers: req.Headers, + AllowedExtraHeaders: req.AllowedExtraHeaders, + AllowOnAllVirtualKeys: req.AllowOnAllVirtualKeys, + } + + // Verify connection and discover tools using the admin's sample + // header values. Discovered tools land on schemasConfig before we + // persist so the DB row includes them from the start — same + // convention as the per-user OAuth branch below. + tools, toolNameMapping, verifyErr := h.mcpManager.VerifyHeadersConnection(ctx, schemasConfig, req.UserHeaders) + if verifyErr != nil { + SendError(ctx, fasthttp.StatusUnprocessableEntity, fmt.Sprintf("Verification failed: %v", verifyErr)) + return + } + schemasConfig.DiscoveredTools = tools + schemasConfig.DiscoveredToolNameMapping = toolNameMapping + + if err := h.store.ConfigStore.CreateMCPClientConfig(ctx, schemasConfig); err != nil { + SendError(ctx, fasthttp.StatusInternalServerError, fmt.Sprintf("Failed to create MCP config: %v", err)) + return + } + if err := h.mcpManager.AddMCPClient(ctx, schemasConfig); err != nil { + if delErr := h.store.ConfigStore.DeleteMCPClientConfig(ctx, schemasConfig.ID); delErr != nil { + logger.Error(fmt.Sprintf("Failed to roll back MCP client config after AddMCPClient failure: %v", delErr)) + } + SendError(ctx, fasthttp.StatusInternalServerError, fmt.Sprintf("Failed to register MCP client: %v", err)) + return + } + h.mcpManager.SetClientTools(schemasConfig.ID, tools, toolNameMapping) + + SendJSON(ctx, map[string]any{ + "status": "success", + "message": fmt.Sprintf("MCP client registered. %d tools discovered. Each user will submit their own headers on first tool use.", len(tools)), + }) + return + } + // Handle per-user OAuth: admin does a test OAuth login to verify the configuration. // Uses the same pending_oauth pattern as server-level OAuth, but on completion we // verify the connection, discover tools, save the client, and discard the admin's token. - if req.AuthType == "per_user_oauth" { + if req.AuthType == string(schemas.MCPAuthTypePerUserOauth) { if req.OauthConfig == nil { SendError(ctx, fasthttp.StatusBadRequest, "OAuth configuration is required when auth_type is 'per_user_oauth'") return @@ -465,7 +575,7 @@ func (h *MCPHandler) addMCPClient(ctx *fasthttp.RequestCtx) { } // Check if server-level OAuth flow is needed - if req.AuthType == "oauth" { + if req.AuthType == string(schemas.MCPAuthTypeOauth) { if req.OauthConfig == nil { SendError(ctx, fasthttp.StatusBadRequest, "OAuth configuration is required when auth_type is 'oauth'") return @@ -710,6 +820,38 @@ func (h *MCPHandler) updateMCPClient(ctx *fasthttp.RequestCtx) { SendError(ctx, fasthttp.StatusBadRequest, fmt.Sprintf("Invalid allowed_extra_headers: %v", err)) return } + // Validate per_user_header_keys only when the request explicitly provides + // the field — otherwise resolvePerUserHeaderKeys carries the existing list + // forward unchanged (already validated at create time). + if req.PerUserHeaderKeys != nil { + // Reject an explicit empty list for per_user_headers clients. + // AuthType is immutable on update (enforced at clientmanager.go:911), + // so existingConfig.AuthType is the reliable gate — clients on other + // auth types may legitimately carry no per_user_header_keys, but for + // per_user_headers an empty schema means the auth mode has nothing + // to collect or validate, which violates the feature contract. + // Without this guard, resolvePerUserHeaderKeys returns [] and the + // resolver errors on every subsequent tool call with "no PerUser- + // HeaderKeys declared" — and MarkMCPPerUserHeaderCredentialsNeedsUpdate + // fires first, flipping all active credentials to needs_update for + // nothing. + if existingConfig.AuthType == schemas.MCPAuthTypePerUserHeaders && len(req.PerUserHeaderKeys) == 0 { + SendError(ctx, fasthttp.StatusBadRequest, "per_user_header_keys must be a non-empty list for per_user_headers clients") + return + } + canonHeaderKeys := make([]string, 0, len(req.PerUserHeaderKeys)) + for i, key := range req.PerUserHeaderKeys { + if strings.TrimSpace(key) == "" { + SendError(ctx, fasthttp.StatusBadRequest, fmt.Sprintf("per_user_header_keys[%d] is empty", i)) + return + } + canonHeaderKeys = append(canonHeaderKeys, strings.ToLower(strings.TrimSpace(key))) + } + if lib.HasDuplicates(canonHeaderKeys) { + SendError(ctx, fasthttp.StatusBadRequest, "per_user_header_keys contains duplicate entries") + return + } + } // OAuth credential rotation is temporarily disabled. if req.OauthConfig != nil { @@ -868,6 +1010,19 @@ func (h *MCPHandler) updateMCPClient(ctx *fasthttp.RequestCtx) { ToolPricing: req.ToolPricing, AllowOnAllVirtualKeys: req.AllowOnAllVirtualKeys, Disabled: req.Disabled, + PerUserHeaderKeys: resolvePerUserHeaderKeys(existingConfig, req), + } + + // If the per-user-headers schema changed, flip every existing active row + // to 'needs_update' so callers are forced to resubmit on next tool use. + // The rows are preserved (status only flips) so the submission UI can + // prefill known values. + if existingConfig.AuthType == schemas.MCPAuthTypePerUserHeaders && + perUserHeaderKeysChanged(existingConfig.PerUserHeaderKeys, schemasConfig.PerUserHeaderKeys) && + h.store.ConfigStore != nil { + if err := h.store.ConfigStore.MarkMCPPerUserHeaderCredentialsNeedsUpdate(ctx, existingConfig.ID); err != nil { + logger.Error(fmt.Sprintf("failed to flip per-user header credentials to needs_update for client %s: %v", existingConfig.ID, err)) + } } // Update MCP client config in memory (always — applies name/tools/header changes, @@ -1455,3 +1610,39 @@ func (h *MCPHandler) completeMCPClientOAuth(ctx *fasthttp.RequestCtx) { } SendJSON(ctx, map[string]any{"status": "success", "message": message}) } + +// resolvePerUserHeaderKeys returns the per-user-header-key list to persist on +// the updated MCP client. If the request explicitly sets the field (even to +// an empty list when the caller is removing all keys), the request wins; +// otherwise the existing schema is preserved. +func resolvePerUserHeaderKeys(existing *schemas.MCPClientConfig, req MCPClientUpdateRequest) []string { + if req.PerUserHeaderKeys != nil { + return req.PerUserHeaderKeys + } + if existing != nil { + return existing.PerUserHeaderKeys + } + return nil +} + +// perUserHeaderKeysChanged reports whether the new key set differs from the +// old set (order-insensitive). Used by updateMCPClient to decide whether to +// flip existing user credentials to 'needs_update'. +func perUserHeaderKeysChanged(oldKeys, newKeys []string) bool { + if len(oldKeys) != len(newKeys) { + return true + } + if len(oldKeys) == 0 { + return false + } + seen := make(map[string]struct{}, len(oldKeys)) + for _, k := range oldKeys { + seen[k] = struct{}{} + } + for _, k := range newKeys { + if _, ok := seen[k]; !ok { + return true + } + } + return false +} diff --git a/transports/bifrost-http/handlers/mcp_per_user_headers.go b/transports/bifrost-http/handlers/mcp_per_user_headers.go new file mode 100644 index 0000000000..82dbe64028 --- /dev/null +++ b/transports/bifrost-http/handlers/mcp_per_user_headers.go @@ -0,0 +1,379 @@ +package handlers + +import ( + "context" + "encoding/json" + "fmt" + "sort" + "strings" + "time" + + "github.com/fasthttp/router" + "github.com/google/uuid" + "github.com/maximhq/bifrost/core/schemas" + "github.com/maximhq/bifrost/framework/temptoken" + "github.com/maximhq/bifrost/transports/bifrost-http/lib" + "github.com/valyala/fasthttp" +) + +// MCPPerUserHeadersHandler exposes the per-user-headers flow-detail, +// flow-submit, and credential-revoke endpoints. It is the storage-side +// companion to the inline-401 MCPAuthRequiredError surfaced by the +// credstore resolver. +// +// Identity scoping: the flow-bound endpoints read (mode, identity) from +// the flow row itself (created server-side by the resolver), so per-user +// submissions are tied to the same identity that triggered the +// inline-401. Admin-side verification + tool discovery during MCP client +// creation happens in the unified POST /api/mcp/client handler — mirrors +// the per-user OAuth shape. +type MCPPerUserHeadersHandler struct { + store *lib.Config + mcpManager MCPManager + tempTokens *temptoken.Service // optional — set when a temp-token service is wired; flowSubmit uses it to revoke the bound token on completion +} + +// NewMCPPerUserHeadersHandler constructs the handler. tempTokens is optional — +// pass nil if the deployment does not run the temp-token service (the +// auth-page URL will then only be usable from a dashboard-authenticated +// browser session). +func NewMCPPerUserHeadersHandler(mcpManager MCPManager, store *lib.Config, tempTokens *temptoken.Service) *MCPPerUserHeadersHandler { + return &MCPPerUserHeadersHandler{ + store: store, + mcpManager: mcpManager, + tempTokens: tempTokens, + } +} + +// RegisterRoutes mounts the per-user-headers routes. +// +// Flow-id-bound endpoints mirror the per-user OAuth surface +// (/api/oauth/per-user/flows/{id}): a pending flow row is created when the +// resolver surfaces the inline-401, its ID rides in the auth-page URL as +// ?flow=, and a temp-token (mcp_headers_auth scope) bound to that +// flow ID is appended as a #t= fragment so anonymous browser +// visitors can complete the submission without a dashboard session. +// +// The DELETE-by-credential-ID route lives under /credential/{id} to +// disambiguate from the flow-ID-keyed routes. +func (h *MCPPerUserHeadersHandler) RegisterRoutes(r *router.Router, middlewares ...schemas.BifrostHTTPMiddleware) { + r.GET("/api/mcp/per-user-headers/flows/{id}", lib.ChainMiddlewares(h.flowDetail, middlewares...)) + r.PUT("/api/mcp/per-user-headers/flows/{id}", lib.ChainMiddlewares(h.flowSubmit, middlewares...)) + r.DELETE("/api/mcp/per-user-headers/credential/{id}", lib.ChainMiddlewares(h.revoke, middlewares...)) +} + +// mcpHeadersFlowDetailResponse is the wire shape for +// GET /api/mcp/per-user-headers/flows/{id}. Mirrors mcpFlowDetailResponse on +// the OAuth side: identity columns + MCP client summary + the schema +// (required keys, admin key names) the submission UI needs to render. +type mcpHeadersFlowDetailResponse struct { + ID string `json:"id"` + FlowMode string `json:"flow_mode"` + Status string `json:"status"` + MCPClient *mcpClientSummary `json:"mcp_client,omitempty"` + UserID *string `json:"user_id,omitempty"` + User *userSummary `json:"user,omitempty"` + VirtualKey *virtualKeySummary `json:"virtual_key,omitempty"` + SessionID *string `json:"session_id,omitempty"` + ExpiresAt string `json:"expires_at"` + CreatedAt string `json:"created_at"` + RequiredHeaderKeys []string `json:"required_header_keys"` + AdminHeaderKeys []string `json:"admin_header_keys,omitempty"` + SubmittedKeys []string `json:"submitted_keys,omitempty"` // Names of keys already on the active credential (no values) + HasActiveCredential bool `json:"has_active_credential"` +} + +// flowDetail returns the pending headers-submission flow row's metadata so +// the auth landing page can render the form. Authorization is via either a +// dashboard session (caller is signed in and DAC-scoped) OR the +// mcp_headers_auth temp token bound to {id} (anonymous browser visitor that +// followed the auth-page URL from a Bifrost API error response). +func (h *MCPPerUserHeadersHandler) flowDetail(ctx *fasthttp.RequestCtx) { + flowID, ok := ctx.UserValue("id").(string) + if !ok || strings.TrimSpace(flowID) == "" { + SendError(ctx, fasthttp.StatusBadRequest, "Invalid flow id") + return + } + flow, err := h.store.ConfigStore.GetMCPPerUserHeaderFlowByID(ctx, flowID) + if err != nil { + SendError(ctx, fasthttp.StatusInternalServerError, fmt.Sprintf("Failed to load flow: %v", err)) + return + } + if flow == nil { + SendError(ctx, fasthttp.StatusNotFound, "Headers submission flow not found") + return + } + + config, cfgErr := h.loadMCPClientConfig(ctx, flow.MCPClientID) + if cfgErr != nil { + SendError(ctx, fasthttp.StatusInternalServerError, cfgErr.Error()) + return + } + + resp := mcpHeadersFlowDetailResponse{ + ID: flow.ID, + FlowMode: flow.FlowMode, + Status: flow.Status, + UserID: flow.UserID, + ExpiresAt: flow.ExpiresAt.UTC().Format(rfc3339Nano), + CreatedAt: flow.CreatedAt.UTC().Format(rfc3339Nano), + RequiredHeaderKeys: append([]string(nil), config.PerUserHeaderKeys...), + AdminHeaderKeys: headerNamesFromConfig(config), + } + if flow.MCPClient != nil { + resp.MCPClient = &mcpClientSummary{ClientID: flow.MCPClient.ClientID, Name: flow.MCPClient.Name} + } else { + resp.MCPClient = &mcpClientSummary{ClientID: flow.MCPClientID} + } + if flow.VirtualKey != nil { + resp.VirtualKey = &virtualKeySummary{ID: flow.VirtualKey.ID, Name: flow.VirtualKey.Name} + } else if flow.VirtualKeyID != nil { + resp.VirtualKey = &virtualKeySummary{ID: *flow.VirtualKeyID} + } + if flow.User != nil { + resp.User = &userSummary{ID: flow.User.ID, Name: flow.User.Name} + } + if flow.FlowMode == string(schemas.MCPAuthModeSession) && flow.SessionID != "" { + sid := flow.SessionID + resp.SessionID = &sid + } + + // Surface a "this credential already exists" hint so the UI can render an + // edit affordance instead of a fresh form. Identity is the flow row's own + // identity column — same convention as OAuth's flowDetail. + if h.store.MCPHeadersProvider != nil { + mode := schemas.MCPAuthMode(flow.FlowMode) + identity := "" + switch mode { + case schemas.MCPAuthModeUser: + if flow.UserID != nil { + identity = *flow.UserID + } + case schemas.MCPAuthModeVK: + if flow.VirtualKeyID != nil { + identity = *flow.VirtualKeyID + } + case schemas.MCPAuthModeSession: + identity = flow.SessionID + } + if identity != "" { + // GetCredentialByMode returns 'active' and 'needs_update' rows + // (orphaned is filtered at the store). SubmittedKeys carries the + // previously-submitted key NAMES — useful regardless of status + // so the auth-page can render a "Previously submitted" hint even + // after an admin schema change has flipped the row. HasActive- + // Credential gates strictly on the row's lifecycle Status: a + // needs_update row is NOT an active credential and the UI + // shouldn't show "you're editing your existing credential" + // copy for it (the user must resubmit, not edit-in-place). + if cred, lookupErr := h.store.MCPHeadersProvider.GetCredentialByMode(ctx, mode, identity, flow.MCPClientID); lookupErr == nil && cred != nil { + resp.HasActiveCredential = cred.Status == schemas.MCPHeadersUserCredentialStatusActive + resp.SubmittedKeys = sortedKeys(cred.Headers) + } + } + } + + SendJSON(ctx, resp) +} + +// flowSubmitRequest is the user-supplied set of header values to persist. +type flowSubmitRequest struct { + Headers map[string]string `json:"headers"` +} + +// flowSubmit consumes a pending headers-submission flow row: verifies the +// caller's values against the upstream, upserts the credential keyed by the +// flow row's (mode, identity), then deletes the flow row and the temp token +// bound to it. Mirrors the OAuth callback's "complete the flow" semantics. +func (h *MCPPerUserHeadersHandler) flowSubmit(ctx *fasthttp.RequestCtx) { + if h.store.MCPHeadersProvider == nil { + SendError(ctx, fasthttp.StatusServiceUnavailable, "per-user headers credential provider is not configured") + return + } + flowID, ok := ctx.UserValue("id").(string) + if !ok || strings.TrimSpace(flowID) == "" { + SendError(ctx, fasthttp.StatusBadRequest, "Invalid flow id") + return + } + var req flowSubmitRequest + if err := json.Unmarshal(ctx.PostBody(), &req); err != nil { + SendError(ctx, fasthttp.StatusBadRequest, fmt.Sprintf("Invalid request format: %v", err)) + return + } + + flow, err := h.store.ConfigStore.GetMCPPerUserHeaderFlowByID(ctx, flowID) + if err != nil { + SendError(ctx, fasthttp.StatusInternalServerError, fmt.Sprintf("Failed to load flow: %v", err)) + return + } + if flow == nil { + SendError(ctx, fasthttp.StatusNotFound, "Headers submission flow not found") + return + } + if !flow.ExpiresAt.IsZero() && flow.ExpiresAt.Before(time.Now()) { + SendError(ctx, fasthttp.StatusGone, "Headers submission flow has expired; restart from the API error link") + return + } + + config, cfgErr := h.loadMCPClientConfig(ctx, flow.MCPClientID) + if cfgErr != nil { + SendError(ctx, fasthttp.StatusInternalServerError, cfgErr.Error()) + return + } + if missing := missingPerUserHeaderValues(config.PerUserHeaderKeys, req.Headers); len(missing) > 0 { + SendError(ctx, fasthttp.StatusBadRequest, fmt.Sprintf("missing values for required keys: %s", strings.Join(missing, ", "))) + return + } + + // Filter to declared keys only — extras get dropped on purpose so a stale + // UI cannot persist values that would never be sent on the wire. + filtered := make(map[string]string, len(config.PerUserHeaderKeys)) + for _, key := range config.PerUserHeaderKeys { + if v, ok := req.Headers[key]; ok { + filtered[key] = v + } + } + + if _, _, verifyErr := h.mcpManager.VerifyHeadersConnection(ctx, config, filtered); verifyErr != nil { + SendError(ctx, fasthttp.StatusUnprocessableEntity, fmt.Sprintf("Verification failed: %v", verifyErr)) + return + } + + mode := schemas.MCPAuthMode(flow.FlowMode) + cred := &schemas.MCPHeadersUserCredential{ + ID: uuid.New().String(), + MCPClientID: flow.MCPClientID, + AuthMode: mode, + Headers: filtered, + Status: schemas.MCPHeadersUserCredentialStatusActive, + } + switch mode { + case schemas.MCPAuthModeUser: + cred.UserID = flow.UserID + case schemas.MCPAuthModeVK: + cred.VirtualKeyID = flow.VirtualKeyID + case schemas.MCPAuthModeSession: + if flow.SessionID != "" { + s := flow.SessionID + cred.SessionID = &s + } + default: + SendError(ctx, fasthttp.StatusInternalServerError, fmt.Sprintf("Flow has unknown auth mode: %s", flow.FlowMode)) + return + } + + if upsertErr := h.store.MCPHeadersProvider.UpsertCredential(ctx, cred); upsertErr != nil { + SendError(ctx, fasthttp.StatusInternalServerError, fmt.Sprintf("Failed to store credential: %v", upsertErr)) + return + } + + // Best-effort: delete the consumed flow row and the temp token bound to + // its ID. Failure to clean up is not fatal — the sweep worker will + // collect both on the next pass. We log but don't surface to the caller. + if delErr := h.store.ConfigStore.DeleteMCPPerUserHeaderFlow(ctx, flow.ID); delErr != nil { + logger.Warn("[mcp/per-user-headers] failed to delete flow %s after successful submit: %v", flow.ID, delErr) + } + if h.tempTokens != nil { + if _, delErr := h.tempTokens.DeleteByResourceID(ctx, temptoken.MCPHeadersAuthScopeName, flow.ID); delErr != nil { + logger.Warn("[mcp/per-user-headers] failed to delete temp token for flow %s: %v", flow.ID, delErr) + } + } + + SendJSON(ctx, map[string]any{ + "status": "success", + "credential_id": cred.ID, + "updated_at": cred.UpdatedAt, + }) +} + +// revoke deletes a credential row by its primary key. Authorization is +// gated by a scoped GetCredentialByID lookup first: in enterprise the DAC +// scope filters the row out for callers that don't own it (returns nil → +// 404 here), so the unscoped delete that follows can never touch a row the +// caller couldn't see. Mirrors the sessions handler's revoke pattern. +func (h *MCPPerUserHeadersHandler) revoke(ctx *fasthttp.RequestCtx) { + if h.store.MCPHeadersProvider == nil { + SendError(ctx, fasthttp.StatusServiceUnavailable, "per-user headers credential provider is not configured") + return + } + if h.store.ConfigStore == nil { + SendError(ctx, fasthttp.StatusServiceUnavailable, "config store is not configured") + return + } + id, ok := ctx.UserValue("id").(string) + if !ok || strings.TrimSpace(id) == "" { + SendError(ctx, fasthttp.StatusBadRequest, "id is required") + return + } + cred, err := h.store.ConfigStore.GetMCPPerUserHeaderCredentialByID(ctx, id) + if err != nil { + SendError(ctx, fasthttp.StatusInternalServerError, fmt.Sprintf("Failed to load credential: %v", err)) + return + } + if cred == nil { + SendError(ctx, fasthttp.StatusNotFound, "credential not found") + return + } + if err := h.store.MCPHeadersProvider.DeleteCredential(ctx, cred.ID); err != nil { + SendError(ctx, fasthttp.StatusInternalServerError, fmt.Sprintf("Failed to revoke credential: %v", err)) + return + } + ctx.SetStatusCode(fasthttp.StatusNoContent) +} + +// loadMCPClientConfig fetches the MCP client config and verifies it is a +// per-user-headers client. Returns a typed error so the handler can pick the +// right HTTP status. +func (h *MCPPerUserHeadersHandler) loadMCPClientConfig(ctx context.Context, mcpClientID string) (*schemas.MCPClientConfig, error) { + if h.store.ConfigStore == nil { + return nil, fmt.Errorf("config store is not configured") + } + row, err := h.store.ConfigStore.GetMCPClientConfigByID(ctx, mcpClientID) + if err != nil { + return nil, fmt.Errorf("failed to load mcp client: %w", err) + } + if row == nil { + return nil, fmt.Errorf("mcp client %s not found", mcpClientID) + } + if row.AuthType != schemas.MCPAuthTypePerUserHeaders { + return nil, fmt.Errorf("mcp client %s is not configured for per-user headers auth", mcpClientID) + } + return row, nil +} + +// missingPerUserHeaderValues returns the names of any required key whose +// value is missing or empty in the supplied map. +func missingPerUserHeaderValues(required []string, values map[string]string) []string { + var missing []string + for _, key := range required { + if v, ok := values[key]; !ok || strings.TrimSpace(v) == "" { + missing = append(missing, key) + } + } + return missing +} + +// headerNamesFromConfig returns just the names (no values) of static admin +// headers on the MCP client config. Used by the submission UI to display +// context the user can't edit. +func headerNamesFromConfig(config *schemas.MCPClientConfig) []string { + if config == nil || len(config.Headers) == 0 { + return nil + } + names := make([]string, 0, len(config.Headers)) + for name := range config.Headers { + names = append(names, name) + } + return names +} + +// sortedKeys returns the keys of m in deterministic order (helpful for stable +// UI rendering — Go map iteration order is random). +func sortedKeys(m map[string]string) []string { + out := make([]string, 0, len(m)) + for k := range m { + out = append(out, k) + } + sort.Strings(out) + return out +} diff --git a/transports/bifrost-http/handlers/mcp_sessions.go b/transports/bifrost-http/handlers/mcp_sessions.go index df12e3fb38..a4792c2262 100644 --- a/transports/bifrost-http/handlers/mcp_sessions.go +++ b/transports/bifrost-http/handlers/mcp_sessions.go @@ -57,21 +57,27 @@ type userSummary struct { Name string `json:"name,omitempty"` } -// mcpSessionRow is the wire shape for both authenticated tokens and pending flows. +// mcpSessionRow is the wire shape for OAuth tokens, OAuth pending flows, +// per-user header credentials, and per-user header pending flows. Kind +// discriminates the row TYPE; AuthKind disambiguates flow rows by auth +// surface (OAuth vs Headers) so the UI can route "Complete authentication" +// to the correct landing page. Per-kind-only fields use omitempty. type mcpSessionRow struct { ID string `json:"id"` - Kind string `json:"kind"` // "token" | "flow" + Kind string `json:"kind"` // "token" | "flow" | "header" + AuthKind string `json:"auth_kind"` // "oauth" | "headers" — disambiguates flow rows; tokens are always "oauth", header creds are always "headers" AuthMode string `json:"auth_mode"` UserID *string `json:"user_id,omitempty"` User *userSummary `json:"user,omitempty"` // Preloaded by enterprise on user-keyed rows; nil in OSS so UI falls back to user_id VirtualKey *virtualKeySummary `json:"virtual_key,omitempty"` MCPClient *mcpClientSummary `json:"mcp_client,omitempty"` SessionID *string `json:"session_id,omitempty"` // Session-mode identity: caller-issued x-bf-mcp-session-id value - Status string `json:"status"` // 'active' | 'orphaned' | 'pending' | 'needs_reauth' - ExpiresAt *string `json:"expires_at,omitempty"` // RFC3339; nil for non-expiring tokens + Status string `json:"status"` // OAuth: 'active' | 'orphaned' | 'pending' | 'needs_reauth'. Headers: 'active' | 'orphaned' | 'needs_update'. + ExpiresAt *string `json:"expires_at,omitempty"` // OAuth-only; nil for non-expiring tokens and always nil for headers CreatedAt string `json:"created_at"` // When the session was first authenticated - LastRefreshedAt *string `json:"last_refreshed_at,omitempty"` // Token rows only; nil if never refreshed - OauthConfigID string `json:"oauth_config_id,omitempty"` + LastRefreshedAt *string `json:"last_refreshed_at,omitempty"` // OAuth token rows only; nil if never refreshed + UpdatedAt *string `json:"updated_at,omitempty"` // Headers rows: timestamp of last submission/edit + OauthConfigID string `json:"oauth_config_id,omitempty"` // OAuth rows only } type mcpSessionsListResponse struct { @@ -88,9 +94,17 @@ func (h *MCPSessionsHandler) list(ctx *fasthttp.RequestCtx) { // used by getVirtualKeys, getPrompts, getTeams, etc. tokens, err := h.store.ConfigStore.ListAllOauthUserTokens(ctx) var flows []tables.TableOauthUserSession + var headerCreds []tables.TableMCPPerUserHeaderCredential + var headerFlows []tables.TableMCPPerUserHeaderFlow if err == nil { flows, err = h.store.ConfigStore.ListAllPendingOauthUserSessions(ctx) } + if err == nil { + headerCreds, err = h.store.ConfigStore.ListAllMCPPerUserHeaderCredentials(ctx) + } + if err == nil { + headerFlows, err = h.store.ConfigStore.ListAllPendingMCPPerUserHeaderFlows(ctx) + } if err != nil { logger.Error("[mcp/sessions] list failed: %v", err) SendError(ctx, fasthttp.StatusInternalServerError, "Failed to list MCP sessions") @@ -107,8 +121,15 @@ func (h *MCPSessionsHandler) list(ctx *fasthttp.RequestCtx) { for _, t := range tokens { tokenBindings[bindingKeyFromToken(t)] = struct{}{} } + // Same de-dup model for headers: a credential row + a pending flow + // for the same binding means the user re-initiated submission for an + // already-stored credential. Surface the credential row only. + headerCredBindings := make(map[sessionBindingKey]struct{}, len(headerCreds)) + for _, c := range headerCreds { + headerCredBindings[bindingKeyFromHeaderCredential(c)] = struct{}{} + } - rows := make([]mcpSessionRow, 0, len(tokens)+len(flows)) + rows := make([]mcpSessionRow, 0, len(tokens)+len(flows)+len(headerCreds)+len(headerFlows)) for _, t := range tokens { rows = append(rows, tokenRow(t)) } @@ -124,6 +145,15 @@ func (h *MCPSessionsHandler) list(ctx *fasthttp.RequestCtx) { } rows = append(rows, flowRow(f)) } + for _, c := range headerCreds { + rows = append(rows, headerCredentialRow(c)) + } + for _, f := range headerFlows { + if _, hasCred := headerCredBindings[bindingKeyFromHeaderFlow(f)]; hasCred { + continue + } + rows = append(rows, headerFlowRow(f)) + } SendJSON(ctx, mcpSessionsListResponse{Sessions: rows}) } @@ -167,8 +197,12 @@ func bindingKeyFromFlow(f tables.TableOauthUserSession) sessionBindingKey { return k } -// reauth starts a fresh OAuth flow for the MCP client backing the given token -// row. Returns the authorize URL the user must visit. +// reauth starts a fresh OAuth flow OR a fresh header-submission flow for +// the MCP client backing the given session row. The row's table determines +// the branch — header credential rows mint a header submission flow, OAuth +// token rows mint an OAuth flow. Returns the URL the caller must visit +// (authorize_url for OAuth, submit_url for headers) along with the new +// flow's ID under session_id (kept identical for UI symmetry). func (h *MCPSessionsHandler) reauth(ctx *fasthttp.RequestCtx) { rowID, ok := ctx.UserValue("id").(string) if !ok || rowID == "" { @@ -178,6 +212,13 @@ func (h *MCPSessionsHandler) reauth(ctx *fasthttp.RequestCtx) { bfCtx, cancel := lib.ConvertToBifrostContext(ctx, h.store) defer cancel() + // Header credential rows and OAuth token rows are both UUIDs in + // separate tables. Try headers first; on miss fall through to OAuth. + if headerCred, _ := h.store.ConfigStore.GetMCPPerUserHeaderCredentialByID(ctx, rowID); headerCred != nil { + h.reauthHeaderCredential(ctx, bfCtx, headerCred) + return + } + tok, err := h.loadRowAuthorizedForCaller(ctx, rowID) if err != nil { // loadRowAuthorizedForCaller already wrote the error response. @@ -233,6 +274,64 @@ func (h *MCPSessionsHandler) reauth(ctx *fasthttp.RequestCtx) { }) } +// reauthHeaderCredential is the header-credential branch of reauth: creates +// a fresh per-user-headers submission flow row keyed to the credential's +// existing identity, then returns the submission URL. Mirrors the OAuth +// branch's call to InitiateUserOAuthFlow — same shape on the wire so the UI +// can render a single "click → redirect" affordance for both kinds. +func (h *MCPSessionsHandler) reauthHeaderCredential(ctx *fasthttp.RequestCtx, bfCtx *schemas.BifrostContext, cred *tables.TableMCPPerUserHeaderCredential) { + if cred.Status == "orphaned" { + SendError(ctx, fasthttp.StatusForbidden, "Access to this MCP has been revoked. Re-submitting headers will not restore access - contact your administrator.") + return + } + provider := h.store.MCPHeadersProvider + if provider == nil { + SendError(ctx, fasthttp.StatusServiceUnavailable, "Per-user headers provider not configured") + return + } + + rowMode := schemas.MCPAuthMode(cred.AuthMode) + identity := "" + switch rowMode { + case schemas.MCPAuthModeUser: + if cred.UserID != nil { + identity = *cred.UserID + } + case schemas.MCPAuthModeVK: + if cred.VirtualKeyID != nil { + identity = *cred.VirtualKeyID + } + case schemas.MCPAuthModeSession: + identity = cred.SessionID + } + if identity == "" { + SendError(ctx, fasthttp.StatusInternalServerError, "Credential row is missing an identity column for its auth mode") + return + } + + baseURL := lib.BuildBaseURL(ctx, h.store.GetMCPExternalClientURL()) + if baseURL == "" { + SendError(ctx, fasthttp.StatusInternalServerError, "Could not derive callback base URL") + return + } + initiation, err := provider.InitiateUserSubmissionFlow(bfCtx, rowMode, identity, cred.MCPClientID, baseURL) + if err != nil { + logger.Error("[mcp/sessions] reauth header flow init failed: cred=%s err=%v", cred.ID, err) + SendError(ctx, fasthttp.StatusInternalServerError, "Failed to initiate header resubmission") + return + } + logger.Debug("[mcp/sessions] reauth header initiated: cred=%s mcp_client=%s mode=%s flow=%s", cred.ID, cred.MCPClientID, rowMode, initiation.FlowID) + SendJSON(ctx, map[string]any{ + // submit_url is set so a future header-specific client can branch + // cleanly; authorize_url stays compatible with the existing UI which + // already does window.location.href = res.authorize_url. + "authorize_url": initiation.FrontendURL, + "submit_url": initiation.FrontendURL, + "session_id": initiation.FlowID, + "kind": "headers", + }) +} + // revoke hard-deletes the local token row and any pending flow rows for the // same identity + MCP client. Upstream revocation against the OAuth provider // is NOT performed — the per-user OAuth template config doesn't carry a @@ -249,6 +348,55 @@ func (h *MCPSessionsHandler) revoke(ctx *fasthttp.RequestCtx) { SendError(ctx, fasthttp.StatusBadRequest, "Invalid session id") return } + // Row IDs from three tables (OAuth tokens, header credentials, header + // flows) are all UUIDs. Try headers first (credential, then pending + // flow); on miss fall through to OAuth tokens. Each branch returns; + // only one delete runs per request. + if headerCred, _ := h.store.ConfigStore.GetMCPPerUserHeaderCredentialByID(ctx, rowID); headerCred != nil { + // Drop pending submission flow rows for the same binding BEFORE + // the credential. If a flow finishes (submit lands) after the + // credential is gone, the upsert would mint a fresh credential and + // undo the revoke — same race protection as the OAuth revoke path. + credMode := schemas.MCPAuthMode(headerCred.AuthMode) + credIdentity := "" + switch credMode { + case schemas.MCPAuthModeUser: + if headerCred.UserID != nil { + credIdentity = *headerCred.UserID + } + case schemas.MCPAuthModeVK: + if headerCred.VirtualKeyID != nil { + credIdentity = *headerCred.VirtualKeyID + } + case schemas.MCPAuthModeSession: + credIdentity = headerCred.SessionID + } + if credIdentity != "" { + if delErr := h.store.ConfigStore.DeleteMCPPerUserHeaderFlowsByModeIdentityAndMCPClient(ctx, credMode, credIdentity, headerCred.MCPClientID); delErr != nil { + logger.Error("[mcp/sessions] clearing header flow rows failed: cred=%s err=%v", rowID, delErr) + SendError(ctx, fasthttp.StatusInternalServerError, "Failed to delete MCP session") + return + } + } + if err := h.store.ConfigStore.DeleteMCPPerUserHeaderCredential(ctx, headerCred.ID); err != nil { + logger.Error("[mcp/sessions] delete header credential failed: id=%s err=%v", rowID, err) + SendError(ctx, fasthttp.StatusInternalServerError, "Failed to delete MCP session") + return + } + logger.Debug("[mcp/sessions] revoked header credential: id=%s mcp_client=%s mode=%s", rowID, headerCred.MCPClientID, headerCred.AuthMode) + ctx.SetStatusCode(fasthttp.StatusNoContent) + return + } + if headerFlow, _ := h.store.ConfigStore.GetMCPPerUserHeaderFlowByID(ctx, rowID); headerFlow != nil { + if err := h.store.ConfigStore.DeleteMCPPerUserHeaderFlow(ctx, headerFlow.ID); err != nil { + logger.Error("[mcp/sessions] delete header flow failed: id=%s err=%v", rowID, err) + SendError(ctx, fasthttp.StatusInternalServerError, "Failed to delete MCP session") + return + } + logger.Debug("[mcp/sessions] revoked header flow: id=%s mcp_client=%s mode=%s", rowID, headerFlow.MCPClientID, headerFlow.FlowMode) + ctx.SetStatusCode(fasthttp.StatusNoContent) + return + } tok, err := h.loadRowAuthorizedForCaller(ctx, rowID) if err != nil { _ = err @@ -492,6 +640,7 @@ func tokenRow(t tables.TableOauthUserToken) mcpSessionRow { row := mcpSessionRow{ ID: t.ID, Kind: "token", + AuthKind: "oauth", AuthMode: t.AuthMode, UserID: t.UserID, Status: t.Status, @@ -532,6 +681,7 @@ func flowRow(f tables.TableOauthUserSession) mcpSessionRow { row := mcpSessionRow{ ID: f.ID, Kind: "flow", + AuthKind: "oauth", AuthMode: f.FlowMode, UserID: f.UserID, Status: f.Status, @@ -559,4 +709,112 @@ func flowRow(f tables.TableOauthUserSession) mcpSessionRow { return row } +// headerCredentialRow maps a per-user MCP header credential row to the wire +// shape. No expires_at / last_refreshed_at / oauth_config_id — those are +// OAuth-specific. updated_at is included so the UI can show "submitted Xm +// ago" / "edited Xm ago". +func headerCredentialRow(c tables.TableMCPPerUserHeaderCredential) mcpSessionRow { + updated := c.UpdatedAt.UTC().Format(rfc3339Nano) + row := mcpSessionRow{ + ID: c.ID, + Kind: "header", + AuthKind: "headers", + AuthMode: c.AuthMode, + UserID: c.UserID, + Status: c.Status, + CreatedAt: c.CreatedAt.UTC().Format(rfc3339Nano), + UpdatedAt: &updated, + } + if c.MCPClient != nil { + row.MCPClient = &mcpClientSummary{ClientID: c.MCPClient.ClientID, Name: c.MCPClient.Name} + } else { + row.MCPClient = &mcpClientSummary{ClientID: c.MCPClientID} + } + if c.VirtualKey != nil { + row.VirtualKey = &virtualKeySummary{ID: c.VirtualKey.ID, Name: c.VirtualKey.Name} + } else if c.VirtualKeyID != nil { + row.VirtualKey = &virtualKeySummary{ID: *c.VirtualKeyID} + } + if c.User != nil { + row.User = &userSummary{ID: c.User.ID, Name: c.User.Name} + } + if c.AuthMode == string(schemas.MCPAuthModeSession) && c.SessionID != "" { + s := c.SessionID + row.SessionID = &s + } + return row +} + +// headerFlowRow maps a pending per-user header submission flow row to the +// wire shape. Mirrors flowRow for OAuth — same "Pending" kind status so the +// UI renders it uniformly with OAuth pending flows. +func headerFlowRow(f tables.TableMCPPerUserHeaderFlow) mcpSessionRow { + exp := f.ExpiresAt.UTC().Format(rfc3339Nano) + row := mcpSessionRow{ + ID: f.ID, + Kind: "flow", + AuthKind: "headers", + AuthMode: f.FlowMode, + UserID: f.UserID, + Status: f.Status, + ExpiresAt: &exp, + CreatedAt: f.CreatedAt.UTC().Format(rfc3339Nano), + } + if f.MCPClient != nil { + row.MCPClient = &mcpClientSummary{ClientID: f.MCPClient.ClientID, Name: f.MCPClient.Name} + } else { + row.MCPClient = &mcpClientSummary{ClientID: f.MCPClientID} + } + if f.VirtualKey != nil { + row.VirtualKey = &virtualKeySummary{ID: f.VirtualKey.ID, Name: f.VirtualKey.Name} + } else if f.VirtualKeyID != nil { + row.VirtualKey = &virtualKeySummary{ID: *f.VirtualKeyID} + } + if f.User != nil { + row.User = &userSummary{ID: f.User.ID, Name: f.User.Name} + } + if f.FlowMode == string(schemas.MCPAuthModeSession) && f.SessionID != "" { + s := f.SessionID + row.SessionID = &s + } + return row +} + +// bindingKeyFromHeaderCredential / bindingKeyFromHeaderFlow build a comparable +// key for de-dup between credentials and their in-flight resubmission flow. +// Mirrors bindingKeyFromToken / bindingKeyFromFlow on the OAuth side. +func bindingKeyFromHeaderCredential(c tables.TableMCPPerUserHeaderCredential) sessionBindingKey { + k := sessionBindingKey{Mode: c.AuthMode, MCPClientID: c.MCPClientID} + switch schemas.MCPAuthMode(c.AuthMode) { + case schemas.MCPAuthModeUser: + if c.UserID != nil { + k.Identity = *c.UserID + } + case schemas.MCPAuthModeVK: + if c.VirtualKeyID != nil { + k.Identity = *c.VirtualKeyID + } + case schemas.MCPAuthModeSession: + k.Identity = c.SessionID + } + return k +} + +func bindingKeyFromHeaderFlow(f tables.TableMCPPerUserHeaderFlow) sessionBindingKey { + k := sessionBindingKey{Mode: f.FlowMode, MCPClientID: f.MCPClientID} + switch schemas.MCPAuthMode(f.FlowMode) { + case schemas.MCPAuthModeUser: + if f.UserID != nil { + k.Identity = *f.UserID + } + case schemas.MCPAuthModeVK: + if f.VirtualKeyID != nil { + k.Identity = *f.VirtualKeyID + } + case schemas.MCPAuthModeSession: + k.Identity = f.SessionID + } + return k +} + const rfc3339Nano = "2006-01-02T15:04:05.999999999Z07:00" diff --git a/transports/bifrost-http/handlers/mcpserver.go b/transports/bifrost-http/handlers/mcpserver.go index 627759fcf8..6412d49be5 100644 --- a/transports/bifrost-http/handlers/mcpserver.go +++ b/transports/bifrost-http/handlers/mcpserver.go @@ -367,10 +367,21 @@ func (h *MCPServerHandler) syncServer(server *server.MCPServer, availableTools [ // Execute the tool via tool executor toolMessage, err := h.toolManager.ExecuteChatMCPTool(ctx, &toolCall) if err != nil { - if err.ExtraFields.MCPAuthRequired != nil { + if authReq := err.ExtraFields.MCPAuthRequired; authReq != nil { + // Two surfaces share this error: per-user OAuth uses + // AuthorizeURL (the upstream provider's authorize page); + // per-user headers uses SubmitURL (the workspace landing + // page where the user submits their header values). + // Pick whichever Kind populated. + url := authReq.AuthorizeURL + action := "connect your account" + if authReq.Kind == schemas.MCPAuthRequiredKindHeaders { + url = authReq.SubmitURL + action = "submit the required headers" + } return mcp.NewToolResultError(fmt.Sprintf( - "Authentication required for %s. Open this URL to connect your account: %s", - err.ExtraFields.MCPAuthRequired.MCPClientName, err.ExtraFields.MCPAuthRequired.AuthorizeURL, + "Authentication required for %s. Open this URL to %s: %s", + authReq.MCPClientName, action, url, )), nil } return mcp.NewToolResultError(fmt.Sprintf("Tool execution failed: %v", bifrost.GetErrorMessage(err))), nil diff --git a/transports/bifrost-http/handlers/temp_token_scopes.go b/transports/bifrost-http/handlers/temp_token_scopes.go index 35b387810b..5a5cb63132 100644 --- a/transports/bifrost-http/handlers/temp_token_scopes.go +++ b/transports/bifrost-http/handlers/temp_token_scopes.go @@ -32,6 +32,25 @@ var mcpAuthScope = temptoken.Scope{ MaxTTL: 15 * time.Minute, } +// mcpHeadersAuthScope declares the routes the mcp_headers_auth scope grants +// access to. The flow ID is substituted into {id} at validation time, +// binding each token to exactly one headers submission flow. Mirrors +// mcpAuthScope structurally — same TTL, same {id} binding pattern. +// +// The page makes flowDetail then flowSubmit; both routes must remain valid +// for multiple requests within the TTL. Invalidation isn't single-use — +// it happens at submission completion when the submit handler deletes the +// flow row (and the token by resource_id alongside it). +var mcpHeadersAuthScope = temptoken.Scope{ + Name: temptoken.MCPHeadersAuthScopeName, + AllowedRoutes: []temptoken.RoutePattern{ + {Method: "GET", Path: "/api/mcp/per-user-headers/flows/{id}"}, + {Method: "PUT", Path: "/api/mcp/per-user-headers/flows/{id}"}, + }, + ResourceIDInPath: "{id}", + MaxTTL: 15 * time.Minute, +} + // RegisterTempTokenScopes registers every scope owned by this handlers // package on the given service. Called at server startup once the service // has been constructed. Returns an error if any scope is invalid or has @@ -43,5 +62,8 @@ func RegisterTempTokenScopes(svc *temptoken.Service) error { if err := svc.Registry().Register(mcpAuthScope); err != nil { return fmt.Errorf("temp_token_scopes: register mcp_auth: %w", err) } + if err := svc.Registry().Register(mcpHeadersAuthScope); err != nil { + return fmt.Errorf("temp_token_scopes: register mcp_headers_auth: %w", err) + } return nil } diff --git a/transports/bifrost-http/lib/config.go b/transports/bifrost-http/lib/config.go index 95361cceb9..b6adac74cf 100644 --- a/transports/bifrost-http/lib/config.go +++ b/transports/bifrost-http/lib/config.go @@ -32,6 +32,7 @@ import ( "github.com/maximhq/bifrost/framework/featureflags" "github.com/maximhq/bifrost/framework/kvstore" "github.com/maximhq/bifrost/framework/logstore" + "github.com/maximhq/bifrost/framework/mcp_headers" "github.com/maximhq/bifrost/framework/mcpcatalog" "github.com/maximhq/bifrost/framework/modelcatalog" "github.com/maximhq/bifrost/framework/oauth2" @@ -353,6 +354,13 @@ type Config struct { TokenRefreshWorker *oauth2.TokenRefreshWorker OAuthSweepWorker *oauth2.PerUserOAuthSweepWorker + // MCPHeadersProvider backs MCPAuthTypePerUserHeaders credential storage. + // Constructed alongside OAuthProvider and passed into the Bifrost core + // init so the per-user-headers resolver can resolve / persist values + // scoped by (auth_mode, identity, mcp_client). + MCPHeadersProvider *mcp_headers.Provider + MCPHeadersSweepWorker *mcp_headers.CredentialSweepWorker + // Async job executor (initialized during setup if LogsStore + governance are available) AsyncJobExecutor *logstore.AsyncJobExecutor // Shared in-memory kvstore for transport-level protocol coordination. @@ -1458,6 +1466,7 @@ func mcpClientConfigToTable(clientConfig *schemas.MCPClientConfig) (configstoreT Disabled: clientConfig.Disabled, DiscoveredTools: clientConfig.DiscoveredTools, DiscoveredToolNameMapping: clientConfig.DiscoveredToolNameMapping, + PerUserHeaderKeys: clientConfig.PerUserHeaderKeys, ConfigHash: clientConfig.ConfigHash, }, nil } @@ -3203,6 +3212,9 @@ func initFrameworkConfig(ctx context.Context, config *Config, configData *Config // Initialize OAuth provider config.OAuthProvider = oauth2.NewOAuth2Provider(config.ConfigStore, logger) + // Initialize per-user-headers credential provider. Storage parallel of + // OAuthProvider for MCPAuthTypePerUserHeaders clients. + config.MCPHeadersProvider = mcp_headers.NewProvider(config.ConfigStore, logger) // Start token refresh worker for automatic OAuth token refresh config.TokenRefreshWorker = oauth2.NewTokenRefreshWorker(config.OAuthProvider, logger) @@ -3217,6 +3229,15 @@ func initFrameworkConfig(ctx context.Context, config *Config, configData *Config config.OAuthSweepWorker.Start(ctx) } + // Start per-user headers credential sweep worker. Parallel of the OAuth + // sweep but only reaps orphaned credential rows (no flow table to sweep). + // Same 30-day retention so admin expectations stay uniform across the two + // per-user auth surfaces. + config.MCPHeadersSweepWorker = mcp_headers.NewCredentialSweepWorker(config.MCPHeadersProvider, 30*24*time.Hour, logger) + if config.MCPHeadersSweepWorker != nil { + config.MCPHeadersSweepWorker.Start(ctx) + } + config.FrameworkConfig = &framework.FrameworkConfig{ Pricing: pricingConfig, } @@ -3682,6 +3703,9 @@ func (c *Config) Close(ctx context.Context) { if c.OAuthSweepWorker != nil { c.OAuthSweepWorker.Stop() } + if c.MCPHeadersSweepWorker != nil { + c.MCPHeadersSweepWorker.Stop() + } if c.KVStore != nil { c.KVStore.Close() } @@ -4838,6 +4862,7 @@ func (c *Config) UpdateMCPClient(ctx context.Context, id string, updatedConfig * c.MCPConfig.ClientConfigs[configIndex].ToolSyncInterval = updatedConfig.ToolSyncInterval c.MCPConfig.ClientConfigs[configIndex].AllowOnAllVirtualKeys = updatedConfig.AllowOnAllVirtualKeys c.MCPConfig.ClientConfigs[configIndex].Disabled = updatedConfig.Disabled + c.MCPConfig.ClientConfigs[configIndex].PerUserHeaderKeys = updatedConfig.PerUserHeaderKeys // Handle disable/enable lifecycle when the Disabled flag toggles and the client // is registered at runtime. We call the core bifrost methods directly (not the diff --git a/transports/bifrost-http/lib/config_test.go b/transports/bifrost-http/lib/config_test.go index 0aac500d8f..77afb4a99c 100644 --- a/transports/bifrost-http/lib/config_test.go +++ b/transports/bifrost-http/lib/config_test.go @@ -1281,6 +1281,51 @@ func (m *MockConfigStore) DeleteOrphanedOauthUserTokens(ctx context.Context, old return 0, nil } +// Per-user MCP header credentials +func (m *MockConfigStore) GetMCPPerUserHeaderCredentialByMode(ctx context.Context, mode schemas.MCPAuthMode, identity, mcpClientID string) (*tables.TableMCPPerUserHeaderCredential, error) { + return nil, nil +} +func (m *MockConfigStore) GetMCPPerUserHeaderCredentialByID(ctx context.Context, id string) (*tables.TableMCPPerUserHeaderCredential, error) { + return nil, nil +} +func (m *MockConfigStore) UpsertMCPPerUserHeaderCredential(ctx context.Context, cred *tables.TableMCPPerUserHeaderCredential) error { + return nil +} +func (m *MockConfigStore) DeleteMCPPerUserHeaderCredential(ctx context.Context, id string) error { + return nil +} +func (m *MockConfigStore) ListAllMCPPerUserHeaderCredentials(ctx context.Context) ([]tables.TableMCPPerUserHeaderCredential, error) { + return nil, nil +} +func (m *MockConfigStore) MarkMCPPerUserHeaderCredentialsNeedsUpdate(ctx context.Context, mcpClientID string) error { + return nil +} +func (m *MockConfigStore) DeleteOrphanedMCPPerUserHeaderCredentials(ctx context.Context, olderThan time.Duration) (int64, error) { + return 0, nil +} +func (m *MockConfigStore) CreateMCPPerUserHeaderFlow(ctx context.Context, flow *tables.TableMCPPerUserHeaderFlow) error { + return nil +} +func (m *MockConfigStore) GetMCPPerUserHeaderFlowByID(ctx context.Context, id string) (*tables.TableMCPPerUserHeaderFlow, error) { + return nil, nil +} +func (m *MockConfigStore) GetMCPPerUserHeaderFlowByModeIdentityAndMCPClient(ctx context.Context, mode schemas.MCPAuthMode, identity, mcpClientID string) (*tables.TableMCPPerUserHeaderFlow, error) { + return nil, nil +} +func (m *MockConfigStore) UpdateMCPPerUserHeaderFlow(ctx context.Context, flow *tables.TableMCPPerUserHeaderFlow) error { + return nil +} +func (m *MockConfigStore) DeleteMCPPerUserHeaderFlowsByModeIdentityAndMCPClient(ctx context.Context, mode schemas.MCPAuthMode, identity, mcpClientID string) error { + return nil +} +func (m *MockConfigStore) DeleteMCPPerUserHeaderFlow(ctx context.Context, id string) error { return nil } +func (m *MockConfigStore) ListAllPendingMCPPerUserHeaderFlows(ctx context.Context) ([]tables.TableMCPPerUserHeaderFlow, error) { + return nil, nil +} +func (m *MockConfigStore) DeleteExpiredMCPPerUserHeaderFlows(ctx context.Context) (int64, error) { + return 0, nil +} + // Routing rules func (m *MockConfigStore) GetRoutingRules(ctx context.Context) ([]tables.TableRoutingRule, error) { return nil, nil diff --git a/transports/bifrost-http/lib/ctx.go b/transports/bifrost-http/lib/ctx.go index 6226c4af4f..b1b13cc241 100644 --- a/transports/bifrost-http/lib/ctx.go +++ b/transports/bifrost-http/lib/ctx.go @@ -620,15 +620,18 @@ func ConvertToBifrostContext(ctx *fasthttp.RequestCtx, store HandlerStore) (*sch }) bifrostCtx.SetValue(schemas.BifrostContextKeyRequestHeaders, allHeaders) - // Build and set OAuth redirect URI for per-user OAuth flows. Bifrost is acting as - // the OAuth client to upstream MCP servers here, so use the client-side override. + // Build and set the MCP callback base URL. Used by per-user OAuth (appends + // /api/oauth/callback) and per-user headers (appends the workspace submit + // path) resolvers when initiating their respective auth flows. Bifrost is + // acting as the OAuth client to upstream MCP servers here, so the client- + // side override applies. var externalClientURL string if store != nil { externalClientURL = store.GetMCPExternalClientURL() } baseURL := BuildBaseURL(ctx, externalClientURL) if baseURL != "" { - bifrostCtx.SetValue(schemas.BifrostContextKeyOAuthRedirectURI, baseURL+"/api/oauth/callback") + bifrostCtx.SetValue(schemas.BifrostContextKeyMCPCallbackBaseURL, baseURL) } bifrostCtx.SetValue(schemas.BifrostContextKeyAllowPerRequestStorageOverride, allowPerRequestStorageOverride) diff --git a/transports/bifrost-http/lib/lib.go b/transports/bifrost-http/lib/lib.go index 75e2c97812..2a27a98f6c 100644 --- a/transports/bifrost-http/lib/lib.go +++ b/transports/bifrost-http/lib/lib.go @@ -15,6 +15,24 @@ func SetLogger(l schemas.Logger) { logger = l } +// HasDuplicates reports whether the slice contains any repeated element. +// Comparison is exact; callers needing case-insensitive or whitespace-tolerant +// semantics should normalize the slice before calling (e.g. lower-case the +// entries for case-insensitive HTTP header names). +func HasDuplicates[T comparable](items []T) bool { + if len(items) < 2 { + return false + } + seen := make(map[T]struct{}, len(items)) + for _, it := range items { + if _, dup := seen[it]; dup { + return true + } + seen[it] = struct{}{} + } + return false +} + // StreamLargeResponseBody extracts the large response reader from context and streams // it directly to the client. Sets status 200, content-type, and content-length headers. // Returns false if the reader is not available (caller should send an error response). diff --git a/transports/bifrost-http/server/server.go b/transports/bifrost-http/server/server.go index 5b1f316f64..c0488e35cb 100644 --- a/transports/bifrost-http/server/server.go +++ b/transports/bifrost-http/server/server.go @@ -105,6 +105,8 @@ type ServerCallbacks interface { UpdateMCPToolManagerConfig(ctx context.Context, maxAgentDepth int, toolExecutionTimeoutInSeconds int, codeModeBindingLevel string, disableAutoToolInject bool) error // VerifyPerUserOAuthConnection verifies an MCP server using a temporary token and discovers tools. VerifyPerUserOAuthConnection(ctx context.Context, config *schemas.MCPClientConfig, accessToken string) (map[string]schemas.ChatTool, map[string]string, error) + // VerifyHeadersConnection verifies an MCP server using user-supplied header values and discovers tools. + VerifyHeadersConnection(ctx context.Context, config *schemas.MCPClientConfig, userHeaders map[string]string) (map[string]schemas.ChatTool, map[string]string, error) // SetClientTools updates the tool map for an existing client. SetClientTools(clientID string, tools map[string]schemas.ChatTool, toolNameMapping map[string]string) ReconnectMCPClient(ctx context.Context, id string) error @@ -277,6 +279,12 @@ func (s *BifrostHTTPServer) EnableMCPClient(ctx context.Context, id string) erro return nil } +// VerifyHeadersConnection delegates to the Bifrost client to verify an MCP +// server with caller-supplied header values and discover its tools. +func (s *BifrostHTTPServer) VerifyHeadersConnection(ctx context.Context, config *schemas.MCPClientConfig, userHeaders map[string]string) (map[string]schemas.ChatTool, map[string]string, error) { + return s.Client.VerifyHeadersConnection(ctx, config, userHeaders) +} + // VerifyPerUserOAuthConnection delegates to the Bifrost client to verify an MCP // server using a temporary access token and discover available tools. func (s *BifrostHTTPServer) VerifyPerUserOAuthConnection(ctx context.Context, config *schemas.MCPClientConfig, accessToken string) (map[string]schemas.ChatTool, map[string]string, error) { @@ -1186,6 +1194,7 @@ func (s *BifrostHTTPServer) RegisterAPIRoutes(ctx context.Context, callbacks Ser providerHandler := handlers.NewProviderHandler(callbacks, s.Config, s.Client) oauthHandler := handlers.NewOAuthHandler(s.Config.OAuthProvider, s.Client, s.Config) mcpHandler := handlers.NewMCPHandler(callbacks, callbacks, s.Client, s.Config, oauthHandler) + mcpPerUserHeadersHandler := handlers.NewMCPPerUserHeadersHandler(callbacks, s.Config, s.TempTokens) mcpSessionsHandler := handlers.NewMCPSessionsHandler(s.Config) configHandler := handlers.NewConfigHandler(callbacks, s.Config) pluginsHandler := handlers.NewPluginsHandler(callbacks, s.Config.ConfigStore) @@ -1196,6 +1205,7 @@ func (s *BifrostHTTPServer) RegisterAPIRoutes(ctx context.Context, callbacks Ser healthHandler.RegisterRoutes(s.Router, middlewares...) providerHandler.RegisterRoutes(s.Router, middlewares...) mcpHandler.RegisterRoutes(s.Router, middlewares...) + mcpPerUserHeadersHandler.RegisterRoutes(s.Router, middlewares...) mcpSessionsHandler.RegisterRoutes(s.Router, middlewares...) configHandler.RegisterRoutes(s.Router, middlewares...) oauthHandler.RegisterRoutes(s.Router, middlewares...) @@ -1418,6 +1428,7 @@ func (s *BifrostHTTPServer) Bootstrap(ctx context.Context) error { MCPPlugins: s.Config.GetLoadedMCPPlugins(), MCPConfig: mcpConfig, OAuth2Provider: s.Config.OAuthProvider, + MCPHeadersProvider: s.Config.MCPHeadersProvider, Logger: logger, KVStore: s.Config.KVStore, }) @@ -1515,6 +1526,11 @@ func (s *BifrostHTTPServer) Bootstrap(ctx context.Context) error { if s.Config.OAuthProvider != nil { s.Config.OAuthProvider.SetTempTokenService(s.TempTokens) } + // Same wiring for the per-user-headers provider — mints + // mcp_headers_auth tokens on the headers submission URL. + if s.Config.MCPHeadersProvider != nil { + s.Config.MCPHeadersProvider.SetTempTokenService(s.TempTokens) + } s.AuthMiddleware, err = handlers.InitAuthMiddleware(s.Config.ConfigStore, s.WSTicketStore, s.TempTokens) if err != nil { s.WSTicketStore.Stop()