first commit
This commit is contained in:
Beyhan Oğur
346
core/mcp/mcp.go
Normal file
346
core/mcp/mcp.go
Normal file
@@ -0,0 +1,346 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/maximhq/bifrost/core/schemas"
|
||||
|
||||
"github.com/mark3labs/mcp-go/server"
|
||||
)
|
||||
|
||||
// ============================================================================
|
||||
// CONSTANTS
|
||||
// ============================================================================
|
||||
|
||||
const (
|
||||
// MCP defaults and identifiers
|
||||
BifrostMCPVersion = "1.0.0" // Version identifier for Bifrost
|
||||
BifrostMCPClientName = "BifrostClient" // Name for internal Bifrost MCP client
|
||||
BifrostMCPClientKey = "bifrostInternal" // Key for internal Bifrost client in clientMap
|
||||
MCPLogPrefix = "[Bifrost MCP]" // Consistent logging prefix
|
||||
MCPClientConnectionEstablishTimeout = 30 * time.Second // Timeout for MCP client connection establishment
|
||||
)
|
||||
|
||||
// ============================================================================
|
||||
// TYPE DEFINITIONS
|
||||
// ============================================================================
|
||||
|
||||
// MCPManager manages MCP integration for Bifrost core.
|
||||
// It provides a bridge between Bifrost and various MCP servers, supporting
|
||||
// both local tool hosting and external MCP server connections.
|
||||
type MCPManager struct {
|
||||
ctx context.Context
|
||||
logger schemas.Logger // Logger instance for this manager
|
||||
oauth2Provider schemas.OAuth2Provider // Provider for OAuth2 functionality
|
||||
toolsManager *ToolsManager // Handler for MCP tools
|
||||
server *server.MCPServer // Local MCP server instance for hosting tools (STDIO-based)
|
||||
clientMap map[string]*schemas.MCPClientState // Map of MCP client names to their configurations
|
||||
mu sync.RWMutex // Read-write mutex for thread-safe operations
|
||||
serverRunning bool // Track whether local MCP server is running
|
||||
healthMonitorManager *HealthMonitorManager // Manager for client health monitors
|
||||
toolSyncManager *ToolSyncManager // Manager for periodic tool synchronization
|
||||
reconnectingClients sync.Map // Tracks in-flight reconnect attempts per client ID (map[string]bool)
|
||||
}
|
||||
|
||||
// MCPToolFunction is a generic function type for handling tool calls with typed arguments.
|
||||
// T represents the expected argument structure for the tool.
|
||||
type MCPToolFunction[T any] func(args T) (string, error)
|
||||
|
||||
// ============================================================================
|
||||
// CONSTRUCTOR AND INITIALIZATION
|
||||
// ============================================================================
|
||||
|
||||
// NewMCPManager creates and initializes a new MCP manager instance.
|
||||
//
|
||||
// Parameters:
|
||||
// - ctx: Context for the MCP manager
|
||||
// - config: MCP configuration including server port and client configs
|
||||
// - oauth2Provider: OAuth2 provider for authentication
|
||||
// - logger: Logger instance for structured logging (uses default if nil)
|
||||
// - codeMode: Optional CodeMode implementation for code execution (e.g., Starlark).
|
||||
// Pass nil if code mode is not needed. The CodeMode's dependencies will be
|
||||
// injected automatically via SetDependencies after the manager is created.
|
||||
//
|
||||
// Returns:
|
||||
// - *MCPManager: Initialized manager instance
|
||||
func NewMCPManager(ctx context.Context, config schemas.MCPConfig, oauth2Provider schemas.OAuth2Provider, logger schemas.Logger, codeMode CodeMode) *MCPManager {
|
||||
if logger == nil {
|
||||
logger = defaultLogger
|
||||
}
|
||||
// Set default values
|
||||
if config.ToolManagerConfig == nil {
|
||||
config.ToolManagerConfig = &schemas.MCPToolManagerConfig{
|
||||
ToolExecutionTimeout: schemas.DefaultToolExecutionTimeout,
|
||||
MaxAgentDepth: schemas.DefaultMaxAgentDepth,
|
||||
}
|
||||
}
|
||||
// Creating new instance
|
||||
manager := &MCPManager{
|
||||
ctx: ctx,
|
||||
logger: logger,
|
||||
clientMap: make(map[string]*schemas.MCPClientState),
|
||||
healthMonitorManager: NewHealthMonitorManager(),
|
||||
toolSyncManager: NewToolSyncManager(config.ToolSyncInterval),
|
||||
oauth2Provider: oauth2Provider,
|
||||
}
|
||||
// Convert plugin pipeline provider functions to the interface expected by ToolsManager
|
||||
var pluginPipelineProvider func() PluginPipeline
|
||||
var releasePluginPipeline func(pipeline PluginPipeline)
|
||||
|
||||
if config.PluginPipelineProvider != nil && config.ReleasePluginPipeline != nil {
|
||||
pluginPipelineProvider = func() PluginPipeline {
|
||||
if pipeline := config.PluginPipelineProvider(); pipeline != nil {
|
||||
if pp, ok := pipeline.(PluginPipeline); ok {
|
||||
return pp
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
releasePluginPipeline = func(pipeline PluginPipeline) {
|
||||
config.ReleasePluginPipeline(pipeline)
|
||||
}
|
||||
}
|
||||
|
||||
manager.toolsManager = NewToolsManager(config.ToolManagerConfig, manager, config.FetchNewRequestIDFunc, pluginPipelineProvider, releasePluginPipeline, oauth2Provider, logger)
|
||||
|
||||
// Set up CodeMode if provided - inject dependencies after manager is created
|
||||
if codeMode != nil {
|
||||
deps := manager.toolsManager.GetCodeModeDependencies()
|
||||
codeMode.SetDependencies(deps)
|
||||
manager.toolsManager.SetCodeMode(codeMode)
|
||||
}
|
||||
|
||||
// Process client configs: create client map entries and establish connections
|
||||
if len(config.ClientConfigs) > 0 {
|
||||
// Add clients in parallel
|
||||
wg := sync.WaitGroup{}
|
||||
wg.Add(len(config.ClientConfigs))
|
||||
for _, clientConfig := range config.ClientConfigs {
|
||||
go func(clientConfig *schemas.MCPClientConfig) {
|
||||
defer wg.Done()
|
||||
if err := manager.AddClient(clientConfig); err != nil {
|
||||
manager.logger.Warn("%s Failed to register MCP client %s: %v", MCPLogPrefix, clientConfig.Name, err)
|
||||
// Retain the entry in Disconnected state and start a health monitor to
|
||||
// recover it automatically. On startup, a connection failure is likely
|
||||
// transient (e.g. autoscaling cold start) — the client was previously
|
||||
// configured and should be recovered without user intervention.
|
||||
manager.mu.Lock()
|
||||
if _, exists := manager.clientMap[clientConfig.ID]; !exists {
|
||||
manager.clientMap[clientConfig.ID] = &schemas.MCPClientState{
|
||||
Name: clientConfig.Name,
|
||||
ExecutionConfig: clientConfig,
|
||||
State: schemas.MCPConnectionStateDisconnected,
|
||||
ToolMap: make(map[string]schemas.ChatTool),
|
||||
ToolNameMapping: make(map[string]string),
|
||||
ConnectionInfo: &schemas.MCPClientConnectionInfo{
|
||||
Type: clientConfig.ConnectionType,
|
||||
},
|
||||
}
|
||||
} else {
|
||||
manager.clientMap[clientConfig.ID].State = schemas.MCPConnectionStateDisconnected
|
||||
}
|
||||
manager.mu.Unlock()
|
||||
isPingAvailable := true
|
||||
if clientConfig.IsPingAvailable != nil {
|
||||
isPingAvailable = *clientConfig.IsPingAvailable
|
||||
}
|
||||
monitor := NewClientHealthMonitor(manager, clientConfig.ID, DefaultHealthCheckInterval, isPingAvailable, manager.logger)
|
||||
manager.healthMonitorManager.StartMonitoring(monitor)
|
||||
}
|
||||
}(clientConfig)
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
manager.logger.Info(MCPLogPrefix + " MCP Manager initialized")
|
||||
return manager
|
||||
}
|
||||
|
||||
// SetPluginPipeline updates the plugin pipeline provider and release function on the manager's
|
||||
// ToolsManager and CodeMode. Call this after attaching an externally-created MCPManager to a Bifrost
|
||||
// instance so that nested tool calls in code mode can run through Bifrost's plugin hooks.
|
||||
func (manager *MCPManager) SetPluginPipeline(provider func() PluginPipeline, release func(PluginPipeline)) {
|
||||
manager.toolsManager.SetPluginPipeline(provider, release)
|
||||
}
|
||||
|
||||
// AddToolsToRequest parses available MCP tools from the context and adds them to the request.
|
||||
// It respects context-based filtering for clients and tools, and returns the modified request
|
||||
// with tools attached.
|
||||
//
|
||||
// Parameters:
|
||||
// - ctx: Context containing optional client/tool filtering keys
|
||||
// - req: The Bifrost request to add tools to
|
||||
//
|
||||
// Returns:
|
||||
// - *schemas.BifrostRequest: The request with tools added
|
||||
func (m *MCPManager) AddToolsToRequest(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) *schemas.BifrostRequest {
|
||||
return m.toolsManager.ParseAndAddToolsToRequest(ctx, req)
|
||||
}
|
||||
|
||||
func (m *MCPManager) GetAvailableTools(ctx *schemas.BifrostContext) []schemas.ChatTool {
|
||||
return m.toolsManager.GetAvailableTools(ctx)
|
||||
}
|
||||
|
||||
// ExecuteToolCall executes a single tool call and returns the result.
|
||||
// This is the primary tool executor and is used by both Chat Completions and Responses APIs.
|
||||
//
|
||||
// The method accepts an MCP request containing either a ChatAssistantMessageToolCall or
|
||||
// ResponsesToolMessage, and returns the appropriate result format based on the request type.
|
||||
//
|
||||
// Parameters:
|
||||
// - ctx: Context for the tool execution
|
||||
// - request: The MCP request containing the tool call (ChatAssistantMessageToolCall or ResponsesToolMessage)
|
||||
//
|
||||
// Returns:
|
||||
// - *schemas.BifrostMCPResponse: The result response containing tool execution output (ChatMessage or ResponsesMessage)
|
||||
// - error: Any error that occurred during tool execution
|
||||
func (m *MCPManager) ExecuteToolCall(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error) {
|
||||
return m.toolsManager.ExecuteTool(ctx, request)
|
||||
}
|
||||
|
||||
// UpdateToolManagerConfig updates the configuration for the tool manager.
|
||||
// This allows runtime updates to settings like execution timeout and max agent depth.
|
||||
//
|
||||
// Parameters:
|
||||
// - config: The new tool manager configuration to apply
|
||||
func (m *MCPManager) UpdateToolManagerConfig(config *schemas.MCPToolManagerConfig) {
|
||||
m.toolsManager.UpdateConfig(config)
|
||||
}
|
||||
|
||||
// CheckAndExecuteAgentForChatRequest checks if the chat response contains tool calls,
|
||||
// and if so, executes agent mode to handle the tool calls iteratively. If no tool calls
|
||||
// are present, it returns the original response unchanged.
|
||||
//
|
||||
// Agent mode enables autonomous tool execution where:
|
||||
// 1. Tool calls are automatically executed
|
||||
// 2. Results are fed back to the LLM
|
||||
// 3. The loop continues until no more tool calls are made or max depth is reached
|
||||
// 4. Non-auto-executable tools are returned to the caller
|
||||
//
|
||||
// This method is available for both Chat Completions and Responses APIs.
|
||||
// For Responses API, use CheckAndExecuteAgentForResponsesRequest().
|
||||
//
|
||||
// Parameters:
|
||||
// - ctx: Context for the agent execution
|
||||
// - req: The original chat request
|
||||
// - response: The initial chat response that may contain tool calls
|
||||
// - makeReq: Function to make subsequent chat requests during agent execution
|
||||
//
|
||||
// Returns:
|
||||
// - *schemas.BifrostChatResponse: The final response after agent execution (or original if no tool calls)
|
||||
// - *schemas.BifrostError: Any error that occurred during agent execution
|
||||
func (m *MCPManager) CheckAndExecuteAgentForChatRequest(
|
||||
ctx *schemas.BifrostContext,
|
||||
req *schemas.BifrostChatRequest,
|
||||
response *schemas.BifrostChatResponse,
|
||||
makeReq func(ctx *schemas.BifrostContext, req *schemas.BifrostChatRequest) (*schemas.BifrostChatResponse, *schemas.BifrostError),
|
||||
executeTool func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error),
|
||||
) (*schemas.BifrostChatResponse, *schemas.BifrostError) {
|
||||
if makeReq == nil {
|
||||
return nil, &schemas.BifrostError{
|
||||
IsBifrostError: false,
|
||||
Error: &schemas.ErrorField{
|
||||
Message: "makeReq is required to execute agent mode",
|
||||
},
|
||||
}
|
||||
}
|
||||
// Check if initial response has tool calls
|
||||
if !hasToolCallsForChatResponse(response) {
|
||||
m.logger.Debug("No tool calls detected, returning response")
|
||||
return response, nil
|
||||
}
|
||||
// Execute agent mode
|
||||
return m.toolsManager.ExecuteAgentForChatRequest(ctx, req, response, makeReq, executeTool)
|
||||
}
|
||||
|
||||
// CheckAndExecuteAgentForResponsesRequest checks if the responses response contains tool calls,
|
||||
// and if so, executes agent mode to handle the tool calls iteratively. If no tool calls
|
||||
// are present, it returns the original response unchanged.
|
||||
//
|
||||
// Agent mode for Responses API works identically to Chat API:
|
||||
// 1. Detects tool calls in the response (function_call messages)
|
||||
// 2. Automatically executes tools in parallel when possible
|
||||
// 3. Feeds results back to the LLM in Responses API format
|
||||
// 4. Continues the loop until no more tool calls or max depth reached
|
||||
// 5. Returns non-auto-executable tools to the caller
|
||||
//
|
||||
// Format Handling:
|
||||
// This method automatically handles format conversions:
|
||||
// - Responses tool calls (ResponsesToolMessage) are converted to Chat format for execution
|
||||
// - Tool execution results are converted back to Responses format (ResponsesMessage)
|
||||
// - All conversions use the adapters in agent_adaptors.go and converters in schemas/mux.go
|
||||
//
|
||||
// This provides full feature parity between Chat Completions and Responses APIs for tool execution.
|
||||
//
|
||||
// Parameters:
|
||||
// - ctx: Context for the agent execution
|
||||
// - req: The original responses request
|
||||
// - response: The initial responses response that may contain tool calls
|
||||
// - makeReq: Function to make subsequent responses requests during agent execution
|
||||
//
|
||||
// Returns:
|
||||
// - *schemas.BifrostResponsesResponse: The final response after agent execution (or original if no tool calls)
|
||||
// - *schemas.BifrostError: Any error that occurred during agent execution
|
||||
func (m *MCPManager) CheckAndExecuteAgentForResponsesRequest(
|
||||
ctx *schemas.BifrostContext,
|
||||
req *schemas.BifrostResponsesRequest,
|
||||
response *schemas.BifrostResponsesResponse,
|
||||
makeReq func(ctx *schemas.BifrostContext, req *schemas.BifrostResponsesRequest) (*schemas.BifrostResponsesResponse, *schemas.BifrostError),
|
||||
executeTool func(ctx *schemas.BifrostContext, request *schemas.BifrostMCPRequest) (*schemas.BifrostMCPResponse, error),
|
||||
) (*schemas.BifrostResponsesResponse, *schemas.BifrostError) {
|
||||
if makeReq == nil {
|
||||
return nil, &schemas.BifrostError{
|
||||
IsBifrostError: false,
|
||||
Error: &schemas.ErrorField{
|
||||
Message: "makeReq is required to execute agent mode",
|
||||
},
|
||||
}
|
||||
}
|
||||
// Check if initial response has tool calls
|
||||
if !hasToolCallsForResponsesResponse(response) {
|
||||
m.logger.Debug("No tool calls detected, returning response")
|
||||
return response, nil
|
||||
}
|
||||
// Execute agent mode
|
||||
return m.toolsManager.ExecuteAgentForResponsesRequest(ctx, req, response, makeReq, executeTool)
|
||||
}
|
||||
|
||||
// Cleanup performs cleanup of all MCP resources including clients and local server.
|
||||
// This function safely disconnects all MCP clients (HTTP, STDIO, and SSE) and
|
||||
// cleans up the local MCP server. It handles proper cancellation of SSE contexts
|
||||
// and closes all transport connections.
|
||||
//
|
||||
// Returns:
|
||||
// - error: Always returns nil, but maintains error interface for consistency
|
||||
func (m *MCPManager) Cleanup() error {
|
||||
// Stop all health monitors first
|
||||
m.healthMonitorManager.StopAll()
|
||||
|
||||
// Stop all tool syncers
|
||||
m.toolSyncManager.StopAll()
|
||||
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
// Disconnect all external MCP clients
|
||||
for id := range m.clientMap {
|
||||
if err := m.removeClientUnsafe(id); err != nil {
|
||||
m.logger.Error("%s Failed to remove MCP client %s: %v", MCPLogPrefix, id, err)
|
||||
}
|
||||
}
|
||||
|
||||
// Clear the client map
|
||||
m.clientMap = make(map[string]*schemas.MCPClientState)
|
||||
|
||||
// Clear local server reference
|
||||
// Note: mark3labs/mcp-go STDIO server cleanup is handled automatically
|
||||
if m.server != nil {
|
||||
m.logger.Info(MCPLogPrefix + " Clearing local MCP server reference")
|
||||
m.server = nil
|
||||
m.serverRunning = false
|
||||
}
|
||||
|
||||
m.logger.Info(MCPLogPrefix + " MCP cleanup completed")
|
||||
return nil
|
||||
}
|
||||
Reference in New Issue
Block a user