diff --git a/core/bifrost.go b/core/bifrost.go index d5d80f2381..fcc612df12 100644 --- a/core/bifrost.go +++ b/core/bifrost.go @@ -1216,6 +1216,17 @@ func (bifrost *Bifrost) GetMCPClients() ([]schemas.MCPClient, error) { return clientsInConfig, nil } +// GetAvailableTools returns the available tools for the given context. +// +// Returns: +// - []schemas.ChatTool: List of available tools +func (bifrost *Bifrost) GetAvailableMCPTools(ctx context.Context) []schemas.ChatTool { + if bifrost.mcpManager == nil { + return nil + } + return bifrost.mcpManager.GetAvailableTools(ctx) +} + // AddMCPClient adds a new MCP client to the Bifrost instance. // This allows for dynamic MCP client management at runtime. // @@ -2079,6 +2090,7 @@ func executeRequestWithRetries[T any]( // Calculate and apply backoff backoff := calculateBackoff(attempts-1, config) + logger.Debug("sleeping for %s before retry", backoff) time.Sleep(backoff) } diff --git a/core/mcp/mcp.go b/core/mcp/mcp.go index 848f3bf6a0..388a4f6ab7 100644 --- a/core/mcp/mcp.go +++ b/core/mcp/mcp.go @@ -80,9 +80,11 @@ func NewMCPManager(ctx context.Context, config schemas.MCPConfig, logger schemas } manager.toolsHandler = NewToolsManager(config.ToolManagerConfig, manager, config.FetchNewRequestIDFunc) // Process client configs: create client map entries and establish connections - for _, clientConfig := range config.ClientConfigs { - if err := manager.AddClient(clientConfig); err != nil { - logger.Warn(fmt.Sprintf("%s Failed to add MCP client %s: %v", MCPLogPrefix, clientConfig.Name, err)) + if len(config.ClientConfigs) > 0 { + for _, clientConfig := range config.ClientConfigs { + if err := manager.AddClient(clientConfig); err != nil { + logger.Warn(fmt.Sprintf("%s Failed to add MCP client %s: %v", MCPLogPrefix, clientConfig.Name, err)) + } } } logger.Info(MCPLogPrefix + " MCP Manager initialized") @@ -103,6 +105,10 @@ func (m *MCPManager) AddToolsToRequest(ctx context.Context, req *schemas.Bifrost return m.toolsHandler.ParseAndAddToolsToRequest(ctx, req) } +func (m *MCPManager) GetAvailableTools(ctx context.Context) []schemas.ChatTool { + return m.toolsHandler.GetAvailableTools(ctx) +} + // ExecuteTool executes a single tool call from a chat assistant message. // It handles tool execution, error handling, and returns the result as a chat message. // diff --git a/core/mcp/toolmanager.go b/core/mcp/toolmanager.go index 78fbde2e85..e0b719ed91 100644 --- a/core/mcp/toolmanager.go +++ b/core/mcp/toolmanager.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "fmt" + "strings" "sync" "sync/atomic" "time" @@ -73,21 +74,8 @@ func NewToolsManager(config *schemas.MCPToolManagerConfig, clientManager ClientM return manager } -// ParseAndAddToolsToRequest parses the available tools per client and adds them to the Bifrost request. -// -// Parameters: -// - ctx: Execution context -// - req: Bifrost request -// - availableToolsPerClient: Map of client name to its available tools -// -// Returns: -// - *schemas.BifrostRequest: Bifrost request with MCP tools added -func (m *ToolsManager) ParseAndAddToolsToRequest(ctx context.Context, req *schemas.BifrostRequest) *schemas.BifrostRequest { - // MCP is only supported for chat and responses requests - if req.ChatRequest == nil && req.ResponsesRequest == nil { - return req - } - +// GetAvailableTools returns the available tools for the given context. +func (m *ToolsManager) GetAvailableTools(ctx context.Context) []schemas.ChatTool { availableToolsPerClient := m.clientManager.GetToolPerClient(ctx) // Flatten tools from all clients into a single slice, avoiding duplicates var availableTools []schemas.ChatTool @@ -133,8 +121,99 @@ func (m *ToolsManager) ParseAndAddToolsToRequest(ctx context.Context, req *schem } } + return availableTools +} + +// buildIntegrationDuplicateCheckMap builds a map of tool names to check for duplicates +// based on the integration user agent. This includes both direct tool names and +// integration-specific naming patterns from existing tools in the request. +// +// Parameters: +// - existingTools: List of existing tools in the request +// - integrationUserAgent: Integration user agent string (e.g., "claude-cli") +// - availableToolsPerClient: Map of client names to their available tools (for reverse pattern matching) +// +// Returns: +// - map[string]bool: Map of tool names/patterns to check against +func buildIntegrationDuplicateCheckMap(existingTools []schemas.ChatTool, integrationUserAgent string) map[string]bool { + duplicateCheckMap := make(map[string]bool) + + // Add direct tool names + for _, tool := range existingTools { + if tool.Function != nil && tool.Function.Name != "" { + duplicateCheckMap[tool.Function.Name] = true + } + } + + // Add integration-specific patterns from existing tools + switch integrationUserAgent { + case "claude-cli": + // Claude CLI uses pattern: mcp__{foreign_name}__{tool_name} + // The middle part is a foreign name we cannot check for, so we extract the last part + // Examples: + // mcp__bifrost__executeToolCode -> executeToolCode + // mcp__bifrost__listToolFiles -> listToolFiles + // mcp__bifrost__readToolFile -> readToolFile + // mcp__calculator__calculator_add -> calculator_add + for _, tool := range existingTools { + if tool.Function != nil && tool.Function.Name != "" { + existingToolName := tool.Function.Name + // Check if existing tool matches Claude CLI pattern: mcp__*__{tool_name} + if strings.HasPrefix(existingToolName, "mcp__") { + // Split on __ and take the last entry (the tool_name) + parts := strings.Split(existingToolName, "__") + if len(parts) >= 3 { + toolName := parts[len(parts)-1] // Last part is the tool name + // Map Claude CLI pattern back to our tool name format + // This handles both regular MCP tools and code mode tools + if toolName != "" { + duplicateCheckMap[toolName] = true + // Also keep the original pattern for direct matching + duplicateCheckMap[existingToolName] = true + } + } + } + } + } + // Add more integration-specific patterns here as needed + // case "another-integration": + // // Add patterns for other integrations + } + + return duplicateCheckMap +} + +// ParseAndAddToolsToRequest parses the available tools per client and adds them to the Bifrost request. +// +// Parameters: +// - ctx: Execution context +// - req: Bifrost request +// - availableToolsPerClient: Map of client name to its available tools +// +// Returns: +// - *schemas.BifrostRequest: Bifrost request with MCP tools added +func (m *ToolsManager) ParseAndAddToolsToRequest(ctx context.Context, req *schemas.BifrostRequest) *schemas.BifrostRequest { + // MCP is only supported for chat and responses requests + if req.ChatRequest == nil && req.ResponsesRequest == nil { + return req + } + + availableTools := m.GetAvailableTools(ctx) + + if len(availableTools) == 0 { + return req + } + + // Get integration user agent for duplicate checking + var integrationUserAgentStr string + integrationUserAgent := ctx.Value(schemas.BifrostContextKey("integration-user-agent")) + if integrationUserAgent != nil { + if str, ok := integrationUserAgent.(string); ok { + integrationUserAgentStr = str + } + } + if len(availableTools) > 0 { - logger.Debug(fmt.Sprintf("%s Adding %d MCP tools to request from %d clients", MCPLogPrefix, len(availableTools), len(availableToolsPerClient))) switch req.RequestType { case schemas.ChatCompletionRequest, schemas.ChatCompletionStreamRequest: // Only allocate new Params if it's nil to preserve caller-supplied settings @@ -144,13 +223,8 @@ func (m *ToolsManager) ParseAndAddToolsToRequest(ctx context.Context, req *schem tools := req.ChatRequest.Params.Tools - // Create a map of existing tool names for O(1) lookup - existingToolsMap := make(map[string]bool) - for _, tool := range tools { - if tool.Function != nil && tool.Function.Name != "" { - existingToolsMap[tool.Function.Name] = true - } - } + // Build integration-aware duplicate check map + duplicateCheckMap := buildIntegrationDuplicateCheckMap(tools, integrationUserAgentStr) // Add MCP tools that are not already present for _, mcpTool := range availableTools { @@ -159,10 +233,13 @@ func (m *ToolsManager) ParseAndAddToolsToRequest(ctx context.Context, req *schem continue } - if !existingToolsMap[mcpTool.Function.Name] { + toolName := mcpTool.Function.Name + + // Check for duplicates using integration-aware logic + if !duplicateCheckMap[toolName] { tools = append(tools, mcpTool) // Update the map to prevent duplicates within MCP tools as well - existingToolsMap[mcpTool.Function.Name] = true + duplicateCheckMap[toolName] = true } } req.ChatRequest.Params.Tools = tools @@ -174,14 +251,22 @@ func (m *ToolsManager) ParseAndAddToolsToRequest(ctx context.Context, req *schem tools := req.ResponsesRequest.Params.Tools - // Create a map of existing tool names for O(1) lookup - existingToolsMap := make(map[string]bool) + // Convert Responses tools to ChatTool format for duplicate checking + existingChatTools := make([]schemas.ChatTool, 0, len(tools)) for _, tool := range tools { if tool.Name != nil { - existingToolsMap[*tool.Name] = true + existingChatTools = append(existingChatTools, schemas.ChatTool{ + Type: schemas.ChatToolTypeFunction, + Function: &schemas.ChatToolFunction{ + Name: *tool.Name, + }, + }) } } + // Build integration-aware duplicate check map + duplicateCheckMap := buildIntegrationDuplicateCheckMap(existingChatTools, integrationUserAgentStr) + // Add MCP tools that are not already present for _, mcpTool := range availableTools { // Skip tools with nil Function or empty Name @@ -189,7 +274,10 @@ func (m *ToolsManager) ParseAndAddToolsToRequest(ctx context.Context, req *schem continue } - if !existingToolsMap[mcpTool.Function.Name] { + toolName := mcpTool.Function.Name + + // Check for duplicates using integration-aware logic + if !duplicateCheckMap[toolName] { responsesTool := mcpTool.ToResponsesTool() // Skip if the converted tool has nil Name if responsesTool.Name == nil { @@ -198,7 +286,7 @@ func (m *ToolsManager) ParseAndAddToolsToRequest(ctx context.Context, req *schem tools = append(tools, *responsesTool) // Update the map to prevent duplicates within MCP tools as well - existingToolsMap[*responsesTool.Name] = true + duplicateCheckMap[toolName] = true } } req.ResponsesRequest.Params.Tools = tools diff --git a/core/mcp/utils.go b/core/mcp/utils.go index cfed7993f0..156909a56c 100644 --- a/core/mcp/utils.go +++ b/core/mcp/utils.go @@ -221,8 +221,6 @@ func canAutoExecuteTool(toolName string, config schemas.MCPClientConfig) bool { func shouldSkipToolForRequest(ctx context.Context, clientName, toolName string) bool { includeTools := ctx.Value(MCPContextKeyIncludeTools) - logger.Debug(fmt.Sprintf("%s Checking if tool %s should be skipped for request: %v", MCPLogPrefix, toolName, includeTools)) - if includeTools != nil { // Try []string first (preferred type) if includeToolsList, ok := includeTools.([]string); ok { diff --git a/docs/features/governance/virtual-keys.mdx b/docs/features/governance/virtual-keys.mdx index 2c4dcfdc58..0644fe7f23 100644 --- a/docs/features/governance/virtual-keys.mdx +++ b/docs/features/governance/virtual-keys.mdx @@ -564,7 +564,7 @@ curl -X POST http://localhost:8080/v1/chat/completions \ { "error": { "type": "virtual_key_required", - "message": "x-bf-vk header is missing" + "message": "virtual key is missing in headers" } } ``` diff --git a/plugins/governance/main.go b/plugins/governance/main.go index 1e2f373bfc..27a6b321c5 100644 --- a/plugins/governance/main.go +++ b/plugins/governance/main.go @@ -342,7 +342,7 @@ func (p *GovernancePlugin) PreHook(ctx *context.Context, req *schemas.BifrostReq Type: bifrost.Ptr("virtual_key_required"), StatusCode: bifrost.Ptr(400), Error: &schemas.ErrorField{ - Message: "x-bf-vk header is missing", + Message: "virtual key is missing in headers", }, }, }, nil diff --git a/plugins/logging/main.go b/plugins/logging/main.go index 5f1df57c5b..07b2c3aa31 100644 --- a/plugins/logging/main.go +++ b/plugins/logging/main.go @@ -163,8 +163,9 @@ func (p *LoggerPlugin) cleanupWorker() { // cleanupOldProcessingLogs removes processing logs older than 30 minutes func (p *LoggerPlugin) cleanupOldProcessingLogs() { - // Calculate timestamp for 30 minutes ago - thirtyMinutesAgo := time.Now().Add(-1 * 30 * time.Minute) + // Calculate timestamp for 30 minutes ago in UTC to match log entry timestamps + thirtyMinutesAgo := time.Now().UTC().Add(-1 * 30 * time.Minute) + p.logger.Debug("cleaning up old processing logs before %s", thirtyMinutesAgo) // Delete processing logs older than 30 minutes using the store if err := p.store.Flush(p.ctx, thirtyMinutesAgo); err != nil { p.logger.Warn("failed to cleanup old processing logs: %v", err) @@ -305,7 +306,7 @@ func (p *LoggerPlugin) PreHook(ctx *context.Context, req *schemas.BifrostRequest CreatedAt: msg.Timestamp, } p.logCallback(initialEntry) - } + } } }(logMsg) diff --git a/transports/bifrost-http/handlers/inference.go b/transports/bifrost-http/handlers/inference.go index d202768ac7..2111105d93 100644 --- a/transports/bifrost-http/handlers/inference.go +++ b/transports/bifrost-http/handlers/inference.go @@ -948,7 +948,6 @@ func (h *CompletionHandler) handleStreamingResponse(ctx *fasthttp.RequestCtx, ge ctx.SetContentType("text/event-stream") ctx.Response.Header.Set("Cache-Control", "no-cache") ctx.Response.Header.Set("Connection", "keep-alive") - ctx.Response.Header.Set("Access-Control-Allow-Origin", "*") // Get the streaming channel stream, bifrostErr := getStream() diff --git a/transports/bifrost-http/handlers/mcp_server.go b/transports/bifrost-http/handlers/mcp_server.go new file mode 100644 index 0000000000..1ccf8dfc52 --- /dev/null +++ b/transports/bifrost-http/handlers/mcp_server.go @@ -0,0 +1,391 @@ +// Package handlers provides HTTP request handlers for the Bifrost HTTP transport. +// This file contains MCP (Model Context Protocol) server implementation for HTTP streaming. +package handlers + +import ( + "bufio" + "context" + "encoding/json" + "fmt" + "slices" + "strings" + "sync" + + "github.com/fasthttp/router" + "github.com/mark3labs/mcp-go/mcp" + "github.com/mark3labs/mcp-go/server" + bifrost "github.com/maximhq/bifrost/core" + "github.com/maximhq/bifrost/core/schemas" + "github.com/maximhq/bifrost/framework/configstore/tables" + "github.com/maximhq/bifrost/plugins/governance" + "github.com/maximhq/bifrost/transports/bifrost-http/lib" + "github.com/valyala/fasthttp" +) + +// MCPToolExecutor interface defines the method needed for executing MCP tools +type MCPToolManager interface { + GetAvailableMCPTools(ctx context.Context) []schemas.ChatTool + ExecuteTool(ctx context.Context, toolCall schemas.ChatAssistantMessageToolCall) (*schemas.ChatMessage, *schemas.BifrostError) +} + +// MCPServerHandler manages HTTP requests for MCP server operations +// It implements the MCP protocol over HTTP streaming (SSE) for MCP clients +type MCPServerHandler struct { + toolManager MCPToolManager + globalMCPServer *server.MCPServer + vkMCPServers map[string]*server.MCPServer // Map of vk value -> mcp server + config *lib.Config + mu sync.RWMutex +} + +// NewMCPServerHandler creates a new MCP server handler instance +func NewMCPServerHandler(ctx context.Context, config *lib.Config, toolManager MCPToolManager) (*MCPServerHandler, error) { + if config == nil { + return nil, fmt.Errorf("config is required") + } + if toolManager == nil { + return nil, fmt.Errorf("tool manager is required") + } + + // Create MCP server instance using mcp-go + globalMCPServer := server.NewMCPServer( + "global", + version, + server.WithToolCapabilities(true), + ) + + handler := &MCPServerHandler{ + toolManager: toolManager, + globalMCPServer: globalMCPServer, + config: config, + vkMCPServers: make(map[string]*server.MCPServer), + } + + if err := handler.SyncAllMCPServers(ctx); err != nil { + return nil, fmt.Errorf("failed to sync all MCP servers: %w", err) + } + + return handler, nil +} + +// RegisterRoutes registers the MCP server route +func (h *MCPServerHandler) RegisterRoutes(r *router.Router, middlewares ...lib.BifrostHTTPMiddleware) { + // MCP server endpoint - supports both POST (JSON-RPC) and GET (SSE) + r.POST("/mcp", lib.ChainMiddlewares(h.handleMCPServer, middlewares...)) + r.GET("/mcp", lib.ChainMiddlewares(h.handleMCPServerSSE, middlewares...)) +} + +// handleMCPServer handles POST requests for MCP JSON-RPC 2.0 messages +func (h *MCPServerHandler) handleMCPServer(ctx *fasthttp.RequestCtx) { + mcpServer, err := h.getMCPServerForRequest(ctx) + if err != nil { + SendError(ctx, fasthttp.StatusUnauthorized, err.Error()) + return + } + + // Convert context + bifrostCtx, cancel := lib.ConvertToBifrostContext(ctx, false) + defer cancel() + + // Use mcp-go server to handle the request + // HandleMessage processes JSON-RPC messages and returns appropriate responses + response := mcpServer.HandleMessage(*bifrostCtx, ctx.PostBody()) + + // Check if response is nil (notification - no response needed) + if response == nil { + ctx.SetStatusCode(fasthttp.StatusOK) + return + } + + // Marshal and send response + responseJSON, err := json.Marshal(response) + if err != nil { + logger.Warn(fmt.Sprintf("Failed to marshal MCP response: %v", err)) + SendError(ctx, fasthttp.StatusInternalServerError, fmt.Sprintf("Failed to encode response: %v", err)) + return + } + + ctx.SetContentType("application/json") + ctx.SetBody(responseJSON) +} + +// handleMCPServerSSE handles GET requests for MCP Server-Sent Events streaming +func (h *MCPServerHandler) handleMCPServerSSE(ctx *fasthttp.RequestCtx) { + _, err := h.getMCPServerForRequest(ctx) + if err != nil { + SendError(ctx, fasthttp.StatusUnauthorized, err.Error()) + return + } + + // Set SSE headers + ctx.SetContentType("text/event-stream") + ctx.Response.Header.Set("Cache-Control", "no-cache") + ctx.Response.Header.Set("Connection", "keep-alive") + + // Convert context + bifrostCtx, cancel := lib.ConvertToBifrostContext(ctx, false) + + // Use streaming response writer + ctx.Response.SetBodyStreamWriter(func(w *bufio.Writer) { + defer func() { + cancel() + _ = w.Flush() + }() + + // Send initial connection message + initMessage := map[string]interface{}{ + "jsonrpc": "2.0", + "method": "connection/opened", + } + if initJSON, err := json.Marshal(initMessage); err == nil { + fmt.Fprintf(w, "data: %s\n\n", initJSON) + w.Flush() + } + + // Wait for context cancellation (client disconnect or server-side cancel) + <-(*bifrostCtx).Done() + }) +} + +// Sync methods for MCP servers + +func (h *MCPServerHandler) SyncAllMCPServers(ctx context.Context) error { + h.mu.Lock() + defer h.mu.Unlock() + availableTools := h.toolManager.GetAvailableMCPTools(ctx) + h.syncServer(h.globalMCPServer, availableTools) + logger.Debug("Synced global MCP server with %d tools", len(availableTools)) + + // initialize vkMCPServers map + if h.config.ConfigStore != nil { + virtualKeys, err := h.config.ConfigStore.GetVirtualKeys(ctx) + if err != nil { + return fmt.Errorf("failed to get virtual keys: %w", err) + } + h.vkMCPServers = make(map[string]*server.MCPServer) + for i := range virtualKeys { + vk := &virtualKeys[i] + h.vkMCPServers[vk.Value] = server.NewMCPServer( + vk.Name, + version, + server.WithToolCapabilities(true), + ) + availableTools := h.fetchToolsForVK(vk) + h.syncServer(h.vkMCPServers[vk.Value], availableTools) + logger.Debug("Synced MCP server for virtual key '%s' with %d tools", vk.Name, len(availableTools)) + } + } + return nil +} + +func (h *MCPServerHandler) SyncVKMCPServer(vk *tables.TableVirtualKey) { + h.mu.Lock() + defer h.mu.Unlock() + vkServer, ok := h.vkMCPServers[vk.Value] + if !ok { + // Add new server + vkServer = server.NewMCPServer( + vk.Name, + version, + server.WithToolCapabilities(true), + ) + h.vkMCPServers[vk.Value] = vkServer + } + availableTools := h.fetchToolsForVK(vk) + h.syncServer(vkServer, availableTools) + h.vkMCPServers[vk.Value] = vkServer + logger.Debug("Synced MCP server for virtual key '%s' with %d tools", vk.Name, len(availableTools)) +} + +func (h *MCPServerHandler) DeleteVKMCPServer(vkValue string) { + h.mu.Lock() + defer h.mu.Unlock() + delete(h.vkMCPServers, vkValue) +} + +func (h *MCPServerHandler) syncServer(server *server.MCPServer, availableTools []schemas.ChatTool) { + // Clear existing tools + toolMap := server.ListTools() + for toolName, _ := range toolMap { + server.DeleteTools(toolName) + } + + // Register tools from all connected clients + for _, tool := range availableTools { + // Only process function tools (skip custom tools) + if tool.Function == nil { + continue + } + + // Capture tool name for closure + toolName := tool.Function.Name + + handler := func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { + // Convert to Bifrost tool call format + toolCallType := "function" + toolCallID := fmt.Sprintf("mcp-%s", toolName) + argsJSON, jsonErr := json.Marshal(request.GetArguments()) + if jsonErr != nil { + return mcp.NewToolResultError(fmt.Sprintf("Failed to marshal tool arguments: %v", jsonErr)), nil + } + toolCall := schemas.ChatAssistantMessageToolCall{ + ID: &toolCallID, + Type: &toolCallType, + Function: schemas.ChatAssistantMessageToolCallFunction{ + Name: &toolName, + Arguments: string(argsJSON), + }, + } + + // Execute the tool via tool executor + toolMessage, err := h.toolManager.ExecuteTool(ctx, toolCall) + if err != nil { + return mcp.NewToolResultError(fmt.Sprintf("Tool execution failed: %v", bifrost.GetErrorMessage(err))), nil + } + + // Extract content from tool message + var resultText string + if toolMessage != nil && toolMessage.Content != nil { + // Handle ContentStr (string content) + if toolMessage.Content.ContentStr != nil { + resultText = *toolMessage.Content.ContentStr + } else if toolMessage.Content.ContentBlocks != nil { + // Handle ContentBlocks (structured content) + for _, block := range toolMessage.Content.ContentBlocks { + if block.Type == schemas.ChatContentBlockTypeText && block.Text != nil { + resultText += *block.Text + } + } + } + } + + // Return result using mcp-go helper + return mcp.NewToolResultText(resultText), nil + } + + // Convert description from *string to string + description := "" + if tool.Function.Description != nil { + description = *tool.Function.Description + } + + // Convert Parameters to mcp.ToolInputSchema + var inputSchema mcp.ToolInputSchema + if tool.Function.Parameters != nil { + inputSchema.Type = tool.Function.Parameters.Type + if tool.Function.Parameters.Properties != nil { + // Convert *map[string]interface{} to map[string]any + props := make(map[string]any) + for k, v := range *tool.Function.Parameters.Properties { + props[k] = v + } + inputSchema.Properties = props + } + if tool.Function.Parameters.Required != nil { + inputSchema.Required = tool.Function.Parameters.Required + } + } else { + // Default to empty object schema if no parameters + inputSchema.Type = "object" + inputSchema.Properties = make(map[string]any) + } + + // Register tool with the server + server.AddTool(mcp.Tool{ + Name: toolName, + Description: description, + InputSchema: inputSchema, + }, handler) + } +} + +// fetchToolsForVK fetches the tools for a given virtual key value. +// vkValue is the virtual key value for the server, if empty, all tools will be fetched for global mcp server. +// Returns a map of tool name to tool. +func (h *MCPServerHandler) fetchToolsForVK(vk *tables.TableVirtualKey) []schemas.ChatTool { + ctx := context.Background() + + if len(vk.MCPConfigs) > 0 { + executeOnlyTools := make([]string, 0) + for _, vkMcpConfig := range vk.MCPConfigs { + if len(vkMcpConfig.ToolsToExecute) == 0 { + // No tools specified in virtual key config - skip this client entirely + continue + } + + // Handle wildcard in virtual key config - allow all tools from this client + if slices.Contains(vkMcpConfig.ToolsToExecute, "*") { + // Virtual key uses wildcard - use client-specific wildcard + executeOnlyTools = append(executeOnlyTools, fmt.Sprintf("%s/*", vkMcpConfig.MCPClient.Name)) + continue + } + + for _, tool := range vkMcpConfig.ToolsToExecute { + if tool != "" { + // Add the tool - client config filtering will be handled by mcp.go + executeOnlyTools = append(executeOnlyTools, fmt.Sprintf("%s/%s", vkMcpConfig.MCPClient.Name, tool)) + } + } + } + + // Set even when empty to exclude tools when no tools are present in the virtual key config + ctx = context.WithValue(ctx, schemas.BifrostContextKey("mcp-include-tools"), executeOnlyTools) + } + + return h.toolManager.GetAvailableMCPTools(ctx) +} + +// Utility methods + +func (h *MCPServerHandler) getMCPServerForRequest(ctx *fasthttp.RequestCtx) (*server.MCPServer, error) { + h.mu.RLock() + defer h.mu.RUnlock() + + h.config.Mu.RLock() + enforceVK := h.config.ClientConfig.EnforceGovernanceHeader + h.config.Mu.RUnlock() + + vk := getVKFromRequest(ctx) + + // Return global MCP server if not enforcing virtual key header and no virtual key is provided + if !enforceVK && vk == "" { + return h.globalMCPServer, nil + } + + // Check if virtual key is provided + if vk == "" { + return nil, fmt.Errorf("virtual key header is required to access MCP server.") + } + + // Check if vk exists in the map + vkServer, ok := h.vkMCPServers[vk] + if !ok { + return nil, fmt.Errorf("virtual key not found.") + } + + return vkServer, nil +} + +func getVKFromRequest(ctx *fasthttp.RequestCtx) string { + if value := strings.TrimSpace(string(ctx.Request.Header.Peek(string(schemas.BifrostContextKeyVirtualKey)))); value != "" { + return value + } + + authHeader := strings.TrimSpace(string(ctx.Request.Header.Peek("Authorization"))) + if authHeader != "" { + if strings.HasPrefix(strings.ToLower(authHeader), "bearer ") { + token := strings.TrimSpace(authHeader[7:]) + if token != "" && strings.HasPrefix(strings.ToLower(token), governance.VirtualKeyPrefix) { + return token + } + } + } + + if apiKey := strings.TrimSpace(string(ctx.Request.Header.Peek("x-api-key"))); apiKey != "" { + if strings.HasPrefix(strings.ToLower(apiKey), governance.VirtualKeyPrefix) { + return apiKey + } + } + + return "" +} diff --git a/transports/bifrost-http/integrations/anthropic.go b/transports/bifrost-http/integrations/anthropic.go index db44cf6ed3..1d6bba8e90 100644 --- a/transports/bifrost-http/integrations/anthropic.go +++ b/transports/bifrost-http/integrations/anthropic.go @@ -173,9 +173,26 @@ func checkAnthropicPassthrough(ctx *fasthttp.RequestCtx, bifrostCtx *context.Con return nil } + headers := extractHeadersFromRequest(ctx) + if len(headers) > 0 { + // Check for User-Agent header (case-insensitive) + var userAgent []string + for key, value := range headers { + if strings.EqualFold(key, "user-agent") { + userAgent = value + break + } + } + if len(userAgent) > 0 { + // Check if it's claude code + if strings.Contains(userAgent[0], "claude-cli") { + *bifrostCtx = context.WithValue(*bifrostCtx, schemas.BifrostContextKey("integration-user-agent"), "claude-cli") + } + } + } + // Check if anthropic oauth headers are present if !isAnthropicAPIKeyAuth(ctx) { - headers := extractHeadersFromRequest(ctx) url := extractExactPath(ctx) if !strings.HasPrefix(url, "/") { url = "/" + url diff --git a/transports/bifrost-http/server/server.go b/transports/bifrost-http/server/server.go index 1e962be5ad..f94bb5e24f 100644 --- a/transports/bifrost-http/server/server.go +++ b/transports/bifrost-http/server/server.go @@ -101,6 +101,7 @@ type BifrostHTTPServer struct { Router *router.Router WebSocketHandler *handlers.WebSocketHandler LogsCleaner *logstore.LogsCleaner + MCPServerHandler *handlers.MCPServerHandler } var logger schemas.Logger @@ -445,17 +446,44 @@ func FindPluginByName[T schemas.Plugin](plugins []schemas.Plugin, name string) ( // AddMCPClient adds a new MCP client to the in-memory store func (s *BifrostHTTPServer) AddMCPClient(ctx context.Context, clientConfig schemas.MCPClientConfig) error { - return s.Config.AddMCPClient(ctx, clientConfig) + if err := s.Config.AddMCPClient(ctx, clientConfig); err != nil { + return err + } + if err := s.MCPServerHandler.SyncAllMCPServers(ctx); err != nil { + logger.Warn("failed to sync MCP servers after adding client: %v", err) + } + return nil +} + +// EditMCPClient edits an MCP client in the in-memory store +func (s *BifrostHTTPServer) EditMCPClient(ctx context.Context, id string, updatedConfig schemas.MCPClientConfig) error { + if err := s.Config.EditMCPClient(ctx, id, updatedConfig); err != nil { + return err + } + if err := s.MCPServerHandler.SyncAllMCPServers(ctx); err != nil { + logger.Warn("failed to sync MCP servers after editing client: %v", err) + } + return nil } // RemoveMCPClient removes an MCP client from the in-memory store func (s *BifrostHTTPServer) RemoveMCPClient(ctx context.Context, id string) error { - return s.Config.RemoveMCPClient(ctx, id) + if err := s.Config.RemoveMCPClient(ctx, id); err != nil { + return err + } + if err := s.MCPServerHandler.SyncAllMCPServers(ctx); err != nil { + logger.Warn("failed to sync MCP servers after removing client: %v", err) + } + return nil } -// EditMCPClient edits an MCP client in the in-memory store -func (s *BifrostHTTPServer) EditMCPClient(ctx context.Context, id string, updatedConfig schemas.MCPClientConfig) error { - return s.Config.EditMCPClient(ctx, id, updatedConfig) +// ExecuteTool executes an MCP tool call and returns the result +func (s *BifrostHTTPServer) ExecuteTool(ctx context.Context, toolCall schemas.ChatAssistantMessageToolCall) (*schemas.ChatMessage, *schemas.BifrostError) { + return s.Client.ExecuteMCPTool(ctx, toolCall) +} + +func (s *BifrostHTTPServer) GetAvailableMCPTools(ctx context.Context) []schemas.ChatTool { + return s.Client.GetAvailableMCPTools(ctx) } // ReloadVirtualKey reloads a virtual key from the in-memory store @@ -487,6 +515,7 @@ func (s *BifrostHTTPServer) ReloadVirtualKey(ctx context.Context, id string) (*t } } } + s.MCPServerHandler.SyncVKMCPServer(preloadedVk) return preloadedVk, nil } @@ -528,6 +557,7 @@ func (s *BifrostHTTPServer) RemoveVirtualKey(ctx context.Context, id string) err } } } + s.MCPServerHandler.DeleteVKMCPServer(preloadedVk.Value) return nil } @@ -909,6 +939,11 @@ func (s *BifrostHTTPServer) RegisterAPIRoutes(ctx context.Context, callbacks Ser healthHandler := handlers.NewHealthHandler(s.Config) providerHandler := handlers.NewProviderHandler(callbacks, s.Config, s.Client) mcpHandler := handlers.NewMCPHandler(callbacks, s.Client, s.Config) + mcpServerHandler, err := handlers.NewMCPServerHandler(ctx, s.Config, s) + if err != nil { + return fmt.Errorf("failed to initialize mcp server handler: %v", err) + } + s.MCPServerHandler = mcpServerHandler configHandler := handlers.NewConfigHandler(callbacks, s.Config) pluginsHandler := handlers.NewPluginsHandler(callbacks, s.Config.ConfigStore) sessionHandler := handlers.NewSessionHandler(s.Config.ConfigStore) @@ -916,6 +951,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...) + mcpServerHandler.RegisterRoutes(s.Router, middlewares...) configHandler.RegisterRoutes(s.Router, middlewares...) if pluginsHandler != nil { pluginsHandler.RegisterRoutes(s.Router, middlewares...) diff --git a/transports/go.mod b/transports/go.mod index eccee260f0..25bda6c531 100644 --- a/transports/go.mod +++ b/transports/go.mod @@ -8,6 +8,7 @@ require ( github.com/fasthttp/router v1.5.4 github.com/fasthttp/websocket v1.5.12 github.com/google/uuid v1.6.0 + github.com/mark3labs/mcp-go v0.41.1 github.com/maximhq/bifrost/core v1.2.30 github.com/maximhq/bifrost/framework v1.1.39 github.com/maximhq/bifrost/plugins/governance v1.3.40 @@ -82,7 +83,6 @@ require ( github.com/klauspost/compress v1.18.1 // indirect github.com/klauspost/cpuid/v2 v2.3.0 // indirect github.com/mailru/easyjson v0.9.1 // indirect - github.com/mark3labs/mcp-go v0.41.1 // indirect github.com/mattn/go-colorable v0.1.14 // indirect github.com/mattn/go-isatty v0.0.20 // indirect github.com/mattn/go-sqlite3 v1.14.32 // indirect diff --git a/ui/app/workspace/mcp-gateway/views/mcpClientSheet.tsx b/ui/app/workspace/mcp-gateway/views/mcpClientSheet.tsx index 57537d0289..af7c23cede 100644 --- a/ui/app/workspace/mcp-gateway/views/mcpClientSheet.tsx +++ b/ui/app/workspace/mcp-gateway/views/mcpClientSheet.tsx @@ -183,8 +183,8 @@ export default function MCPClientSheet({ mcpClient, onClose, onSubmitSuccess }:
- -
+ +
{mcpClient.config.name} diff --git a/ui/app/workspace/mcp-gateway/views/mcpClientsTable.tsx b/ui/app/workspace/mcp-gateway/views/mcpClientsTable.tsx index 0404339978..bb988c77b5 100644 --- a/ui/app/workspace/mcp-gateway/views/mcpClientsTable.tsx +++ b/ui/app/workspace/mcp-gateway/views/mcpClientsTable.tsx @@ -155,7 +155,7 @@ export default function MCPClientsTable({ mcpClients }: MCPClientsTableProps) { {clients.length === 0 && ( - + No clients found. diff --git a/ui/components/ui/sheet.tsx b/ui/components/ui/sheet.tsx index 6b68509f78..63b6fb33c0 100644 --- a/ui/components/ui/sheet.tsx +++ b/ui/components/ui/sheet.tsx @@ -72,11 +72,11 @@ function SheetContent({ className={cn( "bg-background data-[state=open]:animate-in data-[state=closed]:animate-out fixed z-50 flex flex-col shadow-lg transition-all ease-in-out data-[state=closed]:duration-100 data-[state=open]:duration-100", side === "right" && - "data-[state=closed]:slide-out-to-right data-[state=open]:slide-in-from-right top-2 bottom-2 right-0 h-auto w-3/4 border-l rounded-l-lg", + "data-[state=closed]:slide-out-to-right data-[state=open]:slide-in-from-right top-2 right-0 bottom-2 h-auto w-3/4 rounded-l-lg border-l", side === "right" && (!expandable || !expanded) && "sm:max-w-2xl", side === "right" && expandable && expanded && "sm:max-w-5xl", side === "left" && - "data-[state=closed]:slide-out-to-left data-[state=open]:slide-in-from-left top-2 bottom-2 left-0 h-auto w-3/4 border-r rounded-r-lg sm:max-w-sm", + "data-[state=closed]:slide-out-to-left data-[state=open]:slide-in-from-left top-2 bottom-2 left-0 h-auto w-3/4 rounded-r-lg border-r sm:max-w-sm", side === "top" && "data-[state=closed]:slide-out-to-top data-[state=open]:slide-in-from-top inset-x-0 top-0 h-auto border-b", side === "bottom" && "data-[state=closed]:slide-out-to-bottom data-[state=open]:slide-in-from-bottom inset-x-0 bottom-0 h-auto border-t", @@ -91,36 +91,33 @@ function SheetContent({ ); } -function SheetHeader({ className, children, ...props }: React.ComponentProps<"div">) { +function SheetHeader({ + className, + children, + showCloseButton = true, + ...props +}: React.ComponentProps<"div"> & { showCloseButton?: boolean }) { const sheetContext = useSheetContext(); return ( -
+
{sheetContext?.expandable && sheetContext?.side === "right" && ( )} -
- {children} -
- - - Close - +
{children}
+ {showCloseButton && ( + + + Close + + )}
); } @@ -138,4 +135,3 @@ function SheetDescription({ className, ...props }: React.ComponentProps